import json
import os
import pprint
from datetime import datetime, timezone

import click
from colorama import Fore
from huggingface_hub import HfApi, snapshot_download

from src.submission.check_validity import get_model_size
from src.display.utils import ModelType, WeightType, Language
from src.submission.submit import add_new_eval
from src.envs import EVAL_REQUESTS_PATH, QUEUE_REPO

precisions = ("float16", "bfloat16", "8bit (LLM.int8)", "4bit (QLoRA / FP4)", "GPTQ")
model_types = [e.name for e in ModelType]
weight_types = [e.name for e in WeightType]
language_types = [e.name for e in Language]

def main():
    api = HfApi()
    current_time = datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ")
    snapshot_download(repo_id=QUEUE_REPO, revision="main", local_dir=EVAL_REQUESTS_PATH, repo_type="dataset")

    model_name = click.prompt("Enter model name")
    revision = click.prompt("Enter revision", default="main")
    precision = click.prompt("Enter precision", default="float16", type=click.Choice(precisions))
    model_type = click.prompt("Enter model type", type=click.Choice(model_types))
    weight_type = click.prompt("Enter weight type", default="Original", type=click.Choice(weight_types))
    base_model = click.prompt("Enter base model", default="")
    main_language = click.prompt("Enter language type", default="English", type=click.Choice(language_types))

    try:
        model_info = api.model_info(repo_id=model_name, revision=revision)
    except Exception as e:
        print(f"{Fore.RED}Could not find model info for {model_name} on the Hub\n{e}{Fore.RESET}")
        return 1

    if click.confirm("Do you want to continue? This request file will be pushed to the hub"):
        click.echo("continuing...")

        print(add_new_eval(
            model=model_name,
            base_model=base_model,
            revision=revision,
            precision=precision,
            private=False,
            weight_type=weight_type,
            model_type=model_type,
            main_language=main_language,
            source="manual"
        ))
    else:
        click.echo("aborting...")


if __name__ == "__main__":
    main()