"""
FunASR 本地 STT 服务 - 简化版本
只使用 SenseVoiceSmall 模型（内置 VAD）
"""
import os
import sys
import json
import logging
import tempfile
from pathlib import Path

# 添加虚拟环境路径
sys.path.insert(0, '/Users/leo/s2s-env/lib/python3.12/site-packages')

from fastapi import FastAPI, UploadFile, File, HTTPException
from fastapi.responses import JSONResponse, HTMLResponse
import uvicorn

# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# ==================== 配置 ====================
MODEL_ID = "iic/SenseVoiceSmall"
PORT = int(os.environ.get("PORT", "8080"))
HOST = "0.0.0.0"

# ==================== 模型加载 ====================
recognizer = None
model_initialized = False
model_loading = False
loading_error = None

def load_models():
    """加载 FunASR 模型"""
    global recognizer, model_initialized, model_loading, loading_error
    
    if model_loading:
        return "模型正在加载中..."
    
    model_loading = True
    loading_error = None
    
    try:
        from funasr import AutoModel
        
        logger.info("🔄 正在加载 FunASR 模型...")
        logger.info(f"   模型: {MODEL_ID}")
        
        # 只加载 SenseVoiceSmall（内置 VAD）
        recognizer = AutoModel(
            model=MODEL_ID,
            model_revision="master",
            disable_update=True,
        )
        
        logger.info("✅ 模型加载完成")
        model_initialized = True
        
    except Exception as e:
        loading_error = f"模型加载失败: {e}"
        logger.error(loading_error)
    finally:
        model_loading = False
    
    return loading_error

# ==================== FastAPI 应用 ====================
app = FastAPI(
    title="FunASR STT Service",
    description="基于 FunASR SenseVoiceSmall 的本地语音转文字服务",
    version="1.0.0"
)

@app.on_event("startup")
async def startup_event():
    """启动时加载模型"""
    logger.info("🚀 启动 FunASR STT 服务...")
    error = load_models()
    if error:
        logger.warning(f"⚠️  {error}")
    else:
        logger.info("✅ 服务启动完成")

@app.get("/health")
async def health_check():
    """健康检查端点"""
    status = {
        "status": "healthy" if model_initialized else "unhealthy",
        "model_loaded": model_initialized,
        "loading": model_loading,
        "error": loading_error,
        "endpoint": f"http://{HOST}:{PORT}"
    }
    return JSONResponse(content=status)

@app.post("/stream")
async def stream_asr(file: UploadFile = File(...)):
    """
    流式语音识别端点
    """
    if not model_initialized:
        if model_loading:
            raise HTTPException(status_code=503, detail="模型正在加载中，请稍后")
        else:
            raise HTTPException(status_code=503, detail="模型未加载")
    
    try:
        # 读取音频数据
        audio_data = await file.read()
        
        # 保存临时文件
        with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as tmp:
            tmp.write(audio_data)
            tmp_path = tmp.name
        
        try:
            # 使用 FunASR 进行识别
            result = recognizer.generate(input=tmp_path)
            
            # 解析结果
            text = ""
            if result and len(result) > 0:
                text = result[0].get("text", "")
                # 清理文本，移除特殊标签
                import re
                text = re.sub(r'<\|.*?\|>', '', text).strip()
                
            return JSONResponse({
                "text": text,
                "status": "success"
            })
            
        finally:
            # 清理临时文件
            os.unlink(tmp_path)
            
    except Exception as e:
        logger.error(f"识别失败: {e}")
        raise HTTPException(status_code=500, detail=str(e))

@app.post("/recognize")
async def recognize(file: UploadFile = File(...)):
    """非流式语音识别端点"""
    if not model_initialized:
        if model_loading:
            raise HTTPException(status_code=503, detail="模型正在加载中")
        else:
            raise HTTPException(status_code=503, detail="模型未加载")
    
    try:
        audio_data = await file.read()
        
        with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as tmp:
            tmp.write(audio_data)
            tmp_path = tmp.name
        
        try:
            result = recognizer.generate(input=tmp_path)
            
            text = ""
            if result and len(result) > 0:
                text = result[0].get("text", "")
                import re
                text = re.sub(r'<\|.*?\|>', '', text).strip()
                
            return JSONResponse({
                "text": text,
                "result": result
            })
            
        finally:
            os.unlink(tmp_path)
            
    except Exception as e:
        logger.error(f"识别失败: {e}")
        raise HTTPException(status_code=500, detail=str(e))

@app.get("/", response_class=HTMLResponse)
async def root():
    """首页"""
    return f"""
    <!DOCTYPE html>
    <html>
    <head>
        <meta charset="UTF-8">
        <title>FunASR STT Service</title>
        <style>
            body {{ font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', sans-serif; background: #1a1a2e; color: #eee; padding: 40px; }}
            h1 {{ color: #00ffaa; }}
            .status {{ padding: 20px; border-radius: 8px; margin: 20px 0; }}
            .ok {{ background: #1e3a2e; border: 1px solid #00ffaa; }}
            .error {{ background: #3a1e1e; border: 1px solid #ff4444; }}
            code {{ background: #2a2a4a; padding: 2px 6px; border-radius: 4px; }}
        </style>
    </head>
    <body>
        <h1>🎤 FunASR STT Service</h1>
        <div class="status {'ok' if model_initialized else 'error'}">
            <h3>状态: {'✅ 已就绪' if model_initialized else '❌ 未就绪'}</h3>
            <p>模型: SenseVoiceSmall</p>
            <p>端口: {PORT}</p>
            {'<p>错误: ' + str(loading_error) + '</p>' if loading_error else ''}
        </div>
        <h3>API 端点</h3>
        <ul>
            <li><code>POST /stream</code> - 流式语音识别</li>
            <li><code>POST /recognize</code> - 非流式语音识别</li>
            <li><code>GET /health</code> - 健康检查</li>
        </ul>
    </body>
    </html>
    """

if __name__ == "__main__":
    logger.info(f"🚀 启动 FunASR STT 服务，端口: {PORT}")
    uvicorn.run(app, host=HOST, port=PORT)