"""MCP (Model Context Protocol) server for SERPent.

Every route registered on the FastAPI app is automatically re-exposed as an MCP
tool, so a single deployment serves both audiences:

    REST  ->  POST https://<host>/serp/search_arxiv
    MCP   ->  POST https://<host>/mcp        (streamable HTTP transport)

Nothing here duplicates business logic. FastMCP reads the app's OpenAPI schema
and, when a tool is called, dispatches the request back into the same FastAPI
app in-process (via an ASGI transport, no network round-trip), so tools always
stay in sync with the endpoints.

Usage (see the bottom of `app.py`):

    from mcp_server import mount_mcp_server
    mcp = mount_mcp_server(app)

Configuration (all optional, read from the environment):

    MCP_ENABLED        "0"/"false" to disable the MCP server entirely.
    MCP_PATH           Mount path. Default "/mcp".
    MCP_STATELESS      "0" to keep per-session state. Default stateless, which
                       is what you want behind the HF Spaces proxy.
    MCP_ALLOWED_HOSTS  Comma-separated Host allow-list. Default "*".
    MCP_ALLOWED_ORIGINS Comma-separated Origin allow-list. Default "*".
    MCP_EXPOSE_ALL     "1" to expose every endpoint as a tool instead of the
                       curated set (see EXCLUDED_ROUTES below).
"""

from __future__ import annotations

import logging
import os
from typing import Optional

from fastapi import FastAPI
from fastapi.routing import APIRoute
from fastapi.utils import generate_unique_id

from fastmcp import FastMCP
from fastmcp.server.providers.openapi import MCPType, RouteMap
from fastmcp.utilities.lifespan import combine_lifespans

logger = logging.getLogger(__name__)

MCP_SERVER_NAME = "SERPent"

# Endpoints kept out of the MCP tool surface. They stay fully available over
# REST — this only controls what an LLM sees.
#
# Fewer, well-differentiated tools measurably improve tool selection, and both
# groups below are redundant *for an agent*:
#
#   /ops/*        The regular `scrap_patent` / `search_patents` tools already
#                 fall back to EPO OPS on their own. Calling the OPS tools
#                 directly skips that fallback, so a patent missing from OPS
#                 becomes a hard error instead of a Google Patents hit.
#   search_brave  These are `search` with the fallback chain removed. `search`
#   search_bing   already tries DuckDuckGo, then Brave, then Bing, so exposing
#   search_duck   them separately only invites the model to pick a worse path.
#
# Set MCP_EXPOSE_ALL=1 to expose everything again, or edit this list.
EXCLUDED_ROUTES = [
    r"^/ops/.*",
    r"^/serp/search_(brave|bing|duck)$",
]

MCP_INSTRUCTIONS = """\
SERPent gives you web, academic and patent search plus full-text patent retrieval.

Typical workflow:
  1. Search with the backend that matches the question:
     - `search` for general web results (tries DuckDuckGo, then Brave, then Bing)
     - `search_arxiv` for preprints, `search_google_scholar` for academic papers
     - `search_patents` for prior art (falls back to the EPO OPS API)
  2. Feed the patent ids you found to `scrap_patent` (one) or `scrap_patents`
     (many, and much faster than looping) to get title, abstract, claims and
     description.

Every search tool accepts a LIST of queries and runs them concurrently, so send
all your query variations in one call rather than making several calls. Set
`n_results` to control how many hits come back per query.

`search` and `search_patents` already fall back across several backends
internally, including the official EPO OPS patent API, so a single call is the
most thorough option available — there is no lower-level tool to reach for when
one comes back empty. A patent that returns "not found" is genuinely absent
from every backend; move on to the next id rather than retrying.
"""


def _env_flag(name: str, default: bool) -> bool:
    raw = os.getenv(name)
    if raw is None:
        return default
    return raw.strip().lower() not in {"0", "false", "no", "off", ""}


def _env_list(name: str, default: list[str]) -> list[str]:
    raw = os.getenv(name)
    if not raw:
        return default
    return [item.strip() for item in raw.split(",") if item.strip()]


def _tool_names(app: FastAPI) -> dict[str, str]:
    """Map each route's OpenAPI operationId to a clean MCP tool name.

    FastAPI's generated operationIds look like
    `search_arxiv_serp_search_arxiv_post`, which is what an LLM would otherwise
    see. The Python handler name (`search_arxiv`) is a far better tool name, so
    use that instead. Adding an explicit `operation_id=` to a route decorator
    overrides this.
    """
    names: dict[str, str] = {}

    def walk(routes) -> None:
        for route in routes:
            # Depending on the FastAPI version, included routers either flatten
            # into the parent or stay nested, so recurse.
            if isinstance(route, APIRoute):
                generate = getattr(route, "generate_unique_id_function", None)
                # FastAPI stores an unset callable as a DefaultPlaceholder.
                generate = getattr(generate, "value", generate)
                if not callable(generate):
                    generate = generate_unique_id
                operation_id = route.operation_id or generate(route)
                names[operation_id] = route.name
            # Starlette Mounts nest under `.routes`; FastAPI >= 0.141 wraps
            # included routers in `_IncludedRouter`, which nests under
            # `.original_router`.
            sub = getattr(route, "routes", None)
            if sub:
                walk(sub)
            included = getattr(route, "original_router", None)
            if included is not None and getattr(included, "routes", None):
                walk(included.routes)

    walk(app.routes)
    return names


def build_mcp_server(app: FastAPI) -> FastMCP:
    """Build an MCP server exposing the app's routes as tools."""
    if _env_flag("MCP_EXPOSE_ALL", False):
        route_maps = None
        logger.info("MCP exposing all endpoints (MCP_EXPOSE_ALL=1).")
    else:
        route_maps = [
            RouteMap(pattern=pattern, mcp_type=MCPType.EXCLUDE)
            for pattern in EXCLUDED_ROUTES
        ]

    mcp = FastMCP.from_fastapi(
        app=app,
        name=MCP_SERVER_NAME,
        mcp_names=_tool_names(app),
        route_maps=route_maps,
        # Requests are dispatched straight back into this same app in-process.
        httpx_client_kwargs={"timeout": 300.0},
    )
    mcp.instructions = MCP_INSTRUCTIONS
    return mcp


def mount_mcp_server(app: FastAPI, path: Optional[str] = None) -> Optional[FastMCP]:
    """Mount the MCP server onto `app` (default: `/mcp`).

    Must be called after every router has been included, otherwise the missing
    routes will not show up as tools. Returns None when MCP is disabled.
    """
    if not _env_flag("MCP_ENABLED", True):
        logger.info("MCP server disabled (MCP_ENABLED=0).")
        return None

    path = path or os.getenv("MCP_PATH", "/mcp")
    path = "/" + path.strip("/")

    mcp = build_mcp_server(app)

    mcp_app = mcp.http_app(
        path=path,
        stateless_http=_env_flag("MCP_STATELESS", True),
        allowed_hosts=_env_list("MCP_ALLOWED_HOSTS", ["*"]),
        allowed_origins=_env_list("MCP_ALLOWED_ORIGINS", ["*"]),
    )

    # The MCP session manager needs its own lifespan to run. The app already has
    # one (Playwright), so run both.
    app.router.lifespan_context = combine_lifespans(
        app.router.lifespan_context,
        mcp_app.lifespan,
    )

    # The MCP endpoint is `path` *inside* mcp_app, and mcp_app is mounted at the
    # root. Mounting it at `path` instead would make Starlette answer a bare
    # `POST /mcp` with a 307 to `/mcp/`, which several MCP clients choke on.
    # This mount is added last, so all the API routes above still match first.
    app.mount("/", mcp_app)

    logger.info("MCP server mounted at %s (streamable HTTP).", path)
    return mcp