"""MCPServer instance exposing the compliance knowledge base as an MCP tool.

This module is a pure protocol adapter: search_regulations() below calls the
existing AgentConversationService.ask() (the same application service backing
the /api/v1/agent/ask REST endpoint) and reshapes its result into a plain
dict. No new retrieval, ranking, or LLM orchestration logic lives here.
"""

from __future__ import annotations

from mcp.server import MCPServer
from starlette.responses import PlainTextResponse
from starlette.types import ASGIApp, Receive, Scope, Send

from app.config.settings import settings
from app.shared.bootstrap import get_agent_conversation_service, get_jwt_handler

# Single shared MCPServer instance — analogous to the single shared FastAPI
# `app` instance in app/api/main.py. Tools registered via @mcp.tool() below.
# Note: the installed mcp SDK (2.0.0) renamed the older "FastMCP" class to
# "MCPServer" (mcp.server.mcpserver.MCPServer); the .tool()/.streamable_http_app()
# API surface used here is unchanged across that rename.
mcp = MCPServer("ai-regulations")


@mcp.tool()
def search_regulations(query: str, top_k: int = 5) -> dict:
    """Search the compliance knowledge base and return a grounded answer.

    query: Natural-language search question, e.g. "国六排放标准最新要求".
    top_k: Maximum number of cited sources to return (default 5).
    """
    # No session_id is passed: this keeps each call stateless (no
    # ConversationStore reads/writes), matching "search" semantics rather
    # than multi-turn chat semantics.
    _, result = get_agent_conversation_service().ask(query=query, top_k=top_k)
    return {
        "answer": result.answer,
        "sources": [source.__dict__ for source in result.sources],
    }


class MCPAuthMiddleware:
    """Reject unauthenticated requests before they reach the MCP protocol handler.

    Mirrors the existing get_current_user dependency's behavior (auth.py) but
    implemented as raw ASGI middleware, since the mounted MCP app is a plain
    ASGI app, not a FastAPI/APIRouter instance that supports Depends().
    """

    def __init__(self, app: ASGIApp) -> None:
        """Store the wrapped ASGI app to delegate to once auth passes."""
        self.app = app

    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
        """Validate the bearer token for HTTP requests; pass non-HTTP scopes through."""
        # Only HTTP requests carry an Authorization header to check; lifespan
        # and other scope types must always pass through untouched.
        if scope["type"] != "http" or not settings.auth_enabled:
            await self.app(scope, receive, send)
            return

        headers = dict(scope["headers"])
        auth_header = headers.get(b"authorization", b"").decode()
        token = auth_header.removeprefix("Bearer ").strip()
        try:
            get_jwt_handler().decode_token(token)
        except ValueError as exc:
            # Reject before the MCP session/protocol layer ever sees the request.
            response = PlainTextResponse(str(exc), status_code=401)
            await response(scope, receive, send)
            return

        await self.app(scope, receive, send)


def build_mcp_asgi_app() -> ASGIApp:
    """Return the Streamable HTTP ASGI app for the MCP server, auth-guarded.

    streamable_http_path="/" is required here: MCPServer.streamable_http_app()
    registers its own internal route at "/mcp" by default, and this app is
    itself mounted at "/mcp" in api/main.py — without overriding the internal
    path to "/", the effective external path would be the confusing "/mcp/mcp"
    instead of "/mcp".
    """
    asgi_app = mcp.streamable_http_app(streamable_http_path="/")
    asgi_app.add_middleware(MCPAuthMiddleware)
    return asgi_app