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
import redis.asyncio as redis
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'
)


@asynccontextmanager
async def lifespan(app: FastAPI):
    """应用生命周期管理"""
    # 🔥 启动时初始化Redis
    try:
        redis_client = redis.Redis(
            host=settings.REDIS_HOST,
            port=settings.REDIS_PORT,
            db=settings.REDIS_DB,
            password=settings.REDIS_PASSWORD if settings.REDIS_PASSWORD else None,
            decode_responses=False,  # 保持二进制，手动decode
            max_connections=settings.REDIS_MAX_CONNECTIONS,
        )
        
        # 测试连接
        await redis_client.ping()
        
        # 创建Redis缓存服务
        app.state.redis_cache = RedisCacheService(redis_client)
        app.state.redis_client = redis_client
        
    except Exception as e:
        app.state.redis_cache = None
        app.state.redis_client = None
    
    # 启动时初始化数据库
    try:
        await init_database()
    except Exception as e:
        pass
    
    yield
    
    # 关闭时清理资源
    if hasattr(app.state, 'redis_client') and app.state.redis_client:
        await app.state.redis_client.close()
    
    await close_database()


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=True,
        ssl_keyfile=ssl_keyfile,
        ssl_certfile=ssl_certfile,
    )