"""
Script de mise à jour du Dashboard des Modèles Mammouth.ai.
Récupère les tarifs publics de Mammouth.ai et les benchmarks d'Artificial Analysis
pour générer un comparatif basé sur l'efficience (Performance/Prix).
"""
import os
import requests
import json
import time
import re
from dotenv import load_dotenv

load_dotenv("../.env.global")
AIANALASYS_APIKEY = os.getenv("AIANALASYS_APIKEY")

def get_mammouth_models():
    url = "https://api.mammouth.ai/public/models"
    try:
        requests.packages.urllib3.disable_warnings()
        response = requests.get(url, verify=False)
        response.raise_for_status()
        return response.json().get('data', [])
    except Exception as e:
        print(f"Error fetching Mammouth models: {e}")
        return []

def get_aa_data():
    url = "https://artificialanalysis.ai/api/v2/data/llms/models"
    headers = {"x-api-key": AIANALASYS_APIKEY}
    try:
        response = requests.get(url, headers=headers)
        response.raise_for_status()
        return response.json().get('data', [])
    except Exception as e:
        print(f"Error fetching AA data: {e}")
        return []

def clean_id(model_id):
    id_clean = re.sub(r'-\d{4,8}', '', model_id.lower())
    id_clean = id_clean.replace('-latest', '').replace('-preview', '').replace('-instruct', '')
    return id_clean.strip()

def generate_markdown(models_data):
    categories = {}
    for m in models_data:
        cat = m.get('category', 'General')
        if cat not in categories: categories[cat] = []
        categories[cat].append(m)
        
    md = "# Dashboard des Modèles Mammouth.ai\n\n"
    md += "*Analyse comparative basée sur le prix (Mammouth) et la performance (Artificial Analysis).*\n\n"
    md += f"Dernière mise à jour : {time.strftime('%Y-%m-%d %H:%M:%S')}\n\n"
    md += "### Légende :\n"
    md += "- **Note Prix** : 10 = Le moins cher, 0 = Le plus cher.\n"
    md += "- **Efficience** : Ratio Performance / Prix. Un score élevé indique un excellent rapport qualité/prix.\n\n"
    
    order = ['Coding', 'Agents', 'General']
    sorted_cats = sorted(categories.keys(), key=lambda x: order.index(x) if x in order else 99)

    for cat in sorted_cats:
        md += f"## {cat}\n\n"
        md += "| Modèle | Prix (In/Out 1M) | Score | Vitesse | Note Prix | **Efficience** |\n"
        md += "| :--- | :--- | :--- | :--- | :--- | :--- |\n"
        models = categories[cat]
        # Tri par Efficience décroissante
        models.sort(key=lambda x: x.get('efficiency_score', 0), reverse=True)
        for m in models:
            score_str = f"**{m['score']:.1f}**" if m['score'] else "N/A"
            speed_str = f"{m['speed']:.1f}" if m['speed'] else "N/A"
            p_score = f"{m['price_score']:.1f}" if m['price_score'] is not None else "N/A"
            eff_score = f"**{m['efficiency_score']:.1f}**" if m['efficiency_score'] else "N/A"
            md += f"| {m['name']} | ${m['price_in']:.2f}/${m['price_out']:.2f} | {score_str} | {speed_str} | {p_score} | {eff_score} |\n"
        md += "\n"
    return md

def main():
    print("Calcul de l'efficience des modèles...")
    m_models = get_mammouth_models()
    aa_data = get_aa_data()
    
    aa_map = {m.get('slug', '').lower(): m for m in aa_data}
    for aa_m in aa_data:
        aa_map[aa_m.get('name', '').lower()] = aa_m

    enriched = []
    # On calcule d'abord les prix pour déterminer les échelles de note
    temp_list = []
    for m in m_models:
        m_id = m.get('id', '')
        info = m.get('model_info', {})
        if not m_id: continue
        
        price_in = float(info.get('input_cost_per_token', 0)) * 1000000
        price_out = float(info.get('output_cost_per_token', 0)) * 1000000
        # Prix combiné (moyenne pondérée 3:1 comme AA)
        blended_price = (price_in * 0.75) + (price_out * 0.25)
        
        if blended_price > 0:
            temp_list.append((m, blended_price, price_in, price_out))

    if not temp_list: return

    # Calcul des échelles de prix pour la note (Log scale pour mieux différencier)
    min_p = min(x[1] for x in temp_list)
    max_p = max(x[1] for x in temp_list)

    for m_data, b_price, p_in, p_out in temp_list:
        m_id = m_data.get('id', '')
        m_id_clean = clean_id(m_id)
        short_id = m_id_clean.split('/')[-1]
        
        aa_info = aa_map.get(m_id_clean) or aa_map.get(short_id)
        if not aa_info:
            norm = m_id_clean.replace('-', '').replace('.', '')
            for k, v in aa_map.items():
                if k.replace('-', '').replace('.', '') == norm:
                    aa_info = v
                    break

        # Note Prix : 10 pour le moins cher, 0 pour le plus cher
        # Formule : 10 * (1 - (price - min) / (max - min))
        price_score = 10 * (1 - (b_price - min_p) / (max_p - min_p)) if max_p > min_p else 10
        
        score = None
        speed = None
        category = "General"
        if any(x in m_id_clean for x in ['coding', 'code', 'starcoder', 'codestral', 'coder']): category = "Coding"
        elif any(x in m_id_clean for x in ['agent', 'hermes', 'tool', 'function', 'sonar']): category = "Agents"

        if aa_info:
            evals = aa_info.get('evaluations', {})
            score = evals.get('artificial_analysis_coding_index') if category == "Coding" else evals.get('artificial_analysis_intelligence_index')
            if not score: score = evals.get('artificial_analysis_intelligence_index')
            speed = aa_info.get('median_output_tokens_per_second')

        # Efficience : On combine la performance et le prix
        # Si pas de score AA, on base l'efficience uniquement sur le prix (avec un bonus de base)
        efficiency_score = 0
        if score:
            # Score normalisé (0-100) * Note Prix (0-10) / 10
            efficiency_score = (score * price_score) / 10
        else:
            # Modèle sans benchmark : on lui donne une efficience basée sur son prix seul
            efficiency_score = price_score * 2 # Moins prioritaire que ceux avec score

        enriched.append({
            'name': m_id,
            'price_in': p_in,
            'price_out': p_out,
            'score': score,
            'speed': speed,
            'price_score': price_score,
            'efficiency_score': efficiency_score,
            'category': category
        })

    with open("README.md", "w", encoding="utf-8") as f:
        f.write(generate_markdown(enriched))
    print(f"Success! Dashboard updated with {len(enriched)} models.")

if __name__ == "__main__":
    main()