"""Project-local OAuth provider for the MCP SDK: consent, PKCE grants and rotating tokens.

The SDK validates protocol requests, redirects, client credentials and S256 PKCE.
This module stores grants and authenticates the project owner on a separate consent page.
It implements a shared project access code, not a multi-user identity service.
"""

from __future__ import annotations

from contextlib import contextmanager
import hashlib
import hmac
from html import escape
import json
import os
from pathlib import Path
import re
import secrets
import sqlite3
import time
from urllib.parse import parse_qs, urlsplit

from mcp.server.auth.provider import (
    AccessToken,
    AuthorizationCode,
    AuthorizationParams,
    AuthorizeError,
    RefreshToken,
    RegistrationError,
    TokenError,
    construct_redirect_uri,
)
from mcp.shared.auth import OAuthClientInformationFull, OAuthToken
from starlette.responses import HTMLResponse, JSONResponse, RedirectResponse

SCOPE = "zpwiki:read"
SCOPES = [SCOPE, "offline_access"]
CALLBACKS = {"https://claude.ai/api/mcp/auth_callback"}
COOKIE = "__Host-zp-oauth"
PENDING_TTL, CODE_TTL, ACCESS_TTL, REFRESH_TTL = 600, 120, 3600, 30 * 86400


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


class Store:
    def __init__(self, path: Path, namespace: str):
        self.path, self.namespace = path, namespace
        path.parent.mkdir(parents=True, exist_ok=True, mode=0o700)
        if not path.exists():
            fd = os.open(path, os.O_CREAT | os.O_EXCL | os.O_WRONLY, 0o600)
            os.close(fd)
        path.chmod(0o600)
        with self.transaction() as conn:
            conn.execute("""CREATE TABLE IF NOT EXISTS records (
                namespace TEXT NOT NULL, kind TEXT NOT NULL, key TEXT NOT NULL,
                body TEXT NOT NULL, expires REAL NOT NULL, family TEXT NOT NULL DEFAULT '',
                PRIMARY KEY(namespace,kind,key))""")
            conn.execute("CREATE INDEX IF NOT EXISTS oauth_family ON records(namespace,family)")
            conn.execute("CREATE INDEX IF NOT EXISTS oauth_expiry ON records(expires)")

    @contextmanager
    def transaction(self):
        conn = sqlite3.connect(self.path, timeout=5)
        try:
            conn.execute("BEGIN IMMEDIATE")
            yield conn
            conn.commit()
        except BaseException:
            conn.rollback()
            raise
        finally:
            conn.close()

    def get(self, conn, kind, key):
        row = conn.execute(
            "SELECT body FROM records WHERE namespace=? AND kind=? AND key=? AND expires>?",
            (self.namespace, kind, key, time.time()),
        ).fetchone()
        return json.loads(row[0]) if row else None

    def put(self, conn, kind, key, body, expires, family=""):
        conn.execute(
            "INSERT OR REPLACE INTO records VALUES (?,?,?,?,?,?)",
            (self.namespace, kind, key, json.dumps(body), expires, family),
        )

    def delete(self, conn, kind, key):
        conn.execute(
            "DELETE FROM records WHERE namespace=? AND kind=? AND key=?",
            (self.namespace, kind, key),
        )

    def capacity(self, conn, kind, maximum):
        conn.execute("DELETE FROM records WHERE expires<=?", (time.time(),))
        return (
            conn.execute(
                "SELECT count(*) FROM records WHERE namespace=? AND kind=?", (self.namespace, kind)
            ).fetchone()[0]
            < maximum
        )

    def revoke_family(self, conn, family):
        conn.execute("DELETE FROM records WHERE namespace=? AND family=?", (self.namespace, family))


class OAuthProvider:
    def __init__(self, resource: str, static_token: str, password: str, database: Path):
        if (
            len(password) < 32
            or not password.isascii()
            or any(c.isspace() for c in password)
            or hmac.compare_digest(password, static_token)
        ):
            raise ValueError(
                "Configure a separate ZP_MCP_LOGIN_PASSWORD with at least 32 characters"
            )
        self.resource = resource
        url = urlsplit(resource)
        self.origin = url.scheme + "://" + url.netloc
        self.static_token, self.password_hash = static_token, digest(password)
        # A changed public URL or rotated owner credential invalidates previous sessions.
        self.store = Store(database, digest(resource + "\0" + static_token + "\0" + password))

    async def get_client(self, client_id):
        with self.store.transaction() as conn:
            value = self.store.get(conn, "client", client_id)
        return OAuthClientInformationFull.model_validate(value) if value else None

    async def register_client(self, client_info):
        redirects = [str(uri) for uri in client_info.redirect_uris or []]
        if (
            not redirects
            or len(redirects) > 4
            or any(uri not in CALLBACKS for uri in redirects)
            or client_info.token_endpoint_auth_method
            not in {"none", "client_secret_post", "client_secret_basic"}
            or set(client_info.grant_types) != {"authorization_code", "refresh_token"}
            or client_info.response_types != ["code"]
            or not set((client_info.scope or "").split()) <= set(SCOPES)
        ):
            raise RegistrationError(
                "invalid_client_metadata",
                "Only the documented Claude web callback and code flow are supported",
            )
        with self.store.transaction() as conn:
            if not self.store.capacity(conn, "client", 1024):
                raise RegistrationError(
                    "invalid_client_metadata", "Client registration capacity reached"
                )
            self.store.put(
                conn,
                "client",
                client_info.client_id,
                client_info.model_dump(mode="json"),
                time.time() + 90 * 86400,
            )

    async def authorize(self, client, params: AuthorizationParams):
        if params.resource not in {None, self.resource}:
            raise AuthorizeError(
                "invalid_request", "The requested resource does not match this MCP server"
            )
        scopes = params.scopes or (client.scope or SCOPE).split()
        if SCOPE not in scopes or not set(scopes) <= set(SCOPES):
            raise AuthorizeError("invalid_scope", "The read scope is required")
        if not re.fullmatch(r"[A-Za-z0-9_-]{43}", params.code_challenge):
            raise AuthorizeError("invalid_request", "S256 PKCE is required")
        request_id = secrets.token_urlsafe(32)
        value = {
            "client_id": client.client_id,
            "client_name": (client.client_name or "Claude")[:100],
            "params": params.model_copy(
                update={"scopes": scopes, "resource": self.resource}
            ).model_dump(mode="json"),
            "attempts": 0,
            "expires": time.time() + PENDING_TTL,
        }
        with self.store.transaction() as conn:
            if not self.store.capacity(conn, "pending", 1024):
                raise AuthorizeError("temporarily_unavailable", "Please retry later")
            self.store.put(conn, "pending", digest(request_id), value, value["expires"])
        return self.origin + "/oauth/login?request=" + request_id

    async def load_authorization_code(self, client, authorization_code):
        with self.store.transaction() as conn:
            value = self.store.get(conn, "code", digest(authorization_code))
        if not value or value["client_id"] != client.client_id:
            return None
        return AuthorizationCode(code=authorization_code, **value)

    def _issue(self, conn, client_id, scopes, family, refresh_expiry=None):
        now = int(time.time())
        access = secrets.token_urlsafe(48)
        refresh_expiry = refresh_expiry or now + REFRESH_TTL
        expiry = min(now + ACCESS_TTL, refresh_expiry)
        common = {
            "client_id": client_id,
            "scopes": scopes,
            "resource": self.resource,
            "subject": "project-user",
        }
        self.store.put(
            conn, "access", digest(access), {**common, "expires_at": expiry}, expiry, family
        )
        refresh = None
        if "offline_access" in scopes:
            refresh = secrets.token_urlsafe(48)
            self.store.put(
                conn,
                "refresh",
                digest(refresh),
                {**common, "expires_at": refresh_expiry, "family": family, "used": False},
                refresh_expiry,
                family,
            )
        return OAuthToken(
            access_token=access,
            expires_in=expiry - now,
            refresh_token=refresh,
            scope=" ".join(scopes),
        )

    async def exchange_authorization_code(self, client, authorization_code):
        with self.store.transaction() as conn:
            value = self.store.get(conn, "code", digest(authorization_code.code))
            if not value or value["client_id"] != client.client_id:
                raise TokenError("invalid_grant", "The code has expired or was already used")
            self.store.delete(conn, "code", digest(authorization_code.code))
            return self._issue(conn, client.client_id, value["scopes"], secrets.token_hex(24))

    async def load_refresh_token(self, client, refresh_token):
        with self.store.transaction() as conn:
            value = self.store.get(conn, "refresh", digest(refresh_token))
        if not value or value["client_id"] != client.client_id:
            return None
        return RefreshToken(
            token=refresh_token, **{k: v for k, v in value.items() if k not in {"family", "used"}}
        )

    async def exchange_refresh_token(self, client, refresh_token, scopes):
        replay = False
        with self.store.transaction() as conn:
            value = self.store.get(conn, "refresh", digest(refresh_token.token))
            if not value or value["client_id"] != client.client_id:
                raise TokenError("invalid_grant", "Invalid refresh token")
            if value["used"]:
                self.store.revoke_family(conn, value["family"])
                replay = True
            else:
                if SCOPE not in scopes or not set(scopes) <= set(value["scopes"]):
                    raise TokenError("invalid_scope", "Requested scopes exceed the grant")
                value["used"] = True
                self.store.put(
                    conn,
                    "refresh",
                    digest(refresh_token.token),
                    value,
                    value["expires_at"],
                    value["family"],
                )
                # Rotate both credentials; retire old access tokens in the same family.
                conn.execute(
                    "DELETE FROM records WHERE namespace=? AND kind=? AND family=?",
                    (self.store.namespace, "access", value["family"]),
                )
                issued = self._issue(
                    conn, client.client_id, scopes, value["family"], value["expires_at"]
                )
        if replay:
            raise TokenError("invalid_grant", "Refresh token reuse revoked this grant")
        return issued

    async def load_access_token(self, token):
        if hmac.compare_digest(token.encode(), self.static_token.encode()):
            return AccessToken(
                token=token, client_id="local-static", scopes=[SCOPE], resource=self.resource
            )
        with self.store.transaction() as conn:
            value = self.store.get(conn, "access", digest(token))
        return AccessToken(token=token, **value) if value else None

    async def revoke_token(self, token):
        with self.store.transaction() as conn:
            value = self.store.get(conn, "access", digest(token.token)) or self.store.get(
                conn, "refresh", digest(token.token)
            )
            row = conn.execute(
                "SELECT family FROM records WHERE namespace=? AND key=? AND kind IN (?,?)",
                (self.store.namespace, digest(token.token), "access", "refresh"),
            ).fetchone()
            if value and row:
                self.store.revoke_family(conn, row[0])

    def _form(self, request_id, pending, csrf, message="", status=200):
        target = escape(urlsplit(pending["params"]["redirect_uri"]).hostname or "")
        name = escape(pending["client_name"])
        document = f"""<!doctype html><html lang="sk"><meta charset="utf-8">
<meta name="viewport" content="width=device-width,initial-scale=1"><title>Prihlásenie — ZP Agent</title>
<style>body{{font:17px system-ui;max-width:520px;margin:8vh auto;padding:24px;line-height:1.5;background:#f5f5f5;color:#222}}
form{{background:white;padding:24px;border-radius:12px}}input,button{{box-sizing:border-box;width:100%;padding:12px;margin:8px 0;font:inherit}}
.error{{color:#a00}}small{{display:block}}</style>
<h1>Pripojiť ZP Agenta</h1><p>Aplikácia <strong>{name}</strong> žiada prístup k vyhľadávaniu RAG a GraphRAG.</p>
<p>Po potvrdení sa vrátiš na <strong>{target}</strong>.</p><form method="post" action="/oauth/login">
<input type="hidden" name="request" value="{escape(request_id)}"><input type="hidden" name="csrf" value="{escape(csrf)}">
<label for="password">Prístupový kód ZP Agenta</label><input id="password" name="password" type="password" autocomplete="current-password" required autofocus>
<small>Kód ti poskytne správca ZP Agenta. Nezadávaj sem heslo svojho účtu Claude.</small>
<p class="error">{escape(message)}</p><button name="action" value="allow">Prihlásiť a povoliť prístup</button>
<button name="action" value="deny" formnovalidate>Zrušiť</button></form></html>"""
        return HTMLResponse(document, status_code=status)

    async def login(self, request):
        if request.method == "GET":
            request_id = request.query_params.get("request", "")
            with self.store.transaction() as conn:
                pending = self.store.get(conn, "pending", digest(request_id))
                if not pending:
                    return HTMLResponse(
                        "Prihlásenie vypršalo. Začni pripojenie znova v Claude.", status_code=400
                    )
                csrf = secrets.token_urlsafe(32)
                pending["csrf"] = digest(csrf)
                self.store.put(conn, "pending", digest(request_id), pending, pending["expires"])
            response = self._form(request_id, pending, csrf)
            response.set_cookie(
                COOKIE,
                csrf,
                max_age=PENDING_TTL,
                secure=True,
                httponly=True,
                samesite="lax",
                path="/",
            )
            return response
        if request.headers.get("origin") not in {None, self.origin}:
            return JSONResponse({"error": "invalid_origin"}, status_code=403)
        if (
            request.headers.get("content-type", "").split(";")[0]
            != "application/x-www-form-urlencoded"
        ):
            return JSONResponse({"error": "invalid_request"}, status_code=415)
        form = parse_qs(
            (await request.body()).decode("utf-8", errors="replace"), keep_blank_values=True
        )
        if any(len(values) != 1 for values in form.values()):
            return JSONResponse({"error": "invalid_request"}, status_code=400)
        get = lambda name: form.get(name, [""])[0]
        request_id, csrf = get("request"), get("csrf")
        with self.store.transaction() as conn:
            pending = self.store.get(conn, "pending", digest(request_id))
            if (
                not pending
                or not csrf
                or not request.cookies.get(COOKIE)
                or not hmac.compare_digest(csrf.encode(), request.cookies[COOKIE].encode())
                or not hmac.compare_digest(digest(csrf), pending.get("csrf", ""))
            ):
                return JSONResponse({"error": "invalid_csrf_or_expired_login"}, status_code=403)
            if get("action") == "deny":
                self.store.delete(conn, "pending", digest(request_id))
                response = RedirectResponse(
                    construct_redirect_uri(
                        pending["params"]["redirect_uri"],
                        error="access_denied",
                        state=pending["params"].get("state"),
                    ),
                    status_code=303,
                )
                response.delete_cookie(COOKIE, secure=True, httponly=True, samesite="lax")
                return response
            bucket = str(int(time.time()) // 60)
            rate = self.store.get(conn, "login_rate", bucket) or {"count": 0}
            rate["count"] += 1
            self.store.put(conn, "login_rate", bucket, rate, time.time() + 120)
            if rate["count"] > 30 or pending["attempts"] >= 5:
                return self._form(
                    request_id, pending, csrf, "Príliš veľa pokusov. Skús to neskôr.", 429
                )
            if get("action") != "allow" or not hmac.compare_digest(
                digest(get("password")), self.password_hash
            ):
                pending["attempts"] += 1
                self.store.put(conn, "pending", digest(request_id), pending, pending["expires"])
                return self._form(request_id, pending, csrf, "Nesprávny prístupový kód.", 401)
            code = secrets.token_urlsafe(32)
            params = pending["params"]
            value = {
                k: params[k]
                for k in [
                    "scopes",
                    "code_challenge",
                    "redirect_uri",
                    "redirect_uri_provided_explicitly",
                    "resource",
                ]
            }
            value.update(
                client_id=pending["client_id"],
                expires_at=time.time() + CODE_TTL,
                subject="project-user",
            )
            self.store.put(conn, "code", digest(code), value, value["expires_at"])
            self.store.delete(conn, "pending", digest(request_id))
        response = RedirectResponse(
            construct_redirect_uri(params["redirect_uri"], code=code, state=params.get("state")),
            status_code=303,
        )
        response.delete_cookie(COOKIE, secure=True, httponly=True, samesite="lax")
        return response


class OAuthBoundary:
    """Bound auth request bodies, enforce resource binding and harden the consent page."""

    def __init__(self, app, provider):
        self.app, self.provider = app, provider

    async def __call__(self, scope, receive, send):
        if scope["type"] != "http":
            return await self.app(scope, receive, send)
        headers = dict(scope["headers"])
        public_host = urlsplit(self.provider.origin).netloc
        host = headers.get(b"host", b"").decode("latin1")
        if host != public_host and not re.fullmatch(r"(?:127\.0\.0\.1|localhost)(?::\d+)?", host):
            return await JSONResponse({"error": "invalid_host"}, status_code=421)(
                scope, receive, send
            )
        path = scope["path"]
        if scope["method"] == "POST" and path in {
            "/token",
            "/register",
            "/authorize",
            "/revoke",
            "/oauth/login",
        }:
            body = bytearray()
            while True:
                message = await receive()
                if message["type"] == "http.disconnect":
                    return
                body.extend(message.get("body", b""))
                if len(body) > 16384:
                    return await JSONResponse({"error": "request_too_large"}, status_code=413)(
                        scope, receive, send
                    )
                if not message.get("more_body"):
                    break
            if path in {"/token", "/revoke"}:
                if (
                    headers.get(b"content-type", b"").split(b";")[0]
                    != b"application/x-www-form-urlencoded"
                ):
                    return await JSONResponse({"error": "invalid_request"}, status_code=415)(
                        scope, receive, send
                    )
                fields = parse_qs(body.decode("utf-8", errors="replace"), keep_blank_values=True)
                if any(len(v) != 1 for v in fields.values()) or fields.get(
                    "resource", [self.provider.resource]
                ) != [self.provider.resource]:
                    return await JSONResponse({"error": "invalid_request"}, status_code=400)(
                        scope, receive, send
                    )
                # SDK 1.x models an optional revocation secret as a required nullable field.
                # Public clients have no secret; normalize that omission before its parser.
                if path == "/revoke" and "client_secret" not in fields:
                    body.extend(b"&client_secret=")
            if path == "/register":
                try:
                    value = json.loads(body)
                    if not isinstance(value, dict):
                        raise ValueError()
                except (ValueError, UnicodeError):
                    return await JSONResponse(
                        {"error": "invalid_client_metadata"}, status_code=400
                    )(scope, receive, send)
            original_receive = receive
            sent = False

            async def replay():
                nonlocal sent
                if not sent:
                    sent = True
                    return {"type": "http.request", "body": bytes(body), "more_body": False}
                return await original_receive()

            receive = replay

        async def secured(message):
            if message["type"] == "http.response.start":
                existing = list(message.get("headers", []))
                names = {k.lower() for k, v in existing}
                defaults = {
                    b"cache-control": b"no-store",
                    b"pragma": b"no-cache",
                    b"referrer-policy": b"no-referrer",
                    b"x-content-type-options": b"nosniff",
                    b"x-frame-options": b"DENY",
                }
                if path == "/oauth/login":
                    # no-referrer makes native form POSTs send Origin: null (Fetch).
                    # Keep the real origin for this same-origin form while suppressing
                    # the referrer on the redirect to Claude. CSRF checks stay strict.
                    defaults[b"referrer-policy"] = b"same-origin"
                    # Chromium applies form-action to POST redirects too. Permit only
                    # the same fixed callbacks accepted during OAuth registration.
                    callbacks = " ".join(sorted(CALLBACKS))
                    defaults[b"content-security-policy"] = (
                        "default-src 'none'; style-src 'unsafe-inline'; form-action 'self' "
                        + callbacks
                        + "; frame-ancestors 'none'; base-uri 'none'"
                    ).encode()
                message = {
                    **message,
                    "headers": existing + [(k, v) for k, v in defaults.items() if k not in names],
                }
            await send(message)

        return await self.app(scope, receive, secured)