import os
import re
import json
from huggingface_hub import HfApi

# Configura il token di Hugging Face (se necessario)
#TOKEN = "YOUR_HUGGINGFACE_API_TOKEN"
api = HfApi()

# Percorsi delle cartelle
input_folder = "../evalita_llm_models_output/"
output_folder = "../evalita_llm_requests/"

# Creazione della cartella di output se non esiste
os.makedirs(output_folder, exist_ok=True)

# Espressione regolare per trovare il nome del modello
model_pattern = re.compile(r"pretrained=([\w\-./]+)")

# Scansiona i file nella cartella di input
for filename in os.listdir(input_folder):
    file_path = os.path.join(input_folder, filename)

    # Leggi il contenuto del file
    with open(file_path, "r", encoding="utf-8") as f:
        content = f.read()

    # Estrai il nome del modello
    match = model_pattern.search(content)
    if match:
        model_name = match.group(1)
        print(f"Processing model: {model_name}")

        try:
            # Ottieni le informazioni del modello da Hugging Face
            model_info = api.model_info(model_name)

            # Calcola il numero di parametri in miliardi, se disponibile
            num_params = None
            if model_info.safetensors and "BF16" in model_info.safetensors.parameters:
                num_params = model_info.safetensors.parameters["BF16"] / 1e9  # Converti in miliardi

            # Estrai la lingua (può essere una lista, quindi prendiamo la prima se esiste)
            # Estrai e concatena i linguaggi
            language = "_".join(model_info.card_data.get("language", [])) if model_info.card_data else ""

            print(model_info)

            # Costruisci il dizionario con i metadati richiesti
            model_data = {
                "model": model_name,
                "base_model": model_info.config.get("architectures", [""])[0] if model_info.config else "",
                "revision": model_info.sha,
                #"precision": "bfloat16",  # Se disponibile, sostituire con un valore reale
                #"weight_type": "Original",
                #"status": "FINISHED",
                "submitted_time": str(model_info.created_at),
                #"model_type": "pretrained",
                #"likes": model_info.likes,
                #"params": model_info.safetensors_size_in_bytes / 1e9 if model_info.safetensors_size_in_bytes else None,
                #"license": model_info.license,
                #"private": model_info.private,
                "num_params_billion": num_params,  # Numero di parametri in miliardi
                "language": language,  # Lingua estratta
            }

            # Separare il model_name in due parti: prima e dopo "/"
            if "/" in model_name:
                dir_name, file_name = model_name.split("/", 1)
            else:
                dir_name, file_name = model_name, model_name  # Se non c'è "/", usa lo stesso nome

            # Creare la cartella per la prima parte del nome del modello
            model_output_folder = os.path.join(output_folder, dir_name)
            os.makedirs(model_output_folder, exist_ok=True)

            # Salvare il file JSON nella cartella appropriata
            output_file = os.path.join(model_output_folder, f"{file_name}.json")
            with open(output_file, "w", encoding="utf-8") as f:
                json.dump(model_data, f, indent=4)

            print(f"Saved metadata for {model_name} in {output_file}")

        except Exception as e:
            print(f"Error retrieving info for {model_name}: {e}")

        print("Process completed1.")