app/main.py

"""
FastAPI应用主入口 - SPA 模式
"""

import os
from fastapi import FastAPI, Request
from fastapi.staticfiles import StaticFiles
from fastapi.responses import HTMLResponse, JSONResponse, FileResponse
from fastapi.exceptions import RequestValidationError
from starlette.exceptions import HTTPException as StarletteHTTPException
from starlette.middleware.trustedhost import TrustedHostMiddleware
from starlette.middleware import Middleware
from pydantic import ValidationError
from contextlib import asynccontextmanager
import logging
import logging.config
from pathlib import Path

from app.config.settings import settings
from app.api.routes import router as api_router
from app.plugins.manager import get_plugin_manager, set_app_instance
from app.plugins.loader import plugin_loader
from app.plugins.models import PluginStatus
from app.mcp.server import mcp, get_mcp_app
from app.mcp.plugin_registry import init_plugin_registry, get_plugin_registry

# 配置日志
_LOG_CONFIG = {
    "version": 1,
    "disable_existing_loggers": False,
    "formatters": {
        "default": {"format": "%(asctime)s - %(name)s - %(levelname)s - %(message)s"}
    },
    "handlers": {
        "default": {
            "formatter": "default",
            "class": "logging.StreamHandler",
            "stream": "ext://sys.stdout",
        },
        "file": {
            "formatter": "default",
            "class": "logging.handlers.RotatingFileHandler",
            "filename": "logs/app.log",
            "maxBytes": 10485760,
            "backupCount": 5,
            "encoding": "utf-8",
        },
    },
    "loggers": {
        "app": {
            "handlers": ["default", "file"],
            "level": "DEBUG" if settings.DEBUG else "INFO",
        },
        "plugins": {
            "handlers": ["default", "file"],
            "level": "DEBUG" if settings.DEBUG else "INFO",
        },
    },
}

Path("logs").mkdir(exist_ok=True)

logging.config.dictConfig(_LOG_CONFIG)
logger = logging.getLogger(__name__)

# SPA 模式配置
USE_SPA = os.getenv("USE_SPA", "true").lower() == "true"
frontend_dist = Path(__file__).parent.parent / "dist" / "frontend"
static_dir = Path(__file__).parent / "static"


# 注册插件路由 - 在应用初始化时预注册API路由
def pre_register_plugin_api_routes(app: FastAPI):
    """启动前预注册已启用插件的API路由"""
    manager = get_plugin_manager()
    enabled_plugins = manager.get_enabled_plugins()

    for plugin_info in enabled_plugins:
        try:
            api_router = plugin_loader.create_api_router(
                manager.plugin_directory / plugin_info.metadata.name,
                plugin_info,
            )
            if api_router:
                base_path = f"/plugins/{plugin_info.metadata.name}/api"
                app.include_router(api_router, prefix=base_path)
                api_base_path = f"/api/plugins/{plugin_info.metadata.name}"
                app.include_router(api_router, prefix=api_base_path)
                logger.info(f"预注册插件API路由: {base_path}")
                logger.info(f"预注册插件API兼容路由: {api_base_path}")
        except Exception as e:
            logger.error(f"预注册插件API路由失败 {plugin_info.metadata.name}: {e}")


# 生命周期事件处理器
@asynccontextmanager
async def lifespan(app: FastAPI):
    """应用生命周期管理"""
    logger.info(f"{settings.PROJECT_NAME} v{settings.VERSION} 正在启动...")
    logger.info(f"插件目录: {settings.PLUGINS_DIR}")
    logger.info(f"SPA 模式: {USE_SPA}")

    # 0. 创建沙盒目录
    settings.ensure_sandbox_directory()
    logger.info(f"沙盒目录已创建: {settings.SANDBOX_ROOT}")

    # OCR 已拆分为独立插件，默认按需懒加载，避免启动时占用内存。
    if os.getenv("PRELOAD_OCR", "false").lower() == "true":
        try:
            from app.utils.ocr_engine import preload_ocr_engine
            preload_ocr_engine()
            logger.info("OCR 模型预加载完成")
        except Exception as e:
            logger.warning(f"OCR 模型预加载失败（将在首次使用时加载）: {e}")

    # 1. 仅获取启用插件列表
    manager = get_plugin_manager()
    enabled_plugins = manager.get_enabled_plugins()
    enabled_names = [p.metadata.name for p in enabled_plugins]
    logger.info(f"已启用插件: {enabled_names}")

    # 2. 仅加载已启用插件
    for plugin_info in enabled_plugins:
        try:
            if plugin_loader.load_plugin(plugin_info):
                # 加载后更新 API 路由中的插件实例
                plugin_name = plugin_info.metadata.name
                plugin_instance = plugin_loader.get_plugin_instance(plugin_name)
                if plugin_instance:
                    logger.info(f"插件实例状态: {plugin_name} - enabled={plugin_instance.enabled if hasattr(plugin_instance, 'enabled') else 'N/A'}, file_ops={plugin_instance.file_ops is not None if hasattr(plugin_instance, 'file_ops') else 'N/A'}")
                    try:
                        plugin_path = manager.plugin_directory / plugin_name
                        api_module = plugin_loader._import_plugin_module(plugin_path, "api")
                        if api_module and hasattr(api_module, "set_plugin_instance"):
                            api_module.set_plugin_instance(plugin_instance)
                            logger.info(f"已更新插件API实例: {plugin_name}")
                    except Exception as e:
                        logger.warning(f"更新插件API实例失败 {plugin_name}: {e}")
                logger.info(f"已加载插件: {plugin_info.metadata.name}")
            else:
                logger.warning(f"加载插件失败: {plugin_info.metadata.name}")
        except Exception as e:
            logger.error(f"加载插件失败 {plugin_info.metadata.name}: {e}")

    # 3. 初始化插件 MCP 注册器（仅注册已启用插件）
    registry = init_plugin_registry(mcp)
    tool_count = registry.register_plugin_tools_by_name(settings.PLUGINS_DIR, enabled_names)
    from app.plugins.tool_catalog import tool_catalog_service

    tool_catalog_service.scan_tools(settings.PLUGINS_DIR, manager)
    # MCP 注册完成后同步目录，否则 /api/tools 只能看到 plugin.json tools 字段，无法反映真实可调用工具。
    tool_catalog_service.sync_mcp_tools(registry, manager)
    logger.info(f"已注册 {tool_count} 个插件 MCP 工具")

    # 4. 启动缓存清理任务（如果有 cache）
    if "cache" in enabled_names:
        try:
            from plugins.cache.core import init_cache
            await init_cache()
            logger.info("缓存清理任务已启动")
        except Exception as e:
            logger.warning(f"启动缓存清理失败: {e}")

    # 5. 挂载 MCP 应用
    mcp_app = get_mcp_app()
    
    # 创建 MCP 应用的 ASGI 包装器，处理代理头
    async def mcp_app_wrapper(scope, receive, send):
        """包装 MCP ASGI 应用，处理代理头"""
        if scope["type"] == "http":
            # 获取 headers
            headers = scope.get("headers", [])
            headers_dict = {k.decode(): v.decode() for k, v in headers}
            
            # 记录请求信息（调试用）
            logger.debug(f"MCP Request headers: {headers_dict}")
        
        # 调用原始的 MCP 应用
        await mcp_app(scope, receive, send)
    
    app.mount("/mcp", mcp_app_wrapper)
    logger.info("MCP 服务已挂载到 /mcp 路径")

    # 6. 启动 MCP session manager
    async with mcp.session_manager.run():
        yield

    # 关闭时清理
    logger.info(f"{settings.PROJECT_NAME} 正在关闭...")

    # 关闭缓存清理任务
    if "cache" in enabled_names:
        try:
            from plugins.cache.core import shutdown_cache
            await shutdown_cache()
            logger.info("缓存清理任务已停止")
        except Exception as e:
            logger.warning(f"停止缓存清理失败: {e}")

    # 关闭全局 httpx 客户端
    try:
        from app.utils.http_client import close_http_client
        await close_http_client()
        logger.info("httpx 客户端已关闭")
    except Exception as e:
        logger.warning(f"关闭 httpx 客户端失败: {e}")


# 创建FastAPI应用
app = FastAPI(
    title=settings.PROJECT_NAME,
    version=settings.VERSION,
    debug=settings.DEBUG,
    lifespan=lifespan,
)

# 添加信任所有主机的中间件（用于 HuggingFace Spaces 等代理环境）
# 解决 "Invalid Host header" 错误
# 注意：TrustedHostMiddleware 需要放在最前面，以便正确处理代理头
app.add_middleware(TrustedHostMiddleware, allowed_hosts=["*"])


# 处理代理头的中间件
@app.middleware("http")
async def proxy_headers_middleware(request: Request, call_next):
    """处理代理转发头，支持 HuggingFace Spaces 等代理环境"""
    # 获取原始 Host 头
    host = request.headers.get("host", "")
    
    # 记录请求信息（调试用）
    logger.debug(f"Request Host: {host}, URL: {request.url}")
    
    # 继续处理请求
    response = await call_next(request)
    return response

# 包含API路由
app.include_router(api_router)

if static_dir.exists():
    app.mount("/static", StaticFiles(directory=str(static_dir)), name="static")
    logger.info(f"已挂载平台静态资源目录: {static_dir}")

# 设置全局应用实例，用于动态注册插件路由
set_app_instance(app)

# 预注册所有插件API路由（必须在SPA路由之前）
pre_register_plugin_api_routes(app)

# 预先挂载已启用插件的前端静态文件路由
# FastAPI 按注册顺序匹配路由，前端路由需优先于 SPA catch-all
manager_pre = get_plugin_manager()
enabled_plugins_pre = manager_pre.get_enabled_plugins()
for plugin_info in enabled_plugins_pre:
    try:
        plugin_path = manager_pre.plugin_directory / plugin_info.metadata.name
        frontend_path = plugin_path / "frontend"

        if (frontend_path / "index.html").exists():
            static_path = f"/plugins/{plugin_info.metadata.name}/ui"
            app.mount(
                static_path,
                StaticFiles(directory=str(frontend_path), html=True, check_dir=True),
                name=f"{plugin_info.metadata.name}_static",
            )
            logger.info(f"预先挂载插件前端路由: {static_path}")
    except Exception as e:
        logger.error(f"预先挂载插件前端路由失败 {plugin_info.metadata.name}: {e}")


# 异常处理 - SPA fallback 支持
@app.exception_handler(StarletteHTTPException)
async def http_exception_handler(request: Request, exc: StarletteHTTPException):
    """HTTP异常处理，404 时对非 API/MCP 路径返回 SPA 入口页面。"""
    if exc.status_code == 404 and USE_SPA and frontend_dist.exists():
        path = request.url.path
        # API/MCP/static/plugins 路径返回 JSON 404，不返回 SPA
        if not path.startswith(("/api/", "/mcp/", "/static/", "/plugins/")):
            index_html = frontend_dist / "index.html"
            if index_html.exists():
                return FileResponse(str(index_html))
    return JSONResponse(status_code=exc.status_code, content={"detail": exc.detail})


@app.exception_handler(ValidationError)
async def pydantic_validation_error_handler(request: Request, exc: ValidationError):
    """Pydantic数据验证错误处理"""
    errors = []
    for error in exc.errors():
        errors.append(
            {
                "field": ".".join(str(x) for x in error.get("loc", [])),
                "message": error.get("msg", "未知错误"),
            }
        )
    return JSONResponse(
        status_code=422, content={"detail": "数据验证失败", "errors": errors}
    )


@app.exception_handler(RequestValidationError)
async def validation_exception_handler(request: Request, exc: RequestValidationError):
    return JSONResponse(status_code=422, content={"detail": exc.errors()})


@app.exception_handler(Exception)
async def general_exception_handler(request: Request, exc: Exception):
    """通用异常处理"""
    import traceback
    logger.error(f"未处理的异常: {exc}", exc_info=True)
    traceback.print_exc()

    return JSONResponse(
        status_code=500,
        content={
            "detail": "内部服务器错误",
            "error": str(exc) if settings.DEBUG else "未知错误",
        },
    )


@app.get("/favicon.ico", include_in_schema=False)
async def favicon():
    """返回空响应避免404错误"""
    return HTMLResponse(status_code=204)


# SPA 模式：通过 404 异常处理器（见上方 http_exception_handler）实现 fallback
# 不再使用 catch-all 路由，因为它会拦截 GET /mcp/ 等挂载路由
if USE_SPA and frontend_dist.exists():
    # 挂载 assets 目录
    assets_dir = frontend_dist / "assets"
    if assets_dir.exists():
        app.mount("/assets", StaticFiles(directory=str(assets_dir)), name="assets")
        logger.info(f"已挂载前端资源目录: {assets_dir}")

    # SPA 入口已通过 http_exception_handler 处理 404 fallback
    # 不再使用 catch-all 路由（@app.get("/{path:path}")），
    # 因为 catch-all 会拦截 app.mount() 挂载的 MCP 路由（如 GET /mcp/）
else:
    logger.warning(f"SPA 模式未启用或前端未构建。USE_SPA={USE_SPA}, frontend_dist exists={frontend_dist.exists()}")


if __name__ == "__main__":
    import uvicorn

    uvicorn.run(
        "app.main:app",
        host="0.0.0.0",
        port=7860,
        reload=settings.DEBUG,
        proxy_headers=True,  # 启用代理头支持
        forwarded_allow_ips="*",  # 允许所有 IP 的转发头（用于 HuggingFace Spaces 等代理环境）
    )