from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
import os
import sys

# 添加父目录到路径
sys.path.insert(0, '/app')

app = FastAPI(title="IndexTTS 2.5 API")

tts = None

class TTSRequest(BaseModel):
    text: str
    spk_audio_prompt: str = "default.wav"
    lang: str = "ZH"
    emotion: str = "calm"
    speed: float = 1.0
    output_path: str = "/app/output/output.wav"

@app.on_event("startup")
async def load_model():
    global tts
    try:
        from indextts.infer_v2_5 import IndexTTS2
        
        config_path = os.environ.get('CONFIG_PATH', '/checkpoints/config.yaml')
        model_dir = os.environ.get('MODEL_DIR', '/checkpoints')
        
        tts = IndexTTS2(
            cfg_path=config_path,
            model_dir=model_dir,
            use_bf16=False  # CPU 模式使用 FP32
        )
        print(f"✅ IndexTTS 2.5 模型加载完成", file=sys.stderr)
    except Exception as e:
        print(f"⚠️ 模型加载失败: {e}", file=sys.stderr)
        tts = None

@app.post("/tts")
async def generate_speech(req: TTSRequest):
    if tts is None:
        raise HTTPException(status_code=503, detail="模型未加载")
    
    try:
        import torch
        
        # 情感向量映射
        emo_map = {
            "happy": [0, 0, 0, 0, 0, 0, 0, 1],
            "angry": [0, 1, 0, 0, 0, 0, 0, 0],
            "sad": [0, 0, 1, 0, 0, 0, 0, 0],
            "afraid": [0, 0, 0, 1, 0, 0, 0, 0],
            "surprised": [0, 0, 0, 0, 0, 0, 0, 1],
            "calm": [0, 0, 0, 0, 0, 0, 0, 0],
        }
        
        emo_vector = emo_map.get(req.emotion, emo_map["calm"])
        
        tts.infer(
            spk_audio_prompt=req.spk_audio_prompt,
            text=req.text,
            lang=req.lang,
            output_path=req.output_path,
            emo_vector=emo_vector,
            duration_factor=req.speed
        )
        
        return {
            "status": "success",
            "path": req.output_path,
            "text_length": len(req.text)
        }
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.get("/health")
async def health():
    return {
        "status": "healthy",
        "model_loaded": tts is not None
    }

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