import glob
import json
import os
import re
import pickle
from typing import List

import huggingface_hub
from huggingface_hub import HfApi
from tqdm import tqdm
from transformers import AutoModel, AutoConfig
from accelerate import init_empty_weights

from src.display_models.model_metadata_flags import DO_NOT_SUBMIT_MODELS, FLAGGED_MODELS
from src.display_models.model_metadata_type import MODEL_TYPE_METADATA, ModelType, model_type_from_str
from src.display_models.utils import AutoEvalColumn, model_hyperlink

api = HfApi(token=os.environ.get("H4_TOKEN", None))


def get_model_infos_from_hub(leaderboard_data: List[dict]):
    # load cache from disk
    try:
        with open("model_info_cache.pkl", "rb") as f:
            model_info_cache = pickle.load(f)
    except (EOFError, FileNotFoundError):
        model_info_cache = {}
    try:
        with open("model_size_cache.pkl", "rb") as f:
            model_size_cache = pickle.load(f)
    except (EOFError, FileNotFoundError):
        model_size_cache = {}

    for model_data in tqdm(leaderboard_data):
        model_name = model_data["model_name_for_query"]

        if model_name in model_info_cache:
            model_info = model_info_cache[model_name]
        else:
            try:
                model_info = api.model_info(model_name)
                model_info_cache[model_name] = model_info
            except huggingface_hub.utils._errors.RepositoryNotFoundError:
                print("Repo not found!", model_name)
                model_data[AutoEvalColumn.license.name] = None
                model_data[AutoEvalColumn.likes.name] = None
                if model_name not in model_size_cache:
                    model_size_cache[model_name] = get_model_size(model_name, None)
                model_data[AutoEvalColumn.params.name] = model_size_cache[model_name]

        model_data[AutoEvalColumn.license.name] = get_model_license(model_info)
        model_data[AutoEvalColumn.likes.name] = get_model_likes(model_info)
        if model_name not in model_size_cache:
            model_size_cache[model_name] = get_model_size(model_name, model_info)
        model_data[AutoEvalColumn.params.name] = model_size_cache[model_name]
    
    # save cache to disk in pickle format
    with open("model_info_cache.pkl", "wb") as f:
        pickle.dump(model_info_cache, f)
    with open("model_size_cache.pkl", "wb") as f:
        pickle.dump(model_size_cache, f)


def get_model_license(model_info):
    try:
        return model_info.cardData["license"]
    except Exception:
        return "?"


def get_model_likes(model_info):
    return model_info.likes


size_pattern = re.compile(r"(\d\.)?\d+(b|m)")


def get_model_size(model_name, model_info):
    # In billions
    try:
        return round(model_info.safetensors["total"] / 1e9, 3)
    except AttributeError:
        try:
            config = AutoConfig.from_pretrained(model_name, trust_remote_code=False)
            with init_empty_weights():
                model = AutoModel.from_config(config, trust_remote_code=False)
            return round(sum(p.numel() for p in model.parameters() if p.requires_grad) / 1e9, 3)
        except (EnvironmentError, ValueError): # model config not found, likely private  
            try:
                size_match = re.search(size_pattern, model_name.lower())
                size = size_match.group(0)
                return round(float(size[:-1]) if size[-1] == "b" else float(size[:-1]) / 1e3, 3)
            except AttributeError:
                return 0


def get_model_type(leaderboard_data: List[dict]):
    for model_data in leaderboard_data:
        request_files = os.path.join(
            "eval-queue",
            model_data["model_name_for_query"] + "_eval_request_*" + ".json",
        )
        request_files = glob.glob(request_files)

        # Select correct request file (precision)
        request_file = ""
        if len(request_files) == 1:
            request_file = request_files[0]
        elif len(request_files) > 1:
            request_files = sorted(request_files, reverse=True)
            for tmp_request_file in request_files:
                with open(tmp_request_file, "r") as f:
                    req_content = json.load(f)
                    if (
                        req_content["status"] == "FINISHED"
                        and req_content["precision"] == model_data["Precision"].split(".")[-1]
                    ):
                        request_file = tmp_request_file

        try:
            with open(request_file, "r") as f:
                request = json.load(f)
            model_type = model_type_from_str(request["model_type"])
            model_data[AutoEvalColumn.model_type.name] = model_type.value.name
            model_data[AutoEvalColumn.model_type_symbol.name] = model_type.value.symbol  # + ("🔺" if is_delta else "")
        except Exception:
            if model_data["model_name_for_query"] in MODEL_TYPE_METADATA:
                model_data[AutoEvalColumn.model_type.name] = MODEL_TYPE_METADATA[
                    model_data["model_name_for_query"]
                ].value.name
                model_data[AutoEvalColumn.model_type_symbol.name] = MODEL_TYPE_METADATA[
                    model_data["model_name_for_query"]
                ].value.symbol  # + ("🔺" if is_delta else "")
            else:
                model_data[AutoEvalColumn.model_type.name] = ModelType.Unknown.value.name
                model_data[AutoEvalColumn.model_type_symbol.name] = ModelType.Unknown.value.symbol


def flag_models(leaderboard_data: List[dict]):
    for model_data in leaderboard_data:
        if model_data["model_name_for_query"] in FLAGGED_MODELS:
            issue_num = FLAGGED_MODELS[model_data["model_name_for_query"]].split("/")[-1]
            issue_link = model_hyperlink(
                FLAGGED_MODELS[model_data["model_name_for_query"]],
                f"See discussion #{issue_num}",
            )
            model_data[
                AutoEvalColumn.model.name
            ] = f"{model_data[AutoEvalColumn.model.name]} has been flagged! {issue_link}"


def remove_forbidden_models(leaderboard_data: List[dict]):
    indices_to_remove = []
    for ix, model in enumerate(leaderboard_data):
        if model["model_name_for_query"] in DO_NOT_SUBMIT_MODELS:
            indices_to_remove.append(ix)

    for ix in reversed(indices_to_remove):
        leaderboard_data.pop(ix)
    return leaderboard_data


def apply_metadata(leaderboard_data: List[dict]):
    leaderboard_data = remove_forbidden_models(leaderboard_data)
    get_model_type(leaderboard_data)
    get_model_infos_from_hub(leaderboard_data)
    flag_models(leaderboard_data)