import os
import requests
import json
import time
from dotenv import load_dotenv

# Charger .env.global depuis le répertoire parent
load_dotenv("../.env.global")

MAMMOUTH_APIKEY = os.getenv("MAMMOUTH_APIKEY")
AIANALASYS_APIKEY = os.getenv("AIANALASYS_APIKEY")

def get_mammouth_models():
    # Mammouth utilise l'API OpenRouter (revendeur)
    url = "https://openrouter.ai/api/v1/models"
    headers = {"Authorization": f"Bearer {MAMMOUTH_APIKEY}"}
    try:
        response = requests.get(url, headers=headers)
        response.raise_for_status()
        return response.json()['data']
    except Exception as e:
        print(f"Error fetching Mammouth models: {e}")
        return []

def get_aa_data():
    # URL correcte d'après la doc (version v2)
    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()
        # Le diagnostic a montré que les données sont dans 'data'
        return response.json().get('data', [])
    except Exception as e:
        print(f"Error fetching Artificial Analysis data: {e}")
        return []

def generate_markdown(models_data):
    # Trier par catégorie (genre)
    categories = {}
    
    for m in models_data:
        cat = m.get('category', 'General')
        if cat not in categories:
            categories[cat] = []
        categories[cat].append(m)

    md = "# Table des Modèles Mammouth.ai\n\n"
    md += "*Mise à jour automatique via Artificial Analysis & Mammouth API*\n\n"
    md += "Dernière mise à jour : " + time.strftime("%Y-%m-%d %H:%M:%S") + "\n\n"

    # Liste des catégories dans un ordre spécifique
    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:
        models = categories[cat]
        md += f"## {cat}\n\n"
        md += "| Modèle | Prix (In / Out / 1M) | Performance (AA Index) | Vitesse (TPS) |\n"
        md += "| :--- | :--- | :--- | :--- |\n"
        # Trier par performance (AA index)
        models.sort(key=lambda x: x.get('score') or 0, reverse=True)
        for m in models:
            p_in = f"${m['price_in']:.2f}" if m['price_in'] is not None else "N/A"
            p_out = f"${m['price_out']:.2f}" if m['price_out'] is not None else "N/A"
            score = f"**{m['score']:.1f}**" if m['score'] else "N/A"
            speed = f"{m['speed']:.1f}" if m['speed'] else "N/A"
            md += f"| {m['name']} | {p_in} / {p_out} | {score} | {speed} |\n"
        md += "\n"
    
    return md

def main():
    print("Fetching Mammouth models...")
    mammouth_models = get_mammouth_models()
    
    print("Fetching Artificial Analysis data...")
    aa_data = get_aa_data()
    
    # Créer un dictionnaire de mapping pour AA (clé: nom du modèle en minuscule)
    aa_map = {}
    for aa_m in aa_data:
        name = aa_m.get('model_name', '').lower()
        aa_map[name] = aa_m

    enriched_models = []
    
    for m in mammouth_models:
        m_id = m['id']
        m_name = m['name'].lower()
        short_name = m_id.split('/')[-1].lower()
        
        # Mapping logique plus complet
        aa_info = aa_map.get(m_name) or aa_map.get(short_name)
        
        # Si pas de match exact, on cherche par sous-chaîne ou flou
        if not aa_info:
            for key in aa_map:
                if key in m_name or m_name in key or key in short_name or short_name in key:
                    aa_info = aa_map[key]
                    break

        # Extraction des prix Mammouth (prix pour 1 token chez OpenRouter)
        pricing = m.get('pricing', {})
        try:
            price_in = float(pricing.get('prompt', 0)) * 1000000
            price_out = float(pricing.get('completion', 0)) * 1000000
        except (ValueError, TypeError):
            price_in = 0
            price_out = 0
        
        score = None
        speed = None
        # On essaie d'extraire la catégorie de AA, sinon on devine
        category = "General"
        
        if aa_info:
            evals = aa_info.get('evaluations', {})
            # On cherche l'intelligence index
            score = evals.get('artificial_analysis_intelligence_index')
            speed = aa_info.get('median_output_tokens_per_second')
            
        # Détermination de la catégorie (Genre)
        if any(x in m_name or x in short_name for x in ['coding', 'code', 'starcoder', 'stable-code', 'deepseek-coder']):
            category = "Coding"
        elif any(x in m_name or x in short_name for x in ['agent', 'hermes', 'tool']):
            category = "Agents"
        else:
            category = "General"

        enriched_models.append({
            'name': m['name'],
            'price_in': price_in,
            'price_out': price_out,
            'score': score,
            'speed': speed,
            'category': category
        })

    # On ne garde que les modèles qui ont un score de performance OU un prix raisonnable 
    # (Certains modèles sont gratuits ou ont des prix nuls)
    final_list = [m for m in enriched_models if m['price_in'] > 0 or m['score'] is not None]
    
    markdown = generate_markdown(final_list)
    
    with open("README.md", "w", encoding="utf-8") as f:
        f.write(markdown)
    
    print(f"README.md updated with {len(final_list)} models!")

if __name__ == "__main__":
    main()