from fastapi import FastAPI, HTTPException
import subprocess
import os
from pydantic import BaseModel
# from melo.api import TTS  # Ensure you have the correct import path for TTS

app = FastAPI()

# Define the request body model
class TTSRequest(BaseModel):
    text: str
    language: str
    speaker: str
    speed: int

language_dict = {
    "English": "EN",
    "Spanish": "ES",
    "French": "FR",
    "Chinese": "ZH",
    "Japanese": "JP",
    "Korean": "KR"
}

speakers_dict = {
    "Default_Accent": "EN-Default",
    "American_Accent": "EN-US",
    "British_Accent": "EN-BR",
    "Indian_Accent": "EN-IN",
    "Australian_Accent": "EN-AU",
    "Spanish_Accent": "ES",
    "French_Accent": "FR",
    "Chinese_Accent": "ZH",
    "Japanese_Accent": "JP",
    "Korean_Accent": "KO"
}

device = 'auto'  # Will automatically use GPU if available

@app.on_event("startup")
async def install_melotts():
    # Clone the MeloTTS repository
    subprocess.run(["git", "clone", "https://github.com/myshell-ai/MeloTTS.git"], check=True)
    
    # Change directory to MeloTTS
    os.chdir("MeloTTS")
    
    # Install MeloTTS
    subprocess.run(["pip", "install", "-e", "."], check=True)
    
    # Download Unidic
    subprocess.run(["python", "-m", "unidic", "download"], check=True)

@app.post("/generate-tts/")
async def generate_tts(tts_request: TTSRequest):
    # Validate language and speaker
    language_code = language_dict.get(tts_request.language)
    if not language_code:
        raise HTTPException(status_code=400, detail="Invalid language selected.")

    speaker_id = speakers_dict.get(tts_request.speaker)
    if not speaker_id:
        raise HTTPException(status_code=400, detail="Invalid speaker selected.")

    # Initialize the model
    from melo.api import TTS
    try:
        model = TTS(language=language_code, device=device)
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Error initializing TTS model: {e}")

    # Get speaker ID
    try:
        selected_speaker_id = model.hps.data.spk2id[speaker_id]
    except KeyError:
        raise HTTPException(status_code=400, detail="Invalid speaker ID.")

    # Generate TTS
    output_path = 'generated_MeloTTS.wav'
    try:
        model.tts_to_file(tts_request.text, selected_speaker_id, output_path, speed=tts_request.speed)
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Error generating TTS: {e}")

    return {"message": "TTS generated successfully", "output_file": output_path}

@app.get("/")
async def read_root():
    subprocess.run(["ls"], check=True)
    return {"message": "MeloTTS installation started!"}

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