from fastapi import FastAPI, WebSocket, WebSocketDisconnect
from fastapi.middleware.cors import CORSMiddleware
from gradio_client import Client, handle_file
import tempfile
import wave
import os

app = FastAPI()

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

# 🧠 懶加載大腦：防止啟動超時
brain = {"asr": None, "mt": None}

def get_brain(key):
    if brain[key] is None:
        urls = {
            "asr": "https://ai-labs.ilrdf.org.tw/sapolita-kaldi/",
            "mt": "https://ai-labs.ilrdf.org.tw/kari-seejiq-tnpusu-ai-hmjil/"
        }
        brain[key] = Client(urls[key])
    return brain[key]

@app.get("/")
async def root():
    return {"status": "online", "mode": "太魯閣語測試模式"}

@app.websocket("/ws/subtitle")
async def websocket_subtitle(websocket: WebSocket):
    await websocket.accept()
    audio_buffer = bytearray()
    print("📢 太魯閣語專線已啟動")
    
    try:
        while True:
            audio_chunk = await websocket.receive_bytes()
            audio_buffer.extend(audio_chunk)
            
            # 累積 25 萬 bytes (約 3 秒) 提高太魯閣語辨識率
            if len(audio_buffer) >= 250000:
                with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as temp_wav:
                    with wave.open(temp_wav.name, 'wb') as wav_file:
                        wav_file.setnchannels(1)
                        wav_file.setsampwidth(2)
                        wav_file.setframerate(44100)
                        wav_file.writeframes(audio_buffer)
                    temp_path = temp_wav.name

                try:
                    # 1. 太魯閣語 ASR
                    asr_res = get_brain("asr").predict(
                        dialect_id="formosan_trv", 
                        audio_data=handle_file(temp_path), 
                        api_name="/automatic_speech_recognition"
                    )
                    ind_text = str(asr_res).strip()

                    # 2. 太魯閣語翻譯
                    zh_text = ""
                    if ind_text and ind_text not in [".", "...", ""]:
                        mt_brain = get_brain("mt")
                        # 抓取太魯閣語代碼
                        d_code = "trv_Truku" # 太魯閣語固定代碼
                        zh_res = mt_brain.predict(text=ind_text, src_lang=d_code, tgt_lang="zho_Hant", api_name="/translate")
                        zh_text = str(zh_res).strip()

                    await websocket.send_json({
                        "status": "recognizing",
                        "indigenous": ind_text,
                        "chinese": zh_text
                    })
                except Exception as e:
                    print(f"辨識出錯: {e}")

                audio_buffer = bytearray()
                if os.path.exists(temp_path): os.remove(temp_path)
            else:
                if len(audio_buffer) % 60000 == 0:
                    await websocket.send_json({"status": "buffering", "text": "🟢 神獸正在聽太魯閣語..."})
    except Exception: pass

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