import base64
import io
import torch
from fastapi import FastAPI, HTTPException, Request
from transformers import pipeline

app = FastAPI()

# 1. 初始化你的 Whisper 模型的 Pipeline
print("正在載入 Whisper 模型...")
asr = pipeline(
    "automatic-speech-recognition",
    model="openai/whisper-large-v3-turbo",
    torch_dtype=torch.float16,
    device="cuda:0",  # 如果沒有 GPU 請改成 "cpu"
)
print("模型載入完成！")


@app.post("/v1/audio/transcriptions")
async def transcribe_audio(request: Request):
    try:
        content_type = request.headers.get("content-type", "")
        if "multipart/form-data" in content_type:
            # 支援標準的 OpenAI multipart/form-data 格式 (如 reachy_mini 傳送的音訊檔案)
            form = await request.form()
            if "file" not in form:
                raise HTTPException(status_code=400, detail="表單資料中缺少 'file' 欄位")
            file_item = form["file"]
            audio_bytes = await file_item.read()
        else:
            # 支援自訂的 Base64 JSON 格式
            data = await request.json()
            try:
                messages = data["messages"]
                audio_content = messages[0]["content"][0]
                base64_data = audio_content["audio_url"]["url"].split(",")[1]
                audio_bytes = base64.b64decode(base64_data)
            except Exception as e:
                raise HTTPException(status_code=400, detail=f"解析自訂 JSON 失敗: {str(e)}")

        # 執行 Whisper 語音辨識
        result = asr(
            audio_bytes,
            chunk_length_s=30,
            batch_size=8,
            return_timestamps=True,
            generate_kwargs={"language": "english", "task": "transcribe"},
        )

        return {"text": result["text"]}

    except HTTPException as he:
        raise he
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


if __name__ == "__main__":
    import uvicorn

    # 啟動在 4002 埠口
    uvicorn.run(app, host="0.0.0.0", port=4002)