# main.py
import asyncio
import uvicorn
from fastapi import FastAPI, WebSocket, WebSocketDisconnect
from starlette.middleware.cors import CORSMiddleware
# 导入抽离的消息管理模块
from ws_message_manager import ws_queue_manager, ServerMsgType
import logging
from frontend_ws import frontend_websocket_handler
from asr_client import init_asr_pool,  close_asr_pool
from contextlib import asynccontextmanager

# 初始化日志配置（全局生效）
logging.basicConfig(
    level=logging.INFO,  # 日志级别：DEBUG/INFO/WARNING/ERROR
    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",  # 日志格式
    handlers=[
        logging.StreamHandler(),  # 输出到控制台
        logging.FileHandler("ws_server.log", encoding="utf-8")  # 输出到文件（可选）
    ]
)

# 创建日志实例（后续所有logger调用都用这个实例）
logger = logging.getLogger(__name__)  # __name__ 是当前模块名，便于区分日志来源


# FastAPI 启动时初始化 ASR 连接池
@asynccontextmanager
async def lifespan(app: FastAPI):
    # 启动时执行（原 startup 逻辑）
    print(' FastAPI 启动时初始化 ASR 连接池')
    await init_asr_pool()
    yield  # 应用运行中
    # 关闭时执行（可选，比如清理连接池）
    print("应用关闭，开始清理 ASR 连接池...")
    await close_asr_pool()

    # 这里可以添加连接池关闭逻辑（如关闭所有 ASR 连接）


# 初始化FastAPI应用
app = FastAPI(
    title="语音交互网关",
    description="整合前端/ASR/大模型/TTS的异步网关服务",
    version="1.0",
    lifespan=lifespan # 绑定生命周期
)

# 跨域配置
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)


async def process_audio(websocket: WebSocket, audio_data: bytes):
    """处理音频数据，生成多类型消息并推送"""
    try:
        # 1. 模拟ASR识别（生成文字）
        text_result = f"识别结果：音频长度 {len(audio_data)} bytes"
        await ws_queue_manager.send_message(websocket, ServerMsgType.TEXT, text_result)

        # 2. 模拟TTS生成语音（二进制数据）
        voice_data = b"xx"  # 替换为真实TTS输出
        await ws_queue_manager.send_message(websocket, ServerMsgType.VOICE, voice_data)

        # 3. 模拟数字人动画信息
        animation_data = {
            "action": "mouth_move",
            "speed": 1.2,
            "duration": 1000
        }
        await ws_queue_manager.send_message(websocket, ServerMsgType.ANIMATION, animation_data)

        # 4. 模拟前端动作指令
        action_data = {
            "type": "show_loading",
            "status": False
        }
        await ws_queue_manager.send_message(websocket, ServerMsgType.ACTION, action_data)

    except Exception as e:
        # 推送错误信息
        await ws_queue_manager.send_message(websocket, ServerMsgType.ERROR, str(e))


@app.websocket("/ws/audio")
async def websocket_audio(websocket: WebSocket):
    # 完全委托给frontend_websocket_handler处理
    await frontend_websocket_handler(websocket)


# 前端WebSocket路由
# @app.websocket("/ws/audio")
# async def websocket_audio(websocket: WebSocket):
#     # 1. 接受连接并创建专属队列
#     await websocket.accept()
#     await ws_queue_manager.create_queue(websocket)
#
#     try:
#         # 2. 循环接收前端音频数据
#         while True:
#             # 接收二进制音频数据（前端发送的麦克风数据）
#             audio_data = await websocket.receive_bytes()
#             # 异步处理音频（不阻塞接收）
#             asyncio.create_task(process_audio(websocket, audio_data))
#
#     except WebSocketDisconnect:
#         logger.info(f"连接 {id(websocket)} 主动断开")
#     except Exception as e:
#         logger.error(f"连接异常: {e}")
#     finally:
#         # 3. 清理队列和协程
#         await ws_queue_manager.close_queue(websocket)


# 健康检查接口
@app.get("/health")
async def health_check():
    return {
        "status": "ok",
        "service": "voice_gateway",
        "active_connections": len(ws_queue_manager.queue_map)
    }


if __name__ == "__main__":
    uvicorn.run(
        "main:app",
        host="0.0.0.0",
        port=8000,
        reload=True,  # 生产环境关闭
        workers=1  # 异步框架单worker足够，多worker需分布式队列
    )