import json
import os
import re
from datetime import datetime, timezone

from src.envs import API, EVAL_REQUESTS_PATH, QUEUE_REPO, TOKEN
from src.submission.check_validity import (
    already_submitted_models,
    check_model_card,
    get_model_size,
    is_model_on_hub,
)

REQUESTED_MODELS = None
USERS_TO_SUBMISSION_DATES = None
MODEL_ID_PATTERN = re.compile(r"^[\w.-]+/[\w.-]+$")


def response(status: str, message: str) -> dict[str, str]:
    return {"status": status, "message": message}


def add_new_eval(
    model: str,
    base_model: str,
    revision: str,
    precision: str,
    weight_type: str,
    model_type: str,
):
    global REQUESTED_MODELS
    global USERS_TO_SUBMISSION_DATES
    if REQUESTED_MODELS is None:
        REQUESTED_MODELS, USERS_TO_SUBMISSION_DATES = already_submitted_models(EVAL_REQUESTS_PATH)

    model = model.strip()
    base_model = base_model.strip()
    revision = revision.strip() or "main"
    if not MODEL_ID_PATTERN.fullmatch(model):
        return response("error", "Enter a model ID in the form organization/model-name.")

    user_name, model_path = model.split("/", 1)
    precision = precision.split(" ")[0]
    current_time = datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ")

    if not model_type:
        return response("error", "Please select a model type.")

    if weight_type in ["Delta", "Adapter"]:
        if not MODEL_ID_PATTERN.fullmatch(base_model):
            return response("error", "Enter a base model ID in the form organization/model-name.")
        base_model_on_hub, error, _ = is_model_on_hub(
            model_name=base_model,
            revision=revision,
            token=TOKEN,
            test_tokenizer=True,
        )
        if not base_model_on_hub:
            return response("error", f'Base model "{base_model}" {error}')

    if weight_type != "Adapter":
        model_on_hub, error, _ = is_model_on_hub(
            model_name=model,
            revision=revision,
            token=TOKEN,
            test_tokenizer=True,
        )
        if not model_on_hub:
            return response("error", f'Model "{model}" {error}')

    try:
        model_info = API.model_info(repo_id=model, revision=revision)
    except Exception:
        return response("error", "Could not read this model. Check its ID and revision, then try again.")

    model_size = get_model_size(model_info=model_info, precision=precision)
    try:
        license_name = model_info.cardData["license"]
    except Exception:
        return response("error", "Please select a license in your model card.")

    modelcard_ok, error_msg = check_model_card(model)
    if not modelcard_ok:
        return response("error", error_msg)

    eval_entry = {
        "model": model,
        "base_model": base_model,
        "revision": revision,
        "precision": precision,
        "weight_type": weight_type,
        "status": "PENDING",
        "submitted_time": current_time,
        "model_type": model_type,
        "likes": model_info.likes,
        "params": model_size,
        "license": license_name,
        "private": False,
    }

    request_key = f"{model}_{revision}_{precision}"
    if request_key in REQUESTED_MODELS:
        return response("warning", "This model has already been submitted.")

    out_dir = os.path.join(EVAL_REQUESTS_PATH, user_name)
    os.makedirs(out_dir, exist_ok=True)
    out_path = os.path.join(out_dir, f"{model_path}_eval_request_False_{precision}_{weight_type}.json")
    with open(out_path, "w") as file:
        json.dump(eval_entry, file)

    try:
        API.upload_file(
            path_or_fileobj=out_path,
            path_in_repo=os.path.relpath(out_path, EVAL_REQUESTS_PATH),
            repo_id=QUEUE_REPO,
            repo_type="dataset",
            commit_message=f"Add {model} to eval queue",
        )
    except Exception:
        os.remove(out_path)
        return response("error", "The model passed validation, but the queue could not be updated. Please try again.")

    REQUESTED_MODELS.add(request_key)
    return response(
        "success",
        "Your model is in the evaluation queue. It should appear in the pending list shortly.",
    )