from __future__ import annotations

from typing import TYPE_CHECKING

from niannian.mcp.server.tools import (
    register_external_data_tools,
    register_external_query_tools,
    register_memory_tools,
    register_recall_tools,
    register_skill_tools,
    register_system_tools,
    register_vault_tools,
)
from niannian.mcp.service import NianNianService, create_service

if TYPE_CHECKING:
    from mcp.server.fastmcp import FastMCP


def build_fastmcp(svc: NianNianService) -> FastMCP:
    from mcp.server.fastmcp import FastMCP

    mcp = FastMCP(
        "念念",
        stateless_http=True,
        streamable_http_path="/",
    )
    register_memory_tools(mcp, svc)
    register_recall_tools(mcp, svc)
    register_external_data_tools(mcp, svc)
    register_external_query_tools(mcp, svc)
    register_skill_tools(mcp, svc)
    register_system_tools(mcp, svc)
    register_vault_tools(mcp, svc)
    return mcp


def create_mcp_app(svc: NianNianService | None = None):
    """创建 Streamable MCP Starlette 子应用（挂载路径 /mcp）。"""
    service = svc or create_service()
    mcp = build_fastmcp(service)
    return mcp, mcp.streamable_http_app()


def mcp_lifespan(mcp: FastMCP):
    """FastAPI lifespan — 初始化 Streamable MCP session manager。"""
    from contextlib import asynccontextmanager

    @asynccontextmanager
    async def _lifespan(app):
        mcp.streamable_http_app()
        async with mcp.session_manager.run():
            yield

    return _lifespan
