Tolan/backend/main.py

"""
Tolan语音助手后端主应用
"""
import time
from fastapi import FastAPI, WebSocket, WebSocketDisconnect
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse
from loguru import logger
import uvicorn

from config import settings
from connection_pool import connection_pool
from websocket_handler import connection_manager
from performance import performance_monitor

# 创建FastAPI应用
app = FastAPI(
    title=settings.APP_NAME,
    version=settings.APP_VERSION,
    description="实时语音交互API服务"
)

# 配置CORS
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"] if settings.CORS_ALLOW_ALL else settings.CORS_ORIGINS,
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# 配置日志
logger.add(
    f"{settings.LOG_DIR}/app_{{time}}.log",
    rotation=settings.LOG_ROTATION,
    retention=settings.LOG_RETENTION,
    level=settings.LOG_LEVEL,
    format="{time:YYYY-MM-DD HH:mm:ss} | {level} | {message}"
)


@app.on_event("startup")
async def startup_event():
    """应用启动事件"""
    logger.info("=" * 60)
    logger.info(f"{settings.APP_NAME} v{settings.APP_VERSION} 正在启动...")
    logger.info("=" * 60)
    
    # 初始化连接池
    try:
        await connection_pool.initialize()
        logger.success("✓ 连接池初始化成功")
    except Exception as e:
        logger.error(f"✗ 连接池初始化失败: {e}")
    
    logger.info("=" * 60)
    logger.success(f"应用启动完成！")
    logger.info(f"服务地址: http://{settings.HOST}:{settings.PORT}")
    logger.info(f"API文档: http://{settings.HOST}:{settings.PORT}/docs")
    logger.info(f"WebSocket: ws://{settings.HOST}:{settings.PORT}/ws/realtime")
    logger.info("=" * 60)


@app.on_event("shutdown")
async def shutdown_event():
    """应用关闭事件"""
    logger.info("应用正在关闭...")
    
    # 关闭连接池
    try:
        await connection_pool.close_all()
        logger.success("✓ 连接池已关闭")
    except Exception as e:
        logger.error(f"✗ 关闭连接池失败: {e}")
    
    logger.info("应用已关闭")


@app.get("/")
async def root():
    """根路径"""
    return {
        "app": settings.APP_NAME,
        "version": settings.APP_VERSION,
        "status": "running",
        "endpoints": {
            "docs": "/docs",
            "websocket": "/ws/realtime",
            "health": "/health"
        }
    }


@app.get("/health")
async def health_check():
    """健康检查"""
    return {
        "status": "healthy",
        "timestamp": int(time.time() * 1000),
        "connection_pool": {
            "initialized": connection_pool.is_initialized,
            "pool_size": connection_pool.pool_size,
            "idle": connection_pool.idle_connections.qsize(),
            "active": len(connection_pool.active_connections)
        },
        "active_sessions": len(connection_manager.active_connections)
    }


@app.get("/stats")
async def get_stats():
    """获取性能统计"""
    return {
        "status": "ok",
        "timestamp": int(time.time() * 1000),
        "performance": performance_monitor.get_stats(),
        "connection_pool": {
            "pool_size": connection_pool.pool_size,
            "idle": connection_pool.idle_connections.qsize(),
            "active": len(connection_pool.active_connections)
        },
        "sessions": {
            "active": len(connection_manager.active_connections),
            "total": performance_monitor.session_count
        }
    }


@app.websocket("/ws/realtime")
async def websocket_endpoint(websocket: WebSocket):
    """WebSocket实时通信端点"""
    session_id = f"session_{int(time.time() * 1000)}"
    session = None
    
    try:
        # 建立连接
        session = await connection_manager.connect(websocket, session_id)
        logger.info(f"WebSocket连接建立: {session_id}")
        
        # 启动会话
        if not await session.start():
            logger.error(f"会话启动失败: {session_id}")
            return
        
        # 处理流式数据
        await session.process_stream()
        
    except WebSocketDisconnect:
        logger.info(f"客户端主动断开: {session_id}")
        
    except Exception as e:
        logger.error(f"WebSocket错误 [{session_id}]: {e}")
        
        # 发送错误消息
        if session:
            try:
                await session.send_error(str(e))
            except:
                pass
    
    finally:
        # 清理连接
        await connection_manager.disconnect(session_id)


@app.exception_handler(Exception)
async def global_exception_handler(request, exc):
    """全局异常处理"""
    logger.error(f"全局异常: {exc}")
    return JSONResponse(
        status_code=500,
        content={
            "error": "Internal Server Error",
            "message": str(exc)
        }
    )


def start_server():
    """启动服务器"""
    # 检查是否使用SSL
    ssl_keyfile = settings.SSL_KEYFILE
    ssl_certfile = settings.SSL_CERTFILE
    
    import os
    use_ssl = (
        os.path.exists(ssl_keyfile) and 
        os.path.exists(ssl_certfile)
    )
    
    if use_ssl:
        logger.info(f"使用SSL证书启动")
        uvicorn.run(
            "main:app",
            host=settings.HOST,
            port=settings.PORT,
            reload=settings.DEBUG,
            ssl_keyfile=ssl_keyfile,
            ssl_certfile=ssl_certfile,
            log_level=settings.LOG_LEVEL.lower()
        )
    else:
        logger.warning("未找到SSL证书，使用HTTP启动")
        uvicorn.run(
            "main:app",
            host=settings.HOST,
            port=settings.PORT,
            reload=settings.DEBUG,
            log_level=settings.LOG_LEVEL.lower()
        )


if __name__ == "__main__":
    start_server()