# FunASR STT 服务 - 本地 Python 部署（基于 FunAudioLLM/SenseVoiceSmall）
# 参考: https://modelscope.cn/models/iic/SenseVoiceSmall

from funasr import AutoModel
from funasr.utils.postprocess_utils import rich_transcription_postprocess
import uvicorn
from fastapi import FastAPI, UploadFile, File
from fastapi.responses import JSONResponse
import os
import logging

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# 配置
MODEL_ID = os.environ.get("MODEL_ID", "FunAudioLLM/SenseVoiceSmall")
VAD_MODEL = os.environ.get("VAD_MODEL", "fsmn-vad")
VAD_MODEL_REVISION = os.environ.get("VAD_MODEL_REVISION", "v2.0")
PORT = int(os.environ.get("PORT", "8080"))

# 初始化模型
logger.info(f"加载 STT 模型: {MODEL_ID}")
logger.info(f"加载 VAD 模型: {VAD_MODEL}")

model = AutoModel(
    model=MODEL_ID,
    model_revision="master",
    vad_model=VAD_MODEL,
    vad_model_revision=VAD_MODEL_REVISION,
    disable_update=True,
)

app = FastAPI(title="FunASR STT Service")

@app.get("/health")
async def health():
    return {"status": "healthy", "model": MODEL_ID}

@app.post("/stream")
async def stream_asr(file: UploadFile = File(...)):
    """流式语音识别"""
    try:
        content = await file.read()
        
        result = model.generate(
            input=content,
            batch_size_s=300,
            hotword="阿里达摩院",
            postprocessor=rich_transcription_postprocess,
        )
        
        text = result[0]["text"] if result else ""
        return JSONResponse({"text": text, "status": "success"})
        
    except Exception as e:
        logger.error(f"识别失败: {e}")
        return JSONResponse({"error": str(e)}, status_code=500)

@app.post("/recognize")
async def recognize(file: UploadFile = File(...)):
    """非流式语音识别"""
    try:
        content = await file.read()
        result = model.generate(
            input=content,
            cache_length=0,
            hotword="阿里达摩院",
            use_itn=True,
        )
        
        text = result[0]["text"] if result else ""
        return JSONResponse({"text": text, "result": result})
        
    except Exception as e:
        logger.error(f"识别失败: {e}")
        return JSONResponse({"error": str(e)}, status_code=500)

if __name__ == "__main__":
    logger.info(f"🚀 启动 FunASR STT 服务")
    logger.info(f"   模型: {MODEL_ID}")
    logger.info(f"   端口: {PORT}")
    logger.info(f"   地址: http://localhost:{PORT}")
    uvicorn.run(app, host="0.0.0.0", port=PORT)