"""
EPANET MCP Server — Hugging Face Space wrapper
==============================================

Wraps the `epanet-mcp-server` (ePyT/EPANET) FastMCP server and exposes it over
BOTH modern MCP transports simultaneously, behind an optional API key:

    /mcp        -> Streamable HTTP   (HuggingChat, Codex, Perplexity, Gemini, ChatGPT dev-mode)
    /sse        -> HTTP + SSE        (legacy clients / Claude Desktop via mcp-remote)
    /messages/  -> SSE message sink
    /           -> human-readable landing / status page
    /health     -> unauthenticated health probe (returns "ok")

Design notes
------------
* A single Starlette parent app drives the Streamable-HTTP session manager
  lifespan (required) and routes each transport to its own sub-app via a pure
  ASGI dispatcher — no Starlette Mount prefix stripping, no double `/mcp/mcp`.
* Auth is PURE ASGI middleware (not BaseHTTPMiddleware) so it never buffers or
  breaks the streaming SSE / chunked responses.
* Stateful HTTP is kept ON: loaded networks live in an in-process registry and
  must persist across tool calls within one MCP session (single uvicorn worker).
"""

from __future__ import annotations

import contextlib
import hmac
import os
import textwrap

from starlette.applications import Starlette
from starlette.responses import HTMLResponse, JSONResponse, PlainTextResponse
from starlette.routing import Route

from mcp.server.transport_security import TransportSecuritySettings
from epanet_mcp.server import mcp

# ---------------------------------------------------------------------------
# Configuration
# ---------------------------------------------------------------------------
# CLIENT_API_KEY : if set, every MCP request must present it as a bearer token
#                  (Authorization: Bearer <key>) or X-API-Key header.
#                  If unset AND REQUIRE_AUTH is not "true", the server runs OPEN
#                  (handy for ChatGPT connectors, which are OAuth/none only).
API_KEY = os.environ.get("CLIENT_API_KEY", "").strip()
REQUIRE_AUTH = os.environ.get("REQUIRE_AUTH", "").strip().lower() in {"1", "true", "yes"}
AUTH_ENABLED = bool(API_KEY) or REQUIRE_AUTH

# Keep session state so load_network -> get_summary -> run_sim share one registry.
mcp.settings.stateless_http = False

# HF Spaces sit behind a proxy, so the inbound Host header is the public
# *.hf.space domain. FastMCP's DNS-rebinding protection allow-lists only
# localhost by default and would reject every real request with 421
# "Invalid Host header". Disable it (HF terminates TLS and controls Host).
mcp.settings.transport_security = TransportSecuritySettings(
    enable_dns_rebinding_protection=False,
)

streamable_app = mcp.streamable_http_app()  # serves /mcp  (+ owns session mgr lifespan)
sse_app = mcp.sse_app()                      # serves /sse + /messages/

PROTECTED_PREFIXES = ("/mcp", "/sse", "/messages")


# ---------------------------------------------------------------------------
# Pure-ASGI API key middleware
# ---------------------------------------------------------------------------
def _extract_key(headers: list[tuple[bytes, bytes]]) -> str | None:
    hdr = {k.lower(): v for k, v in headers}
    auth = hdr.get(b"authorization")
    if auth:
        text = auth.decode("latin-1").strip()
        if text.lower().startswith("bearer "):
            return text[7:].strip()
        return text
    xkey = hdr.get(b"x-api-key")
    if xkey:
        return xkey.decode("latin-1").strip()
    return None


class APIKeyMiddleware:
    def __init__(self, app):
        self.app = app

    async def __call__(self, scope, receive, send):
        if scope["type"] != "http" or not AUTH_ENABLED:
            await self.app(scope, receive, send)
            return

        path = scope.get("path", "")
        if not any(path == p or path.startswith(p + "/") or path.startswith(p)
                   for p in PROTECTED_PREFIXES):
            await self.app(scope, receive, send)
            return

        presented = _extract_key(scope.get("headers", []))
        ok = presented is not None and API_KEY != "" and hmac.compare_digest(presented, API_KEY)
        if not ok:
            resp = JSONResponse(
                {"error": "unauthorized",
                 "detail": "Provide the API key via 'Authorization: Bearer <key>' "
                           "or the 'X-API-Key' header."},
                status_code=401,
            )
            await resp(scope, receive, send)
            return

        await self.app(scope, receive, send)


# ---------------------------------------------------------------------------
# Transport dispatcher (pure ASGI)
# ---------------------------------------------------------------------------
class MCPDispatcher:
    """Route HTTP requests to the right transport sub-app by path prefix."""

    def __init__(self, streamable, sse, fallback):
        self.streamable = streamable
        self.sse = sse
        self.fallback = fallback

    async def __call__(self, scope, receive, send):
        if scope["type"] != "http":
            # lifespan is handled by the parent Starlette app's lifespan_context
            await self.fallback(scope, receive, send)
            return
        path = scope.get("path", "")
        # Normalise a trailing slash on the transport roots — Gemini / ChatGPT
        # sometimes append one (e.g. "/mcp/"), but the sub-app route is "/mcp".
        if path in ("/mcp/", "/sse/"):
            path = path.rstrip("/")
            scope = dict(scope)
            scope["path"] = path
        if path == "/mcp" or path.startswith("/mcp/"):
            await self.streamable(scope, receive, send)
        elif path == "/sse" or path.startswith("/messages"):
            await self.sse(scope, receive, send)
        else:
            await self.fallback(scope, receive, send)


# ---------------------------------------------------------------------------
# Landing / health routes (the fallback app)
# ---------------------------------------------------------------------------
def _base_url(request) -> str:
    # Honour HF's proxy headers so printed URLs use the public https host.
    proto = request.headers.get("x-forwarded-proto", request.url.scheme)
    host = request.headers.get("x-forwarded-host", request.headers.get("host", request.url.netloc))
    return f"{proto}://{host}"


async def health(request):
    return PlainTextResponse("ok")


async def landing(request):
    base = _base_url(request)
    auth_line = (
        "🔒 API key required — send <code>Authorization: Bearer &lt;key&gt;</code>"
        if AUTH_ENABLED else
        "🔓 Open access — no API key configured (set <code>CLIENT_API_KEY</code> to lock it down)"
    )
    tool_count = len(getattr(mcp._tool_manager, "_tools", {})) if hasattr(mcp, "_tool_manager") else "40+"
    html = textwrap.dedent(f"""\
    <!doctype html><html><head><meta charset="utf-8">
    <title>EPANET MCP Server</title>
    <meta name="viewport" content="width=device-width, initial-scale=1">
    <style>
      body{{font-family:system-ui,-apple-system,Segoe UI,Roboto,sans-serif;max-width:760px;
           margin:40px auto;padding:0 20px;line-height:1.55;color:#0f172a;background:#f8fafc}}
      code{{background:#e2e8f0;padding:1px 5px;border-radius:4px;font-size:.9em}}
      pre{{background:#0f172a;color:#e2e8f0;padding:14px 16px;border-radius:8px;overflow:auto;font-size:.85em}}
      h1{{margin-bottom:.2em}} .sub{{color:#475569;margin-top:0}}
      .grid{{display:grid;grid-template-columns:auto 1fr;gap:6px 16px;margin:12px 0}}
      .k{{color:#0369a1;font-weight:600}} a{{color:#0369a1}}
      .pill{{display:inline-block;background:#dbeafe;color:#1e40af;border-radius:999px;padding:2px 10px;font-size:.8em}}
    </style></head><body>
    <h1>💧 EPANET MCP Server</h1>
    <p class="sub">Water-distribution network modelling over MCP, powered by
       <a href="https://github.com/KIOS-Research/EPyT">ePyT</a> / EPANET.
       <span class="pill">{tool_count} tools</span></p>
    <p>{auth_line}</p>
    <h3>Endpoints</h3>
    <div class="grid">
      <span class="k">Streamable HTTP</span><span><code>{base}/mcp</code> — HuggingChat, Codex, Perplexity, Gemini, ChatGPT (dev mode)</span>
      <span class="k">SSE</span><span><code>{base}/sse</code> — legacy clients / Claude Desktop via <code>mcp-remote</code></span>
      <span class="k">Health</span><span><code>{base}/health</code></span>
    </div>
    <h3>Quick connect (Streamable HTTP)</h3>
    <pre>{{
  "mcpServers": {{
    "epanet": {{
      "url": "{base}/mcp",
      "headers": {{ "Authorization": "Bearer &lt;YOUR_KEY&gt;" }}
    }}
  }}
}}</pre>
    <p>Bundled ePyT benchmark networks + custom networks under
       <code>/home/user/app/networks/</code> (e.g. <code>net3.inp</code>, <code>ky4.inp</code>,
       <code>Net6.inp</code>, <code>BIWS.inp</code>). Try:
       <em>"List bundled networks, then load net3.inp and summarise it."</em></p>
    </body></html>""")
    return HTMLResponse(html)


fallback_app = Starlette(routes=[
    Route("/", landing),
    Route("/health", health),
])

dispatcher = MCPDispatcher(streamable_app, sse_app, fallback_app)
guarded = APIKeyMiddleware(dispatcher)


# ---------------------------------------------------------------------------
# Parent app: drives the Streamable-HTTP session-manager lifespan
# ---------------------------------------------------------------------------
@contextlib.asynccontextmanager
async def lifespan(app):
    async with streamable_app.router.lifespan_context(app):
        yield


app = Starlette(lifespan=lifespan)
app.mount("/", guarded)


if __name__ == "__main__":
    import uvicorn

    port = int(os.environ.get("PORT", "7860"))
    uvicorn.run(app, host="0.0.0.0", port=port, log_level="info")