import uvicorn
from fastapi import FastAPI, HTTPException
from fastapi.responses import FileResponse
from pydantic import BaseModel
import ChatTTS
import scipy.io.wavfile as wavfile
import numpy as np
import os

app = FastAPI(title="Local ChatTTS Server (Fixed Index)")

print("正在載入 ChatTTS 模型...")
chat = ChatTTS.Chat()
chat.load()
print("模型載入完成！")

class TTSRequest(BaseModel):
    text: str

@app.post("/v1/audio/speech")
async def text_to_speech(request: TTSRequest):
    script_dir = os.path.dirname(os.path.abspath(__file__))
    record_dir = os.path.join(script_dir, "../record")
    os.makedirs(record_dir, exist_ok=True)
    output_path = os.path.join(record_dir, "output.wav")
    try:
        # 1. 進行推理 (使用 InferCodeParams 對象，限制最大 token 數，避免無限生成)
        params_infer_code = ChatTTS.Chat.InferCodeParams(
            max_new_token=384
        )
        res = chat.infer(
            [request.text], 
            use_decoder=True,
            params_infer_code=params_infer_code
        )

        # 2. 智慧型格式剝離：確保拿到最裡面的純音訊數據
        if isinstance(res, list):
            audio_data = res[0]
        else:
            audio_data = res

        if isinstance(audio_data, list) or (hasinstance := hasattr(audio_data, 'ndim') and audio_data.ndim > 1):
            if hasattr(audio_data, 'ndim') and audio_data.ndim > 1:
                audio_data = audio_data[0]
            else:
                audio_data = audio_data[0]

        # 3. 強制轉換成標準 1D float32 numpy 陣列並拉平
        audio_data = np.array(audio_data, dtype=np.float32).flatten()

        # 4. 寫入 WAV 檔案
        wavfile.write(output_path, 24000, audio_data)

        if os.path.exists(output_path):
            return FileResponse(output_path, media_type="audio/wav", filename="speech.wav")
        else:
            raise HTTPException(status_code=500, detail="語音檔案生成失敗")

    except Exception as e:
        raise HTTPException(status_code=500, detail=f"TTS 處理失敗: {str(e)}")

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