#!/usr/bin/env python3
"""ai-image 入口：REST + Web 控制台 + MCP。

启动方式：
    python main.py              # Web + REST + MCP(http) 同端口
    python main.py --stdio      # 仅 MCP stdio（给 MCP 客户端当子进程用）
"""

from __future__ import annotations

import argparse
import contextlib

from fastapi import FastAPI
from fastapi.staticfiles import StaticFiles

from app import config, db
from app.mcp_server import build_http_app
from app.routers import api_gallery, api_generate, api_models, api_providers, api_skills

config.ensure_dirs()
db.init_db()

mcp_app = build_http_app(path="/mcp")


@contextlib.asynccontextmanager
async def lifespan(_: FastAPI):
    async with mcp_app.lifespan(mcp_app):
        yield


app = FastAPI(title="ai-image", version="2.0.0", lifespan=lifespan)

app.include_router(api_providers.router)
app.include_router(api_models.router)
app.include_router(api_generate.router)
app.include_router(api_gallery.router)
app.include_router(api_skills.router)


@app.get("/health")
def health():
    return {"status": "ok", "output_dir": str(config.OUTPUT_DIR)}


app.mount("/output", StaticFiles(directory=config.OUTPUT_DIR), name="output")
app.mount("/", StaticFiles(directory=config.WEB_DIR, html=True), name="web")


class RootDispatcher:
    """按路径前缀把请求分发给 MCP 或 Web 应用。

    Web 静态资源挂在根路径，MCP 需要独占 /mcp，因此不能靠 mount 处理。
    """

    def __init__(self, web_app, mcp_asgi, prefix: str = "/mcp"):
        self.web_app = web_app
        self.mcp_asgi = mcp_asgi
        self.prefix = prefix

    async def __call__(self, scope, receive, send):
        path = scope.get("path", "")
        if scope["type"] == "http" and (path == self.prefix or path.startswith(self.prefix + "/")):
            await self.mcp_asgi(scope, receive, send)
        else:
            await self.web_app(scope, receive, send)


dispatcher = RootDispatcher(app, mcp_app)


def main() -> None:
    parser = argparse.ArgumentParser(description="ai-image 服务")
    parser.add_argument("--stdio", action="store_true", help="以 MCP stdio 模式运行")
    parser.add_argument("--host", default=config.HOST)
    parser.add_argument("--port", type=int, default=config.PORT)
    args = parser.parse_args()

    if args.stdio:
        from app.mcp_server import mcp

        mcp.run()
        return

    import uvicorn

    uvicorn.run(dispatcher, host=args.host, port=args.port)


if __name__ == "__main__":
    main()