#!/usr/bin/env python3
"""Verify the live OAuth+MCP flow without opening Claude or calling a language model.

Reads the owner code from local .env/environment, never prints credentials and
never follows the final redirect to the Claude account. Test grants are revoked.
"""

from __future__ import annotations

import asyncio
import argparse
import base64
import hashlib
import os
from pathlib import Path
import re
import secrets
import sys
from urllib.parse import parse_qs, urlsplit

import httpx
from dotenv import dotenv_values
from mcp import ClientSession
from mcp.client.streamable_http import streamable_http_client

ROOT = Path(__file__).resolve().parents[2]
CALLBACK = "https://claude.ai/api/mcp/auth_callback"


def require(condition, message):
    if not condition:
        raise RuntimeError(message)


async def browser_login(login_url, origin, password, engine):
    """Use a fresh browser profile and an actual form POST, with no forged Origin."""
    from playwright.async_api import async_playwright

    async with async_playwright() as playwright:
        # Only the project's server bypasses the deliberately unreachable proxy.
        # The browser can attempt Claude's redirect (testing CSP) but cannot send it.
        browser = await getattr(playwright, engine).launch(
            headless=True,
            proxy={"server": "http://127.0.0.1:9", "bypass": urlsplit(origin).hostname},
        )
        try:
            context = await browser.new_context()
            seen = {}
            callback_ready = asyncio.Event()

            def callback_request(request):
                if request.url.startswith(CALLBACK + "?"):
                    seen["url"] = request.url
                    seen["referrer"] = request.headers.get("referer")
                    callback_ready.set()

            page = await context.new_page()
            page.on("request", callback_request)
            await page.goto(login_url)
            await page.locator("input[name=password]").fill(password)
            is_form = (
                lambda request: request.method == "POST" and request.url == origin + "/oauth/login"
            )
            async with (
                page.expect_request(is_form) as request_info,
                page.expect_response(
                    lambda response: response.request.method == "POST"
                    and response.url == origin + "/oauth/login"
                ) as response_info,
            ):
                await page.locator("button[value=allow]").click(no_wait_after=True)
            request, response = await request_info.value, await response_info.value
            actual_origin = await request.header_value("origin")
            print(
                "Browser form: Origin=" + str(actual_origin) + ", HTTP " + str(response.status),
                flush=True,
            )
            require(actual_origin == origin, "Browser form sent an unexpected Origin")
            require(response.status == 303, "Browser login form failed")
            await asyncio.wait_for(callback_ready.wait(), timeout=10)
            require(
                seen.get("url") and seen.get("referrer") is None,
                "Callback or referrer protection failed",
            )
            return seen["url"]
        finally:
            await browser.close()


async def check(browser=None):
    values = dotenv_values(ROOT / ".env")
    resource = os.getenv("ZP_MCP_PUBLIC_URL") or values.get("ZP_MCP_PUBLIC_URL", "")
    password = os.getenv("ZP_MCP_LOGIN_PASSWORD") or values.get("ZP_MCP_LOGIN_PASSWORD", "")
    url = urlsplit(resource)
    require(
        url.scheme == "https"
        and url.hostname
        and url.path == "/mcp"
        and not url.query
        and not url.fragment
        and url.username is None
        and url.password is None,
        "Configure the HTTPS MCP URL",
    )
    require(password and len(password) >= 32, "Configure the local OAuth login code")
    origin = url.scheme + "://" + url.netloc
    async with httpx.AsyncClient(timeout=30, follow_redirects=False) as http:
        response = await http.post(resource, json={})
        require(
            response.status_code == 401
            and "resource_metadata=" in response.headers.get("www-authenticate", ""),
            "OAuth discovery challenge failed",
        )
        metadata_url = re.search(
            r'resource_metadata="([^"]+)"', response.headers["www-authenticate"]
        )[1]
        require(
            metadata_url == origin + "/.well-known/oauth-protected-resource/mcp",
            "Unexpected metadata destination",
        )
        response = await http.get(metadata_url)
        require(
            response.status_code == 200 and response.json().get("resource") == resource,
            "Resource metadata failed",
        )
        response = await http.get(origin + "/.well-known/oauth-authorization-server")
        require(response.status_code == 200, "Authorization metadata failed")
        metadata = response.json()
        require(
            all(
                metadata.get(key) == origin + path
                for key, path in {
                    "authorization_endpoint": "/authorize",
                    "registration_endpoint": "/register",
                    "token_endpoint": "/token",
                }.items()
            ),
            "Unexpected authorization endpoint destination",
        )
        require(
            "S256" in metadata.get("code_challenge_methods_supported", [])
            and "none" in metadata.get("token_endpoint_auth_methods_supported", []),
            "Missing public-client PKCE support",
        )
        print("OAuth discovery and S256 PKCE metadata: OK", flush=True)

        response = await http.post(
            origin + "/register",
            json={
                "client_name": "ZP Agent connection verification",
                "redirect_uris": [CALLBACK],
                "token_endpoint_auth_method": "none",
                "grant_types": ["authorization_code", "refresh_token"],
                "response_types": ["code"],
                "scope": "zpwiki:read offline_access",
            },
        )
        require(response.status_code == 201, "Dynamic client registration failed")
        client_id = response.json()["client_id"]
        verifier, state = secrets.token_urlsafe(48), secrets.token_urlsafe(24)
        challenge = (
            base64.urlsafe_b64encode(hashlib.sha256(verifier.encode()).digest())
            .decode()
            .rstrip("=")
        )
        response = await http.get(
            origin + "/authorize",
            params={
                "response_type": "code",
                "client_id": client_id,
                "redirect_uri": CALLBACK,
                "code_challenge": challenge,
                "code_challenge_method": "S256",
                "scope": "zpwiki:read offline_access",
                "state": state,
                "resource": resource,
            },
        )
        require(response.status_code == 302, "Authorization request failed")
        login = urlsplit(response.headers["location"])
        require(
            login.scheme + "://" + login.netloc == origin and login.path == "/oauth/login",
            "Unexpected login destination",
        )
        if browser:
            callback_url = await browser_login(
                response.headers["location"], origin, password, browser
            )
        else:
            response = await http.get(response.headers["location"])
            require(
                response.status_code == 200
                and "HttpOnly" in response.headers.get("set-cookie", "")
                and "Secure" in response.headers.get("set-cookie", ""),
                "Secure login page failed",
            )
            csrf = re.search(r'name="csrf" value="([^"]+)"', response.text)
            require(csrf is not None, "Missing login CSRF field")
            form = {
                "request": parse_qs(login.query)["request"][0],
                "csrf": csrf[1],
                "action": "allow",
            }
            rejected = await http.post(
                origin + "/oauth/login",
                data={**form, "password": "invalid-test-code"},
                headers={"Origin": origin},
            )
            require(rejected.status_code == 401, "Invalid owner code was not rejected")
            response = await http.post(
                origin + "/oauth/login",
                data={**form, "password": password},
                headers={"Origin": origin},
            )
            require(response.status_code == 303, "Owner login and consent failed")
            callback_url = response.headers["location"]
        callback = urlsplit(callback_url)
        require(
            callback.scheme + "://" + callback.netloc + callback.path == CALLBACK,
            "Unexpected callback destination",
        )
        query = parse_qs(callback.query)
        require(query.get("state") == [state] and query.get("code"), "OAuth state or code failed")
        body = {
            "grant_type": "authorization_code",
            "client_id": client_id,
            "code": query["code"][0],
            "code_verifier": verifier,
            "redirect_uri": CALLBACK,
            "resource": resource,
        }
        wrong = await http.post(
            origin + "/token", data={**body, "code_verifier": "incorrect-verifier"}
        )
        require(
            wrong.status_code == 400 and wrong.json().get("error") == "invalid_grant",
            "Invalid PKCE verifier accepted",
        )
        response = await http.post(origin + "/token", data=body)
        require(response.status_code == 200, "Authorization code exchange failed")
        token = response.json()
        refresh = token.get("refresh_token")
        require(
            refresh and 0 < token.get("expires_in", 0) <= 3600,
            "Missing limited token or refresh grant",
        )
        try:
            replay = await http.post(origin + "/token", data=body)
            require(replay.status_code == 400, "Authorization code replay accepted")
            print(
                "DCR, owner login, consent, state, code exchange and replay rejection: OK",
                flush=True,
            )
            async with httpx.AsyncClient(
                timeout=60, headers={"Authorization": "Bearer " + token["access_token"]}
            ) as authorized:
                async with streamable_http_client(resource, http_client=authorized) as (
                    read,
                    write,
                    _,
                ):
                    async with ClientSession(read, write) as session:
                        await session.initialize()
                        names = {item.name for item in (await session.list_tools()).tools}
                        require(
                            names == {"zpwiki_rag", "zpwiki_graphrag", "zpwiki_health"},
                            "MCP tool discovery failed",
                        )
                        result = await session.call_tool("zpwiki_health")
                        health = result.structuredContent or {}
                        require(
                            not result.isError and health.get("ready"), "Backend readiness failed"
                        )
                        versions = [
                            health["tools"][engine].get("index_version")
                            for engine in ["rag", "graphrag"]
                        ]
                        require(
                            versions[0] is not None
                            and versions[0] == versions[1] == health.get("index_version"),
                            "Different tool data versions",
                        )
                        for engine in ["rag", "graphrag"]:
                            invalid = await session.call_tool("zpwiki_" + engine, {"query": ""})
                            require(invalid.isError, "Invalid tool query accepted")
            print("OAuth-authenticated MCP, both tools and common index version: OK", flush=True)
            renewed = await http.post(
                origin + "/token",
                data={
                    "grant_type": "refresh_token",
                    "client_id": client_id,
                    "refresh_token": refresh,
                    "resource": resource,
                },
            )
            require(renewed.status_code == 200, "Refresh failed")
            new = renewed.json()
            require(new["refresh_token"] != refresh, "Refresh token did not rotate")
            retired = await http.get(
                origin + "/health", headers={"Authorization": "Bearer " + token["access_token"]}
            )
            require(retired.status_code == 401, "Old access token still active after refresh")
            refresh = new["refresh_token"]
            revoked = await http.post(
                origin + "/revoke", data={"client_id": client_id, "token": refresh}
            )
            require(revoked.status_code == 200, "Revocation failed")
            denied = await http.get(
                origin + "/health", headers={"Authorization": "Bearer " + new["access_token"]}
            )
            require(denied.status_code == 401, "Revoked grant still active")
            print("Refresh rotation and grant revocation: OK; external LLM calls: 0", flush=True)
        finally:
            await http.post(origin + "/revoke", data={"client_id": client_id, "token": refresh})


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--browser",
        choices=["firefox", "chromium"],
        help="Submit the login form in an isolated Playwright browser",
    )
    args = parser.parse_args()
    try:
        asyncio.run(check(args.browser))
    except Exception as error:
        # Never include HTTP bodies, credentials, token responses or callback URLs.
        print(
            "OAuth verification failed: "
            + (str(error) if type(error) is RuntimeError else type(error).__name__),
            file=sys.stderr,
        )
        raise SystemExit(1)