from fastapi import FastAPI, Response, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from melo.api import TTS
import os

app = FastAPI()

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)

# Stockage des modèles chargés
models = {}
device = 'cpu'

print("⏳ Initialisation de MeloTTS...")

try:
    # Chargement initial FR et EN
    print("  - Chargement FR...")
    models['FR'] = TTS(language='FR', device=device)
    
    print("  - Chargement EN...")
    models['EN'] = TTS(language='EN', device=device)
    
    print("✅ MeloTTS Prêt (FR + EN) !")
except Exception as e:
    print(f"❌ Erreur au chargement : {e}")

@app.post("/tts")
async def generate_speech(data: dict):
    text = data.get("text", "")
    lang = data.get("lang", "FR").upper() # Par défaut FR
    speed = data.get("speed", 1.0)

    if not text:
        raise HTTPException(status_code=400, detail="Texte vide")

    # 1. Chargement Dynamique du modèle (si pas encore en mémoire)
    if lang not in models:
        try:
            print(f"⏳ Chargement dynamique de {lang}...")
            models[lang] = TTS(language=lang, device=device)
        except Exception as e:
            raise HTTPException(status_code=400, detail=f"Langue {lang} non supportée: {e}")

    model = models[lang]
    speaker_ids = model.hps.data.spk2id
    
    # 2. SELECTION DE LA VOIX (Selon la Doc Officielle)
    # Melo utilise des clés spécifiques. L'objet speaker_ids ne supporte PAS .get()
    # Il faut utiliser les crochets [] directement.
    
    try:
        spk_id = None
        
        if lang == 'EN':
            # La doc dit: 'EN-Default', 'EN-US', 'EN-BR', 'EN_INDIA', 'EN-AU'
            # On prend l'accent Américain par défaut pour 'EN'
            spk_id = speaker_ids['EN-US']
        elif lang == 'FR':
            # La doc dit: 'FR'
            spk_id = speaker_ids['FR']
        elif lang == 'ES':
            # La doc dit: 'ES'
            spk_id = speaker_ids['ES']
        elif lang == 'ZH':
            spk_id = speaker_ids['ZH']
        elif lang == 'JP':
            spk_id = speaker_ids['JP']
        elif lang == 'KR':
            spk_id = speaker_ids['KR']
        else:
            # Pour les autres langues, on tente la clé langue directement
            # Si ça échoue, le bloc 'except' l'attrapera
            spk_id = speaker_ids[lang]
            
    except Exception as e:
        # Diagnostic précis si la clé n'existe pas
        print(f"⚠️ Erreur Speaker ID: {e}")
        # Tentative ultime : accès par attribut (certaines versions de HParams le permettent)
        try:
            spk_id = getattr(speaker_ids, lang)
        except:
             return Response(content=f"Erreur: Impossible de trouver la voix pour {lang}", status_code=500)

    output_path = f"/tmp/output_{os.getpid()}.wav"

    try:
        # 3. Génération
        model.tts_to_file(text, spk_id, output_path, speed=speed)
        
        with open(output_path, "rb") as f:
            audio_data = f.read()
            
        return Response(content=audio_data, media_type="audio/wav")

    except Exception as e:
        print(f"❌ Erreur génération : {e}")
        return Response(content=str(e), status_code=500)

@app.get("/")
def home():
    return {"status": "MeloTTS Ready", "loaded": list(models.keys())}