Tolan/backend/main.py

"""
主应用入口
"""
import logging
from typing import Optional
from fastapi import FastAPI, WebSocket, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from contextlib import asynccontextmanager
from config import settings, get_dashscope_domain
from websocket_handler import WebSocketHandler
from database import init_database, close_database
from services.redis_cache_service import RedisCacheService
from services.user_service import user_service

# 配置日志
logging.basicConfig(
    level=getattr(logging, settings.LOG_LEVEL.upper(), logging.INFO),
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    datefmt='%Y-%m-%d %H:%M:%S'
)

logger = logging.getLogger(__name__)


@asynccontextmanager
async def lifespan(app: FastAPI):
    """应用生命周期管理"""
    # 启动时初始化数据库（包含MySQL和Redis，自动启动健康检查任务）
    try:
        await init_database()
        
        # 获取Redis客户端（已配置健康检查），用于验证连通性
        from database import db_manager
        await db_manager.get_redis()
        
        # 创建Redis缓存服务，按需获取最新连接
        app.state.redis_cache = RedisCacheService(db_manager.get_redis)
        
    except Exception as e:
        logger.error(f"❌ 数据库初始化失败: {e}")
        app.state.redis_cache = None
    
    # 🔥 启动HTTP连接保活任务（保持连接活跃，避免长时间不交互后响应变慢）
    try:
        from http_pool import http_pool
        from ai_service import AIService
        
        # 创建AIService实例用于保活（共享连接池）
        ai_service = AIService()
        http_pool.start_keepalive(ai_service=ai_service)
    except Exception as e:
        logger.warning(f"⚠️ HTTP连接保活任务启动失败: {e}")
    
    yield
    
    # 关闭时清理资源
    await close_database()
    
    # 🔥 关闭HTTP连接池（类似SQLAlchemy的Engine.dispose()）
    from http_pool import http_pool
    http_pool.close()


app = FastAPI(
    title="Tolan Backend API",
    lifespan=lifespan
)

# 配置CORS
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # 生产环境应该限制具体域名
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)


@app.get("/")
async def root():
    """健康检查"""
    return {"status": "ok", "message": "Tolan Backend API"}


class CreateUserRequest(BaseModel):
    """创建用户请求"""
    username: str = ""
    bot_name: Optional[str] = None


class UpdateUsernameRequest(BaseModel):
    """更新用户名请求"""
    username: str


class UpdateBotNameRequest(BaseModel):
    """更新机器人名字请求"""
    bot_name: str


@app.post("/api/users/create")
async def create_user(request: CreateUserRequest):
    """
    创建游客用户
    
    Returns:
        {
            "user_id": "uuid",
            "username": "",
            "user_type": "guest",
            "bot_name": "机器人名字（可选）"
        }
    """
    try:
        user = await user_service.create_guest_user(
            username=request.username,
            bot_name=request.bot_name
        )
        return {
            "success": True,
            "data": user
        }
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"创建用户失败: {str(e)}")


@app.get("/api/users/{user_id}")
async def get_user(user_id: str):
    """
    获取用户信息
    
    Args:
        user_id: 用户ID
    
    Returns:
        {
            "success": true,
            "data": {
                "user_id": "uuid",
                "username": "",
                "user_type": "guest",
                "bot_name": "机器人名字（可选）",
                ...
            }
        }
    """
    try:
        user = await user_service.get_user(user_id)
        if user:
            return {
                "success": True,
                "data": user
            }
        else:
            raise HTTPException(status_code=404, detail="用户不存在")
    except HTTPException:
        raise
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"获取用户信息失败: {str(e)}")


@app.put("/api/users/{user_id}/username")
async def update_username(user_id: str, request: UpdateUsernameRequest):
    """
    更新用户名
    
    Args:
        user_id: 用户ID
        request: 更新请求（包含username）
    
    Returns:
        {
            "success": true,
            "message": "更新成功"
        }
    """
    try:
        success = await user_service.update_username(user_id, request.username)
        if success:
            return {
                "success": True,
                "message": "用户名更新成功"
            }
        else:
            raise HTTPException(status_code=500, detail="更新用户名失败")
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"更新用户名失败: {str(e)}")


@app.put("/api/users/{user_id}/bot-name")
async def update_bot_name(user_id: str, request: UpdateBotNameRequest):
    """
    更新用户给机器人取的名字
    
    Args:
        user_id: 用户ID
        request: 更新请求（包含bot_name）
    
    Returns:
        {
            "success": true,
            "message": "更新成功"
        }
    """
    try:
        success = await user_service.update_bot_name(user_id, request.bot_name)
        if success:
            return {
                "success": True,
                "message": "机器人名字更新成功"
            }
        else:
            raise HTTPException(status_code=500, detail="更新机器人名字失败")
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"更新机器人名字失败: {str(e)}")


@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
    """WebSocket端点"""
    # 从查询参数获取user_id（tolan_user_id）
    user_id = websocket.query_params.get("user_id")
    
    # 记录连接信息
    import logging
    logger = logging.getLogger(__name__)
    logger.info(f"🔌 WebSocket连接: user_id={user_id}, query_params={dict(websocket.query_params)}")
    
    # 🔥 传入Redis缓存服务和user_id
    redis_cache = getattr(app.state, 'redis_cache', None)
    handler = WebSocketHandler(websocket, user_id=user_id, redis_cache=redis_cache)
    await handler.handle_connection()


if __name__ == "__main__":
    import uvicorn
    import os
    import socket
    from pathlib import Path
    
    # 🔥 DNS预解析（减少首次连接延迟）
    dashscope_domain = get_dashscope_domain()
    try:
        socket.getaddrinfo(dashscope_domain, 443)
    except Exception as e:
        pass
    
    # SSL证书路径配置
    # 优先使用环境变量配置的证书路径（生产环境）
    ssl_keyfile = settings.SSL_KEYFILE
    ssl_certfile = settings.SSL_CERTFILE
    
    # 如果环境变量未配置，使用本地开发证书
    if not ssl_keyfile or not ssl_certfile:
        ssl_dir = Path(__file__).parent.parent / "ssl"
        ssl_keyfile = str(ssl_dir / "key.pem")
        ssl_certfile = str(ssl_dir / "cert.pem")
    
    # 检查证书是否存在
    if ssl_keyfile and ssl_certfile:
        if not os.path.exists(ssl_keyfile) or not os.path.exists(ssl_certfile):
            ssl_keyfile = None
            ssl_certfile = None
    
    uvicorn.run(
        "main:app",
        host=settings.host,
        port=settings.port,
        reload=settings.DEBUG,
        ssl_keyfile=ssl_keyfile,
        ssl_certfile=ssl_certfile,
    )