SupervisorAI/app.py

"""
FastAPI主应用
将原来的main.py重命名为app.py
"""

import time
from contextlib import asynccontextmanager
from fastapi import FastAPI, Request
from fastapi.exceptions import RequestValidationError
from fastapi.middleware.cors import CORSMiddleware
from fastapi.middleware.trustedhost import TrustedHostMiddleware
from fastapi.responses import JSONResponse
from fastapi.openapi.docs import (
    get_swagger_ui_html,
    get_swagger_ui_oauth2_redirect_html,
    get_redoc_html,
)
from fastapi.staticfiles import StaticFiles

from api.routes import face_features
from api.routes.algorithm_router import router as algorithm_router, sync_videofacebiz_params, \
    sync_videofacebiz_blacklist, sync_videofaceprisonbiz_params, sync_videofaceprisonbiz_blacklist
from api.errors import (
    APIError,
    validation_exception_handler,
    api_error_handler,
    generic_exception_handler
)
from config import settings
from database.connection import init_database
from database.connection import db_manager
from rtsp.service import rtsp_server


# 生命周期管理
@asynccontextmanager
async def lifespan(app: FastAPI):
    """
    应用生命周期管理

    - 启动时：初始化数据库
    - 关闭时：清理资源
    """
    # 启动时
    print("🚀 start algorithm service...")
    print(f"📊 db: {settings.DATABASE_NAME}")
    print(f"🔧 debug mode: {settings.DEBUG}")

    # 初始化数据库
    init_database()

    # 数据库健康检查
    if db_manager.health_check():
        print("✅ 数据库连接正常")
    else:
        print("❌ 数据库连接失败")

    # 启动 RTSP 服务（如果启用）
    if settings.RTSP_ENABLED:
        print("📹 启动 RTSP 服务...")
        rtsp_server.start()
        # 将 RTSP 服务实例保存到应用状态
        app.state.rtsp_server = rtsp_server
        
        # 自动同步VideoFaceBiz参数和黑名单
        print("🔄 自动同步VideoFaceBiz参数和黑名单...")
        try:
            params_updated = sync_videofacebiz_params()
            blacklist_loaded = sync_videofacebiz_blacklist()
            print(f"✅ 自动同步完成 - 参数更新: {params_updated}个, 黑名单加载: {blacklist_loaded}个")
        except Exception as e:
            print(f"⚠️  自动同步失败: {e}")
    else:
        print("⚠️  RTSP 服务未启用")

    # 自动同步VideoFacePrisonBiz参数和黑名单
    print("🔄 自动同步VideoFacePrisonBiz参数和黑名单...")
    try:
        params_updated = sync_videofaceprisonbiz_params()
        blacklist_loaded = sync_videofaceprisonbiz_blacklist()
        print(f"✅ 自动同步完成 - 参数更新: {params_updated}个, 黑名单加载: {blacklist_loaded}个")
    except Exception as e:
        print(f"⚠️  自动同步失败: {e}")

    yield

    # 关闭时
    print("🛑 algorithm service stopped...")
    
    # 停止 RTSP 服务
    if settings.RTSP_ENABLED:
        print("🛑 停止 RTSP 服务...")
        rtsp_server.stop()
    
    db_manager.close()


# 创建FastAPI应用
app = FastAPI(
    title=settings.PROJECT_NAME,
    version=settings.PROJECT_VERSION,
    description=settings.PROJECT_DESCRIPTION,
    openapi_url=f"{settings.API_V1_PREFIX}/openapi.json",
    docs_url=None,  # 自定义docs
    redoc_url=None,  # 自定义redoc
    lifespan=lifespan
)


# 自定义API文档页面
@app.get("/docs", include_in_schema=False)
async def custom_swagger_ui_html():
    return get_swagger_ui_html(
        openapi_url=app.openapi_url,
        title=f"{app.title} - Swagger UI",
        oauth2_redirect_url=app.swagger_ui_oauth2_redirect_url,
        swagger_js_url="https://unpkg.com/swagger-ui-dist@5/swagger-ui-bundle.js",
        swagger_css_url="https://unpkg.com/swagger-ui-dist@5/swagger-ui.css",
    )


@app.get(app.swagger_ui_oauth2_redirect_url, include_in_schema=False)
async def swagger_ui_redirect():
    return get_swagger_ui_oauth2_redirect_html()


@app.get("/redoc", include_in_schema=False)
async def redoc_html():
    return get_redoc_html(
        openapi_url=app.openapi_url,
        title=f"{app.title} - ReDoc",
        redoc_js_url="https://unpkg.com/redoc@next/bundles/redoc.standalone.js",
    )


# 中间件配置
app.add_middleware(
    CORSMiddleware,
    allow_origins=settings.BACKEND_CORS_ORIGINS,
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

app.add_middleware(
    TrustedHostMiddleware,
    allowed_hosts=["*"] if settings.DEBUG else ["localhost", "127.0.0.1", "0.0.0.0"]
)


# 请求计时中间件
@app.middleware("http")
async def add_process_time_header(request: Request, call_next):
    """
    添加请求处理时间头
    """
    start_time = time.time()
    response = await call_next(request)
    process_time = time.time() - start_time
    response.headers["X-Process-Time"] = str(process_time)
    return response


# 异常处理器
app.add_exception_handler(RequestValidationError, validation_exception_handler)
app.add_exception_handler(APIError, api_error_handler)
app.add_exception_handler(Exception, generic_exception_handler)


# 根路由
@app.get("/")
async def root():
    """
    根路径
    """
    return {
        "message": "algorithm service",
        "version": settings.PROJECT_VERSION,
        "docs": "/docs",
        "api_prefix": settings.API_V1_PREFIX
    }


@app.get("/health")
async def health_check():
    """
    健康检查端点
    """
    # 检查数据库连接
    db_healthy = db_manager.health_check()

    return {
        "status": "healthy" if db_healthy else "unhealthy",
        "database": "connected" if db_healthy else "disconnected",
        "timestamp": time.time()
    }


# 注册API路由
app.include_router(
    face_features.router,
    prefix=settings.API_V1_PREFIX
)

app.include_router(
    algorithm_router,
    prefix=settings.API_V1_PREFIX
)


# 自定义404处理器
@app.exception_handler(404)
async def not_found_handler(request: Request, exc):
    """
    自定义404错误处理器
    """
    return JSONResponse(
        status_code=404,
        content={
            "error": {
                "code": "NOT_FOUND",
                "message": f"请求的资源不存在: {request.url.path}"
            }
        }
    )


# 导出应用实例
__all__ = ["app"]