import torch
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from TTS.api import TTS
import uvicorn
import tempfile
from fastapi.responses import FileResponse


#curl -X POST -H "Content-Type: application/json" -d '{"message": "Hello ! How are you ?"}' https://micksoftware-vctk.hf.space/tts

print("Starting...")

# Charger le modèle TTS
MODEL_NAME = "tts_models/en/vctk/vits"
tts = TTS(MODEL_NAME)

# Initialiser l'application FastAPI
app = FastAPI()

print("Test du modèle TTS...")
try:
    tts.tts_to_file(text="Test audio", speaker="p339", speaker_wav=None,
                    age=22, gender="F", emotion="neutral", file_path="test.wav")
    print("Le modèle fonctionne correctement.")
except Exception as e:
    print("Erreur lors du test du modèle :", e)


SPEAKER_MAPPING = {
    "Eden": "p339",
    "voice2": "p340",
    # Ajoutez d'autres correspondances si nécessaire
}


class TextInput(BaseModel):
    text: str
    emotion: str = "neutral"
    speaker: str = "Eden"


@app.post("/tts")
def generate_tts(input: TextInput):
    """Génère un fichier audio à partir du texte donné."""
    if not input.text.strip():
        raise HTTPException(status_code=400, detail="Le texte ne peut pas être vide.")
    
    speaker_code = SPEAKER_MAPPING.get(input.speaker)
    if speaker_code is None:
        raise HTTPException(
            status_code=400, detail=f"Le speaker '{input.speaker}' n'est pas reconnu.")


    try:
        print("Generation starts:", input.text.strip(), "avec emotion:",
              input.emotion, "et speaker:", input.speaker)

        with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as tmpfile:
            output_path = tmpfile.name

        print("tmp file ready")

        
        # Générer l'audio
        tts.tts_to_file(
            text=input.text,
            speaker=speaker_code,
            speaker_wav=None,
            age=22,
            gender="F",
            emotion=input.emotion,
            file_path=output_path
        )
        print("tmp file generated")

        return FileResponse(output_path, media_type="audio/wav", filename="output.wav")
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=7860)