#!/usr/bin/env python3
"""Check a deployed knowledge Space: the gate fails closed, the surface matches the descriptor.

C9 of the Phase 1 plan, run against the live Space rather than the code. ``stage: RUNNING`` only
says the container came up; this signs in the way the orchestrator does and asks real questions.

Two tokens, both the caller's HF token, as the orchestrator sends them:

* ``token=`` — Space *access*. It gets a caller past HF's own wall on a private Space; a public
  Space ignores it. The app never reads it.
* ``headers["x-orchestrator-token"]`` — the machine *identity* the allow-list checks. HF's edge
  strips ``x-hf-*``, so only this custom header reaches the gate.

Checks, in order (each prints one line; payloads are never dumped):

1. no identity header        → ``denied`` / ``no_token_header``  (with ``--anonymous``: no token at all,
                                for a public Space, which is what a stranger sends)
2. a token HF cannot resolve  → ``denied`` / ``token_unresolved``
3. ``view_api()`` named endpoints and their params == ``deploy/orchestrator_registration.yaml``
4. the caller's own identity  → ``ok`` on ``annotate_genes`` for ELAVL1
5. with ``--lookups``: one real lookup per endpoint → ``ok`` (keyless public sources, no model call)

Exit codes: 0 all passed, 1 a check failed, 2 could not reach the Space, 3 no token.

    python scripts/check_live_surface.py --space anne-voigt/pdac-knowledge-agent-dev --lookups
    python scripts/check_live_surface.py --space anne-voigt/pdac-knowledge-agent --anonymous
"""

from __future__ import annotations

import argparse
import json
import os
import sys
from pathlib import Path

import yaml

ROOT = Path(__file__).resolve().parent.parent
DESCRIPTOR = ROOT / "deploy" / "orchestrator_registration.yaml"
DEV_SPACE = "anne-voigt/pdac-knowledge-agent-dev"
# Shaped like an HF token so the gate asks HF about it, and HF says no.
_UNRESOLVABLE_TOKEN = "hf_" + "0" * 34

LOOKUPS = {
    "drugs_for_targets": {"genes": "KRAS"},
    "trials_for_genes": {"genes": "CD40", "condition": "pancreatic cancer", "status": "RECRUITING"},
    "annotate_genes": {"genes": "ELAVL1"},
    "pathways_for_genes": {"genes": "ISG15"},
    "targets_for_disease": {"disease": "MONDO_0005184", "size": "5"},
}


def _token() -> str | None:
    token = os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN")
    if token:
        return token
    try:
        from huggingface_hub import get_token

        return get_token()
    except Exception:  # noqa: BLE001 — an old hub, or no cache; the caller reports it
        return None


def _call(client, endpoint: str, **kwargs) -> dict:
    return json.loads(client.predict(**kwargs, api_name=f"/{endpoint}"))


def _describe(payload: dict) -> str:
    """The cause in a few named fields, never the rows."""
    if payload.get("status") == "ok":
        return (
            f"status='ok' rows={len(payload.get('rows', []))} "
            f"unresolved={len(payload.get('unresolved', []))} "
            f"source_errors={[e.get('reason') for e in payload.get('source_errors', [])]}"
        )
    keys = ("status", "machine_auth", "reason", "message")
    return " ".join(f"{k}={payload[k]!r}" for k in keys if k in payload)[:300]


def _expect_denied(client, why: str, label: str) -> bool:
    payload = _call(client, "annotate_genes", genes="ELAVL1")
    passed = payload.get("status") == "denied" and payload.get("machine_auth") == why
    print(f"{'PASS' if passed else 'FAIL'}  {label}: {_describe(payload)}")
    return passed


def _surface_matches(client) -> bool:
    descriptor = yaml.safe_load(DESCRIPTOR.read_text())
    expected = {t["api_name"]: list(t["params"]) for t in descriptor["tools"]}
    named = client.view_api(return_format="dict", print_info=False)["named_endpoints"]
    # Gradio's own endpoints start with "/_" (hf_oauth adds /_check_login_status); listed, not compared.
    internal = sorted(name for name in named if name.startswith("/_"))
    live = {
        name: [p.get("parameter_name") or p.get("label") for p in spec.get("parameters", [])]
        for name, spec in named.items()
        if name not in internal
    }
    passed = live == expected
    print(
        f"{'PASS' if passed else 'FAIL'}  view_api == descriptor: {sorted(live)}"
        f" (gradio internals: {internal})"
    )
    if not passed:
        for name in sorted(set(live) | set(expected)):
            if live.get(name) != expected.get(name):
                print(f"      {name}: live={live.get(name)} descriptor={expected.get(name)}")
    return passed


def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__.split("\n")[0])
    ap.add_argument("--space", default=DEV_SPACE, help=f"Space id (default {DEV_SPACE})")
    ap.add_argument(
        "--anonymous", action="store_true", help="check 1 sends no token at all (public Space)"
    )
    ap.add_argument("--lookups", action="store_true", help="one real lookup per endpoint")
    args = ap.parse_args()

    token = _token()
    if not token:
        print("check: no HF token (set HF_TOKEN or `hf auth login`)", file=sys.stderr)
        return 3

    from gradio_client import Client

    try:
        no_identity = Client(args.space, token=None if args.anonymous else token, verbose=False)
        unresolvable = Client(
            args.space,
            token=token,
            headers={"x-orchestrator-token": _UNRESOLVABLE_TOKEN},
            verbose=False,
        )
        caller = Client(
            args.space, token=token, headers={"x-orchestrator-token": token}, verbose=False
        )
    except Exception as exc:  # noqa: BLE001 — any failure here means "could not reach it"
        print(f"check: could not open {args.space}: {type(exc).__name__}: {exc}", file=sys.stderr)
        return 2

    print(f"check: {args.space}")
    results = [
        _expect_denied(
            no_identity,
            "no_token_header",
            "no token at all" if args.anonymous else "no identity header",
        ),
        _expect_denied(unresolvable, "token_unresolved", "unresolvable token"),
        _surface_matches(caller),
    ]
    authorized = _call(caller, "annotate_genes", genes="ELAVL1")
    results.append(authorized.get("status") == "ok")
    print(f"{'PASS' if results[-1] else 'FAIL'}  own identity: {_describe(authorized)}")

    if args.lookups:
        for endpoint, kwargs in LOOKUPS.items():
            payload = _call(caller, endpoint, **kwargs)
            results.append(payload.get("status") == "ok" and payload.get("kind") == "annotation")
            print(f"{'PASS' if results[-1] else 'FAIL'}  {endpoint}: {_describe(payload)}")

    return 0 if all(results) else 1


if __name__ == "__main__":
    sys.exit(main())