import os
import shutil
import hashlib
from fastapi import FastAPI, Form
from fastapi.middleware.cors import CORSMiddleware
from fastapi.staticfiles import StaticFiles
from fastapi.responses import JSONResponse
import google.generativeai as genai
from gradio_client import Client
import uvicorn

app = FastAPI()

# 💡 1. 開放跨網域連線 (CORS)
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"], 
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# 💡 2. 準備靜態檔案資料夾
if not os.path.exists("static"):
    os.makedirs("static")
app.mount("/static", StaticFiles(directory="static"), name="static")

# 💡 3. API Keys & Models 初始化 (對接 2026 最新 Gemini 3.5 Flash)
GOOGLE_API_KEY = os.getenv('GEMINI_KEY', '請填入您的GEMINI_KEY')
genai.configure(api_key=GOOGLE_API_KEY)
MODEL_NAME = 'gemini-3.5-flash' 
model = genai.GenerativeModel(MODEL_NAME) 

# 💡 4. 初始化原語會的翻譯與 TTS 伺服器
tts_client = Client("https://ai-labs.ilrdf.org.tw/hnang-kari-ai-asi-sluhay/")
mt_client  = Client("https://ai-labs.ilrdf.org.tw/kari-seejiq-tnpusu-ai-hmjil/")

TRIBE_CONFIG = {
    "阿美": {"mt": "阿美"}, "泰雅": {"mt": "泰雅"}, "排灣": {"mt": "排灣"}, "布農": {"mt": "布農"},
    "卑南": {"mt": "卑南"}, "魯凱": {"mt": "魯凱"}, "鄒": {"mt": "鄒"}, "賽夏": {"mt": "賽夏"},
    "雅美": {"mt": "雅美"}, "邵": {"mt": "邵"}, "噶瑪蘭": {"mt": "噶瑪蘭"}, "太魯閣": {"mt": "太魯閣"},
    "撒奇萊雅": {"mt": "撒奇萊雅"}, "賽德克": {"mt": "賽德克"}, "拉阿魯哇": {"mt": "拉阿魯哇"}, "卡那卡那富": {"mt": "卡那卡那富"}
}

def get_clean_value(res):
    if isinstance(res, dict) and 'value' in res: return res['value']
    if isinstance(res, list) and len(res) > 0: return res[0]
    return res

def format_native_text(text, tribe):
    if not text: return text
    text = text.strip()
    if tribe in ["賽夏", "噶瑪蘭", "泰雅"]: return text[0].lower() + text[1:]
    else: return text[0].upper() + text[1:]

@app.get("/")
def read_root():
    return {"status": "ILRDF Avatar Text Brain is Running Successfully!", "model": MODEL_NAME}

@app.post("/api/chat")
async def avatar_chat(
    tribe: str = Form(...), 
    text: str = Form(...)  # 👈 接收前端秒傳過來的中文文字
):
    try:
        config = TRIBE_CONFIG.get(tribe)
        if not config:
            return JSONResponse({"error": f"不支援的族語: {tribe}"}, status_code=400)

        zh_in = text.strip()

        # --- 步驟 A: AI 大腦生成對話 (Gemini 3.5 Flash) ---
        prompt = f"你現在是與我對話的{tribe}族朋友。請用中文聊天。回覆規則：1.口吻生活化親切。2.不要說教解釋。3.字數一定要少（15字內），限一個短句。\n\n使用者說：{zh_in}"
        response = model.generate_content(prompt)
        ai_zh = response.text.strip()

        # --- 步驟 B: 翻譯回族語 ---
        back_code = get_clean_value(mt_client.predict(ethnicity=config["mt"], api_name="/lambda_1"))
        ai_native = get_clean_value(mt_client.predict(text=ai_zh, src_lang="zho_Hant", tgt_lang=back_code, api_name="/translate_1"))
        ai_native = format_native_text(ai_native, tribe)

        # --- 步驟 C: 語音合成 (TTS) ---
        speaker = get_clean_value(tts_client.predict(ethnicity=config["mt"], api_name="/lambda"))
        if tribe == "阿美": speaker = "阿美_秀姑巒_女聲1"
        temp_tts = tts_client.predict(ref=speaker, gen_text_input=ai_native, api_name="/default_speaker_tts")

        unique_str = f"reply_{hashlib.md5(ai_native.encode()).hexdigest()[:8]}.wav"
        out_filepath = f"static/{unique_str}" 
        shutil.move(temp_tts, out_filepath)

        # --- 步驟 D: 組合完整語音網址 ---
        base_url = os.getenv('SPACE_HOST', '') 
        if base_url:
            base_url = f"https://{base_url}"
        else:
            base_url = "http://127.0.0.1:7860"

        audio_url = f"{base_url}/{out_filepath}"

        return {
            "user_zh_text": zh_in,
            "ai_zh_text": ai_zh,
            "ai_native_text": ai_native,
            "audio_url": audio_url
        }

    except Exception as e:
        return JSONResponse({"error": str(e)}, status_code=500)

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