#!/usr/bin/env python3
"""Inventory recorded requests; opt in to a paced drift check; credentials require explicit source-specific opt-in."""

from __future__ import annotations

import argparse
import hashlib
import ipaddress
import json
import os
import re
import runpy
import sys
import time
import urllib.error
import urllib.request
from collections import Counter
from pathlib import Path
from urllib.parse import parse_qsl, urlsplit

ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT / "scripts"))
from replay_drift import (  # noqa: E402
    provider_policy_id,
    provider_policy_ids,
    product_headers,
    recording_provider,
    reply_qualification,
    reply_shape,
    shape_drift,
)

SENSITIVE = re.compile(r"api.?key|token|password|secret|authorization|credential", re.I)
MAX_REPLY_BYTES = 32 * 1024 * 1024
KEYED_S2_RECORDINGS = 5


def default_interval(keyed_s2: bool = False) -> float:
    # Derive the selected S2 schedule from the runtime provider policy.
    # Other inventory checks retain its conservative unkeyed interval.
    text = (ROOT / "src/sources/rate_limit.rs").read_text()
    pattern = (
        r"fn s2_min_interval\([^}]+?if has_s2_api_key\s*\{([^}]+)"
        if keyed_s2
        else r"fn s2_min_interval\([^}]+\}\s*else\s*\{([^}]+)"
    )
    function = re.search(pattern, text)
    seconds = re.search(
        r"Duration::from_secs\((\d+)\)", function[1] if function else ""
    )
    if not seconds:
        raise ValueError("cannot derive provider pacing")
    return float(seconds[1])


def disabled_ids(registry: list[dict], raw: str) -> set[str]:
    records = {record["id"]: record for record in registry}
    disabled = set()
    for token in filter(None, re.split(r"[,|]", raw.lower())):
        token = token.strip()
        if token in records:
            disabled.add(token)
        elif token in {"commercial-restricted", "license-required", "non-commercial"}:
            disabled.update(
                record["id"]
                for record in registry
                if record.get("commercial_terms")
                and (
                    token == "commercial-restricted"
                    or record["commercial_terms"] == token
                )
            )
        else:
            raise ValueError("unknown disabled source or group")
    if os.environ.get("BIOMCP_DISABLE_KEGG") == "1":
        disabled.add("kegg")
    return disabled


def credential_free(value: object) -> bool:
    if isinstance(value, dict):
        return all(
            not SENSITIVE.search(str(key)) and credential_free(child)
            for key, child in value.items()
        )
    if isinstance(value, list):
        return all(credential_free(child) for child in value)
    return True


def request_metadata(
    key: str, entry: dict, provider: str, routes: dict
) -> tuple[str, str, bytes | None, dict]:
    method, path = key.split(" ", 1)
    if method not in {"GET", "HEAD", "POST"}:
        raise ValueError("unsupported request method")
    url = entry["upstream"]
    parsed = urlsplit(url)
    if (
        parsed.username
        or parsed.password
        or parsed.fragment
        or any(SENSITIVE.search(name) for name, _ in parse_qsl(parsed.query))
    ):
        raise ValueError("credential-bearing URL is forbidden")
    try:
        loopback = ipaddress.ip_address(parsed.hostname or "").is_loopback
    except ValueError:
        loopback = parsed.hostname == "localhost"
    allowed = [base for name, base in routes.items() if name == provider]
    if not loopback and (
        parsed.scheme != "https"
        or not any(
            url.split("?", 1)[0] == base.split("?", 1)[0]
            or url.split("?", 1)[0].startswith(base.split("?", 1)[0].rstrip("/") + "/")
            for base in allowed
        )
    ):
        raise ValueError("upstream is outside the provider's declared endpoint")
    headers = entry.get("request_headers", {})
    if not isinstance(headers, dict) or any(
        name.lower() not in {"accept", "content-type"}
        or not isinstance(value, str)
        or "\n" in value
        or "\r" in value
        for name, value in headers.items()
    ):
        raise ValueError(
            "only nonsecret accept/content-type request headers are supported"
        )
    body = None
    if method == "POST":
        raw = entry.get("request_body")
        if not isinstance(raw, str):
            raise ValueError("recording has no reconstructable request body")
        try:
            safe = credential_free(json.loads(raw))
        except ValueError:
            safe = False
        if not safe:
            raise ValueError("request body is not credential-free JSON")
        body = raw.encode()
        if (
            "#" not in path
            or hashlib.sha256(body).hexdigest()[:16] != path.rsplit("#", 1)[1]
        ):
            raise ValueError("request body does not match recorded request hash")
    elif entry.get("request_body"):
        raise ValueError("unexpected request body")
    return (
        method,
        url,
        body,
        {
            **product_headers(ROOT),
            **{name.lower(): value for name, value in headers.items()},
        },
    )


class NoRedirect(urllib.request.HTTPRedirectHandler):
    def redirect_request(self, req, fp, code, msg, headers, newurl):
        return None


def keyed_shape_drift(expected: dict, actual: dict) -> list[str]:
    # Paths come only from the committed recording, never new live field names.
    # Type labels come from the parser's fixed vocabulary, never body values.
    kinds = {"null", "boolean", "object", "array", "number", "string", "element"}
    details = []
    for path, types in sorted(expected.items()):
        if path in actual and types == actual[path]:
            continue
        label = "[sensitive recorded path]" if SENSITIVE.search(path) else path
        wanted = ",".join(sorted(types & kinds)) or "signature"
        if path not in actual:
            details.append(f"missing field {label}: expected {wanted}")
        else:
            live = ",".join(sorted(actual[path] & kinds)) or "signature"
            details.append(f"changed type {label}: expected {wanted}; live {live}")
    return details


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--live",
        action="store_true",
        help="send eligible requests to their live providers",
    )
    parser.add_argument("--replay-dir", type=Path, default=ROOT / "testdata/replay")
    parser.add_argument("--disable-sources", default="")
    parser.add_argument(
        "--allow-restricted",
        default="",
        help="comma-separated sources whose access is explicitly authorized",
    )
    parser.add_argument(
        "--source",
        choices=["semantic-scholar"],
        help="select only Semantic Scholar recordings",
    )
    parser.add_argument(
        "--keyed-semantic-scholar",
        action="store_true",
        help="use inherited S2_API_KEY for the selected release recordings",
    )
    args = parser.parse_args(argv)
    try:
        registry = json.loads((ROOT / "docs/reference/sources.json").read_text())
        disabled = disabled_ids(
            registry,
            args.disable_sources + "," + os.environ.get("BIOMCP_DISABLE_SOURCES", ""),
        )
        records = {entry["id"]: entry for entry in registry}
        authorized = set(filter(None, args.allow_restricted.split(",")))
        if authorized - records.keys():
            raise ValueError("unknown restricted source authorization")
        helper = runpy.run_path(str(ROOT / "tools/replay-server"))
        overrides = {
            name: helper["prefix_for"](name)
            for name in helper["provider_routes"](ROOT / "src/sources")
        }
        identities = provider_policy_ids(ROOT / "src/sources", overrides, set(records))
        routes = {
            helper["prefix_for"](name): url
            for name, url in helper["provider_routes"](ROOT / "src/sources").items()
        }
        manifest = json.loads((args.replay_dir / "manifest.json").read_text())
        if manifest.get("schema_version") != 1:
            raise ValueError("unsupported manifest schema")
        entries = sorted(manifest["routes"].items())
        if args.source:
            entries = [
                (key, entry)
                for key, entry in entries
                if provider_policy_id(
                    recording_provider(key, entry, routes), identities
                )
                == args.source
            ]
        if args.keyed_semantic_scholar:
            if args.source != "semantic-scholar" or len(entries) != KEYED_S2_RECORDINGS:
                raise ValueError("keyed check requires the selected release recordings")
            for key, entry in entries:
                provider = recording_provider(key, entry, routes)
                method, url, _, _ = request_metadata(key, entry, provider, routes)
                parsed = urlsplit(url)
                declared = urlsplit(routes[provider])
                if parsed.scheme != "https" or (parsed.scheme, parsed.netloc) != (
                    declared.scheme,
                    declared.netloc,
                ):
                    raise ValueError("keyed check requires the declared HTTPS origin")
                if (method, parsed.path) not in {
                    ("GET", "/graph/v1/paper/search"),
                    ("POST", "/graph/v1/paper/batch"),
                }:
                    raise ValueError("keyed check requires read-only paper endpoints")
            key_available = bool(os.environ.get("S2_API_KEY", "").strip())
            if not key_available or any(
                not 32 <= ord(c) <= 126 for c in os.environ["S2_API_KEY"]
            ):
                raise ValueError("keyed check requires a valid inherited S2_API_KEY")
        interval = default_interval(args.keyed_semantic_scholar)
    except (ValueError, OSError, KeyError):
        print(json.dumps({"error": "invalid inventory or policy configuration"}))
        return 2
    counts: Counter = Counter(total=len(entries))
    stopped = {}
    last_start = None
    opener = urllib.request.build_opener(urllib.request.ProxyHandler({}), NoRedirect())
    for index, (key, entry) in enumerate(entries, 1):
        # Never print URLs, request bodies, response bodies or raw exceptions.
        provider = recording_provider(key, entry, routes)
        result = {"recording": index, "provider": provider, "file": entry.get("file")}
        state, details, added_paths = "blocked", [], []
        content_type = entry.get("headers", {}).get("content-type", "")
        qualification = reply_qualification(content_type)
        try:
            source_id = provider_policy_id(provider, identities)
            if source_id in disabled:
                state = "disabled"
            else:
                record = records.get(source_id)
                if not record:
                    raise ValueError("provider has no source policy")
                if record.get("commercial_terms") and source_id not in authorized:
                    raise ValueError(
                        "restricted provider needs explicit access authorization or disabling"
                    )
                if record.get("bioMcp_auth") == "required_env":
                    raise ValueError(
                        "provider requires credentials; this checker never sends credentials"
                    )
                method, url, body, headers = request_metadata(
                    key, entry, provider, routes
                )
                content_type = entry.get("headers", {}).get("content-type")
                if not content_type:
                    raise ValueError("recording has no historical content-type")
                path = (args.replay_dir / entry["file"]).resolve()
                if not path.is_relative_to(args.replay_dir.resolve()):
                    raise ValueError("response file is outside recording directory")
                expected = reply_shape(path.read_bytes(), content_type)
                if provider in stopped:
                    raise ValueError("provider stopped after rate limit or Retry-After")
                if not args.live:
                    state = "ready"
                else:
                    if last_start is not None:
                        time.sleep(max(0, interval - (time.monotonic() - last_start)))  # noqa: TID251 -- provider pacing
                    last_start = time.monotonic()
                    if args.keyed_semantic_scholar:
                        headers["x-api-key"] = os.environ["S2_API_KEY"]
                    request = urllib.request.Request(
                        url, data=body, headers=headers, method=method
                    )
                    try:
                        response = opener.open(request, timeout=90)
                    except urllib.error.HTTPError as error:
                        response = error
                    with response:
                        counts["checked"] += 1
                        if args.keyed_semantic_scholar:
                            result["http_status"] = response.status
                        if response.status == 429:
                            stopped[provider] = "HTTP 429"
                        elif response.headers.get("Retry-After"):
                            stopped[provider] = "Retry-After"
                        if response.status != entry["status"]:
                            details.append(
                                f"changed status: recorded {entry['status']}; live {response.status}"
                                if args.keyed_semantic_scholar
                                else "changed status"
                            )
                        if response.headers.get("Content-Type", "").split(";", 1)[
                            0
                        ].strip().lower() != content_type.split(";", 1)[
                            0
                        ].strip().lower() or (
                            qualification == "status/type/signature only"
                            and response.headers.get("Content-Type", "") != content_type
                        ):
                            details.append("changed content-type")
                        actual_body = response.read(MAX_REPLY_BYTES + 1)
                        if len(actual_body) > MAX_REPLY_BYTES:
                            raise ValueError("live response exceeds size limit")
                        try:
                            actual = reply_shape(
                                actual_body, response.headers.get("Content-Type", "")
                            )
                            if qualification == "structure":
                                details.extend(
                                    keyed_shape_drift(expected, actual)
                                    if args.keyed_semantic_scholar
                                    else shape_drift(expected, actual)
                                )
                                added_paths = sorted(actual.keys() - expected.keys())
                            elif expected != actual:
                                details.append("changed file signature")
                        except (ValueError, UnicodeError):
                            details.append(
                                "live response has unsupported or invalid shape"
                            )
                    state = "drift" if details else "passed"
        except (ValueError, KeyError, OSError) as error:
            # Only our fixed validation messages are safe to expose.
            details = (
                [str(error)]
                if type(error) is ValueError
                else ["request or recording unavailable"]
            )
        if args.keyed_semantic_scholar:
            if state == "blocked":
                details = [
                    f"provider stopped after {stopped[provider]}"
                    if provider in stopped
                    else "recording blocked"
                ]
            if provider in stopped:
                result["stop_reason"] = stopped[provider]
            # New live keys can echo credentials; only recorded paths are safe.
            added_paths = []
        counts[state] += 1
        print(
            json.dumps(
                {
                    **result,
                    "state": state,
                    "qualification": qualification,
                    "details": details,
                    "added_paths": added_paths,
                },
                sort_keys=True,
            )
        )
    print(
        json.dumps(
            {
                "summary": {
                    name: counts[name]
                    for name in (
                        "total",
                        "checked",
                        "passed",
                        "ready",
                        "disabled",
                        "blocked",
                        "drift",
                    )
                }
            },
            sort_keys=True,
        )
    )
    return 1 if counts["blocked"] or counts["drift"] else 0


if __name__ == "__main__":
    raise SystemExit(main())
