"""mcp-gate: OAuth sign-in for a remote MCP server, with exe.dev as the login.

From https://exe.dev/deploy/mcp-server. claude.ai, ChatGPT and other chat apps
connect to remote MCP servers only with OAuth. mcp-gate is a small OAuth
authorization server plus a reverse proxy: it listens on port 8000, sends
people to exe.dev to sign in, asks them to approve the app, and passes
approved requests to the MCP server on UPSTREAM. It keeps no passwords.

It lets a request through when it has
  - an access token mcp-gate issued (Authorization: Bearer ...), or
  - an X-ExeDev-Email in ALLOWED_EMAILS, which exe.dev's proxy sets for a
    signed-in browser or a VM token (and strips from everything else), or
  - an X-Exedev-Source-Vm, which exe.dev sets on calls from your other VMs
    through a VM-to-VM integration.
The MCP server then sees X-ExeDev-Email (or X-Exedev-Source-Vm) as it would
behind exe.dev's proxy alone.

Settings come from the environment (see ~/mcp-gate.env):
  ALLOWED_EMAILS  exe.dev accounts that may sign in, comma-separated (required)
  UPSTREAM        the MCP server (default http://127.0.0.1:8001)
  PUBLIC_URL      this server's address (default https://<hostname>.exe.xyz)
  GATE_DB         SQLite file for clients and tokens (default ~/data/mcp-gate.db)
  PORT            port to listen on (default 8000)
"""

import html
import ipaddress
import json
import os
import secrets
import socket
import sqlite3
import time
from hashlib import sha256
from pathlib import Path
from urllib.parse import quote, urlsplit

import httpx2
import uvicorn
from mcp.server.auth.provider import (
    AccessToken,
    AuthorizationCode,
    AuthorizationParams,
    RefreshToken,
    construct_redirect_uri,
)
from mcp.server.auth.routes import create_auth_routes
from mcp.server.auth.settings import ClientRegistrationOptions, RevocationOptions
from mcp.shared.auth import OAuthClientInformationFull, OAuthToken
from pydantic import AnyHttpUrl
from starlette.applications import Starlette
from starlette.background import BackgroundTask
from starlette.requests import Request
from starlette.responses import HTMLResponse, JSONResponse, RedirectResponse, Response, StreamingResponse
from starlette.routing import Route

ALLOWED = {e.strip().lower() for e in os.environ.get("ALLOWED_EMAILS", "").split(",") if e.strip()}
UPSTREAM = os.environ.get("UPSTREAM", "http://127.0.0.1:8001").rstrip("/")
PUBLIC_URL = os.environ.get("PUBLIC_URL", f"https://{socket.gethostname()}.exe.xyz").rstrip("/")
DB = Path(os.environ.get("GATE_DB") or Path.home() / "data" / "mcp-gate.db")
PORT = int(os.environ.get("PORT") or 8000)

ACCESS_TTL = 3600  # an hour; clients refresh silently
REFRESH_TTL = 90 * 86400  # 90 days without use and you sign in again
HOP = {"connection", "keep-alive", "transfer-encoding", "te", "trailer", "upgrade",
       "proxy-authorization", "proxy-authenticate", "host", "content-length"}


def digest(token: str) -> str:
    return sha256(token.encode()).hexdigest()


class Provider:
    """OAuth for the MCP SDK's /authorize, /token and /register handlers.

    Clients and tokens live in SQLite (tokens only as hashes), so restarts
    don't sign anyone out. Pending sign-ins and codes live a few minutes in memory.
    """

    def __init__(self) -> None:
        DB.parent.mkdir(parents=True, exist_ok=True)
        with self.db() as con:
            con.executescript("""
                CREATE TABLE IF NOT EXISTS clients (client_id TEXT PRIMARY KEY, info TEXT NOT NULL);
                CREATE TABLE IF NOT EXISTS tokens (
                    hash TEXT PRIMARY KEY, kind TEXT NOT NULL, client_id TEXT NOT NULL,
                    scopes TEXT NOT NULL, subject TEXT NOT NULL, resource TEXT, expires_at INTEGER NOT NULL
                );
            """)
        self.pending: dict[str, tuple[float, OAuthClientInformationFull, AuthorizationParams]] = {}
        self.codes: dict[str, AuthorizationCode] = {}
        self.cimd: dict[str, tuple[float, OAuthClientInformationFull]] = {}

    def db(self) -> sqlite3.Connection:
        return sqlite3.connect(DB, timeout=10)

    # Clients: registered (RFC 7591), or identified by a metadata document URL (CIMD).
    async def get_client(self, client_id: str) -> OAuthClientInformationFull | None:
        if client_id.startswith("https://"):
            return await self.metadata_document(client_id)
        with self.db() as con:
            row = con.execute("SELECT info FROM clients WHERE client_id = ?", (client_id,)).fetchone()
        return OAuthClientInformationFull.model_validate_json(row[0]) if row else None

    async def register_client(self, client_info: OAuthClientInformationFull) -> None:
        with self.db() as con:
            con.execute("INSERT OR REPLACE INTO clients VALUES (?, ?)",
                        (client_info.client_id, client_info.model_dump_json()))

    async def metadata_document(self, url: str) -> OAuthClientInformationFull | None:
        cached = self.cimd.get(url)
        if cached and cached[0] > time.time():
            return cached[1]
        host = urlsplit(url).hostname or ""
        try:
            ipaddress.ip_address(host)
            return None  # no IP addresses
        except ValueError:
            pass
        if host == "localhost" or host.endswith(".int.exe.xyz") or "." not in host:
            return None  # nothing internal: integration hosts would add your credentials
        try:
            async with httpx2.AsyncClient(timeout=5, follow_redirects=False) as http:
                r = await http.get(url, headers={"Accept": "application/json"})
            doc = r.json() if r.status_code == 200 and len(r.content) < 65536 else {}
            if doc.get("client_id") != url:
                return None
            client = OAuthClientInformationFull(
                client_id=url,
                client_name=doc.get("client_name"),
                client_uri=doc.get("client_uri"),
                redirect_uris=doc["redirect_uris"],
                grant_types=[g for g in doc.get("grant_types", []) if g in ("authorization_code", "refresh_token")]
                or ["authorization_code", "refresh_token"],
                response_types=["code"],
                token_endpoint_auth_method="none",
                scope=doc.get("scope"),
            )
        except Exception:
            return None
        self.cimd[url] = (time.time() + 3600, client)
        return client

    # /authorize: remember the request, then send the browser to the consent page.
    async def authorize(self, client: OAuthClientInformationFull, params: AuthorizationParams) -> str:
        now = time.time()
        self.pending = {k: v for k, v in self.pending.items() if v[0] > now}
        txn = secrets.token_urlsafe(24)
        self.pending[txn] = (now + 600, client, params)
        return f"{PUBLIC_URL}/consent?txn={txn}"

    def approve(self, txn: str, email: str) -> str:
        _, client, params = self.pending.pop(txn)
        code = secrets.token_urlsafe(32)
        self.codes[code] = AuthorizationCode(
            code=code, scopes=params.scopes or [], expires_at=time.time() + 300,
            client_id=client.client_id, code_challenge=params.code_challenge,
            redirect_uri=params.redirect_uri,
            redirect_uri_provided_explicitly=params.redirect_uri_provided_explicitly,
            resource=params.resource, subject=email,
        )
        return construct_redirect_uri(str(params.redirect_uri), code=code, state=params.state)

    def deny(self, txn: str) -> str:
        _, _, params = self.pending.pop(txn)
        return construct_redirect_uri(str(params.redirect_uri), error="access_denied", state=params.state)

    # /token
    async def load_authorization_code(self, client: OAuthClientInformationFull, code: str) -> AuthorizationCode | None:
        c = self.codes.get(code)
        return c if c and c.client_id == client.client_id else None

    async def exchange_authorization_code(self, client: OAuthClientInformationFull, code: AuthorizationCode) -> OAuthToken:
        self.codes.pop(code.code, None)
        return self.issue(client.client_id, code.scopes, code.subject or "", code.resource)

    async def load_refresh_token(self, client: OAuthClientInformationFull, refresh_token: str) -> RefreshToken | None:
        t = self.lookup(refresh_token, "refresh")
        if not t or t["client_id"] != client.client_id:
            return None
        return RefreshToken(token=refresh_token, client_id=t["client_id"], scopes=t["scopes"],
                            expires_at=t["expires_at"], resource=t["resource"], subject=t["subject"])

    async def exchange_refresh_token(self, client: OAuthClientInformationFull, refresh_token: RefreshToken,
                                     scopes: list[str]) -> OAuthToken:
        self.forget(refresh_token.token)  # refresh tokens are single use
        return self.issue(client.client_id, scopes or refresh_token.scopes, refresh_token.subject or "",
                          refresh_token.resource)

    async def load_access_token(self, token: str) -> AccessToken | None:
        t = self.lookup(token, "access")
        if not t:
            return None
        return AccessToken(token=token, client_id=t["client_id"], scopes=t["scopes"],
                           expires_at=t["expires_at"], resource=t["resource"], subject=t["subject"])

    async def revoke_token(self, token: AccessToken | RefreshToken) -> None:
        self.forget(token.token)

    def issue(self, client_id: str, scopes: list[str], subject: str, resource: str | None) -> OAuthToken:
        access, refresh, now = secrets.token_urlsafe(32), secrets.token_urlsafe(32), int(time.time())
        with self.db() as con:
            con.execute("DELETE FROM tokens WHERE expires_at < ?", (now,))
            for token, kind, ttl in ((access, "access", ACCESS_TTL), (refresh, "refresh", REFRESH_TTL)):
                con.execute("INSERT INTO tokens VALUES (?, ?, ?, ?, ?, ?, ?)",
                            (digest(token), kind, client_id, " ".join(scopes), subject, resource, now + ttl))
        return OAuthToken(access_token=access, token_type="Bearer", expires_in=ACCESS_TTL,
                          refresh_token=refresh, scope=" ".join(scopes) or None)

    def lookup(self, token: str, kind: str) -> dict | None:
        with self.db() as con:
            row = con.execute("SELECT client_id, scopes, subject, resource, expires_at FROM tokens"
                              " WHERE hash = ? AND kind = ? AND expires_at > ?",
                              (digest(token), kind, int(time.time()))).fetchone()
        if not row:
            return None
        return {"client_id": row[0], "scopes": row[1].split(), "subject": row[2],
                "resource": row[3], "expires_at": row[4]}

    def forget(self, token: str) -> None:
        with self.db() as con:
            con.execute("DELETE FROM tokens WHERE hash = ?", (digest(token),))


provider = Provider()


def page(title: str, body: str, status: int = 200) -> HTMLResponse:
    return HTMLResponse(
        f"<!doctype html><meta name=viewport content='width=device-width'><title>{title}</title>"
        "<style>body{font:16px/1.5 system-ui,sans-serif;max-width:32rem;margin:4rem auto;padding:0 1rem}"
        "button{font:inherit;padding:.5rem 1.25rem;margin-right:.5rem}code{font-size:.9em}</style>"
        f"<h1>{title}</h1>{body}", status_code=status)


async def consent(request: Request) -> Response:
    """Sign in with exe.dev, then approve (or deny) the app that asked."""
    txn = request.query_params.get("txn") or (await request.form()).get("txn") or ""
    email = (request.headers.get("x-exedev-email") or "").lower()
    if not email:  # exe.dev signs the browser in and sends it back here
        return RedirectResponse(f"/__exe.dev/login?redirect={quote(f'/consent?txn={txn}', safe='')}", 303)
    if txn not in provider.pending:
        return page("Sign-in expired", "<p>Start again from the app you were connecting.</p>", 400)
    if email not in ALLOWED:
        return page("Not allowed", f"<p>{html.escape(email)} may not use this server. "
                    "Its owner can add you to <code>ALLOWED_EMAILS</code>.</p>"
                    "<form method=post action='/__exe.dev/logout'><button>Use another account</button></form>", 403)
    _, client, params = provider.pending[txn]
    if request.method == "POST":
        if request.headers.get("origin") != PUBLIC_URL:  # only this page may approve
            return page("Not allowed", "<p>That approval didn't come from this page.</p>", 403)
        form = await request.form()
        url = provider.approve(txn, email) if form.get("approve") else provider.deny(txn)
        return RedirectResponse(url, 302)
    name = html.escape(client.client_name or "An app")
    dest = html.escape(urlsplit(str(params.redirect_uri)).netloc or str(params.redirect_uri))
    return page(f"Connect {name}?", f"""
        <p><b>{name}</b> wants to use the MCP server at <code>{html.escape(PUBLIC_URL)}</code>
        as <b>{html.escape(email)}</b>. It can call every tool the server has.</p>
        <p>After you approve, you go back to <code>{dest}</code>.
        Approve only if you just started connecting from there.</p>
        <form method=post><input type=hidden name=txn value="{html.escape(txn)}">
        <button name=approve value=1>Approve</button><button name=deny value=1>Deny</button></form>""")


def authorization_server_metadata(request: Request) -> JSONResponse:
    return JSONResponse({
        "issuer": PUBLIC_URL,
        "authorization_endpoint": f"{PUBLIC_URL}/authorize",
        "token_endpoint": f"{PUBLIC_URL}/token",
        "registration_endpoint": f"{PUBLIC_URL}/register",
        "revocation_endpoint": f"{PUBLIC_URL}/revoke",
        "response_types_supported": ["code"],
        "grant_types_supported": ["authorization_code", "refresh_token"],
        "code_challenge_methods_supported": ["S256"],
        "token_endpoint_auth_methods_supported": ["none", "client_secret_post", "client_secret_basic"],
        "client_id_metadata_document_supported": True,
    })


def protected_resource_metadata(request: Request) -> JSONResponse:
    path = request.path_params.get("path", "")
    return JSONResponse({
        "resource": PUBLIC_URL + (f"/{path}" if path else ""),
        "authorization_servers": [PUBLIC_URL],
        "bearer_methods_supported": ["header"],
    })


async def caller(request: Request) -> dict[str, str] | None:
    """The identity headers to send upstream, or None if the request isn't allowed."""
    if vm := request.headers.get("x-exedev-source-vm"):
        return {"X-Exedev-Source-Vm": vm}
    if (email := (request.headers.get("x-exedev-email") or "").lower()) in ALLOWED:
        return {"X-ExeDev-Email": email}
    scheme, _, token = request.headers.get("authorization", "").partition(" ")
    if scheme.lower() == "bearer" and (t := await provider.load_access_token(token.strip())):
        if t.subject in ALLOWED:  # removing someone from ALLOWED_EMAILS cuts them off
            return {"X-ExeDev-Email": t.subject}
    return None


upstream = httpx2.AsyncClient(timeout=httpx2.Timeout(30, read=None))


async def proxy(request: Request) -> Response:
    who = await caller(request)
    if who is None:
        metadata = f"{PUBLIC_URL}/.well-known/oauth-protected-resource{request.url.path.rstrip('/')}"
        return JSONResponse(
            {"error": "invalid_token", "error_description": "sign in with OAuth"}, 401,
            headers={"WWW-Authenticate": f'Bearer error="invalid_token", resource_metadata="{metadata}"'})
    headers = {k: v for k, v in request.headers.items()
               if k not in HOP and k != "authorization" and not k.startswith("x-exedev-")}
    headers |= who | {"X-Forwarded-Host": request.headers.get("host", ""), "X-Forwarded-Proto": "https"}
    url = UPSTREAM + request.url.path + (f"?{request.url.query}" if request.url.query else "")
    req = upstream.build_request(request.method, url, headers=headers, content=request.stream())
    try:
        r = await upstream.send(req, stream=True)
    except httpx2.TransportError as e:
        return JSONResponse({"error": f"MCP server unreachable at {UPSTREAM}: {e}"}, 502)
    return StreamingResponse(r.aiter_raw(), r.status_code, background=BackgroundTask(r.aclose),
                             headers={k: v for k, v in r.headers.items() if k not in HOP})


methods = ["GET", "POST", "PUT", "PATCH", "DELETE", "OPTIONS", "HEAD"]
routes = [
    Route("/.well-known/oauth-authorization-server", authorization_server_metadata),
    Route("/.well-known/openid-configuration", authorization_server_metadata),
    Route("/.well-known/oauth-protected-resource", protected_resource_metadata),
    Route("/.well-known/oauth-protected-resource/{path:path}", protected_resource_metadata),
    Route("/consent", consent, methods=["GET", "POST"]),
    *create_auth_routes(provider, issuer_url=AnyHttpUrl(PUBLIC_URL),
                        client_registration_options=ClientRegistrationOptions(enabled=True),
                        revocation_options=RevocationOptions(enabled=True)),
    Route("/{path:path}", proxy, methods=methods),
]
app = Starlette(routes=routes)

if __name__ == "__main__":
    if not ALLOWED:
        raise SystemExit("mcp-gate: set ALLOWED_EMAILS to the exe.dev accounts that may sign in")
    uvicorn.run(app, host="0.0.0.0", port=PORT, proxy_headers=False)