#!/usr/bin/env python3
"""Smoke-test the knowledge endpoints against the LIVE sources, and report the reliability week.

Phase 3 of the roadmap plan (``docs/plans/PLAN-phase3-review-hardening.md``): at least 95 % of
scheduled smoke probes succeed over one week, and every failure surfaces as a structured error,
never as prose. This script is the probe and the report; WHERE and HOW OFTEN it runs is the
project lead's decision (the plan's D1–D4), so it schedules nothing itself.

A run asks six fixed questions, one or more per source, chosen so all six sources the endpoints
reach are exercised (HGNC, Open Targets, ChEMBL, UniProt, ClinicalTrials.gov, Reactome). Each
probe is a recorded Phase 0 scenario, so ``tests/test_smoke_live.py`` replays the same probes
against the cassettes with sockets blocked: an expectation that the recording day does not meet
fails offline, before it can fail a scheduled run.

Two targets:

* ``--target local`` (default): the endpoints in-process, with the response cache OFF, so every
  call reaches the source. This is what the reliability week counts: the sources and this
  repo's handling of them. Keyless; no model call (``build_tooluniverse`` refuses to start with
  a model-provider key set, as the Space does).
* ``--target space``: the same probes through a deployed Space's machine endpoints, signed in
  the way the orchestrator signs in (``x-orchestrator-token``; the caller must be on the
  Space's allow-list). The Space keeps its own cache, so most rows come back ``cache: hit``:
  this checks the Space path, not the sources, and the record says how many rows were hits.

What a probe passes on: status as expected, ``kind: "annotation"``, every row's ten provenance
fields complete, no ``source_errors``, at least ``min_rows`` rows, and the stable entities the
probe names (KRAS → sotorasib; ELAVL1 → Q15717) present. A failure is classified: the three
payload source-error reasons, a payload ``error``, an unexpected status, too few rows, a missing
entity, a broken contract, or ``unstructured`` — an exception, a non-JSON reply or an unknown
status, which is the one kind the success measure forbids outright.

A record holds counts, classes and timings, never row content: the rows carry source licences
(ChEMBL is share-alike) and the record is not the place to redistribute them.

    python scripts/smoke_live.py                         # one local run, appended to .smoke/
    python scripts/smoke_live.py --target space --space anne-voigt/pdac-knowledge-agent-dev
    python scripts/smoke_live.py --report .smoke/results.jsonl --days 7

Exit codes: 0 every probe passed (``--report``: the window meets the measure), 1 a probe failed
(the window does not), 2 the harness could not run (no network, a refused startup, no token).
"""

from __future__ import annotations

import argparse
import datetime
import json
import math
import os
import subprocess
import sys
import time
from collections import Counter, defaultdict
from collections.abc import Callable, Iterable
from dataclasses import dataclass
from pathlib import Path
from typing import Any

ROOT = Path(__file__).resolve().parent.parent
if str(ROOT) not in sys.path:
    sys.path.insert(0, str(ROOT))

from src.payloads import SOURCE_ERROR_REASONS, STATUSES
from src.provenance import missing_fields

RECORD_SCHEMA = "knowledge-smoke/v1"
# A scheduled Job keeps no files, so it prints the record on one line behind this marker and
# `scripts/smoke_job.py fetch` reads it back out of the Job's logs.
RECORD_MARKER = "SMOKE_RECORD "
DEFAULT_OUT = ROOT / ".smoke" / "results.jsonl"
DEV_SPACE = "anne-voigt/pdac-knowledge-agent-dev"
TARGET_SUCCESS = 0.95
WEEK_DAYS = 7

# Failure classes, most specific first. The three source-error reasons are the payload contract's
# own (src.payloads.SOURCE_ERROR_REASONS); `unstructured` is the one the success measure forbids.
FAILURE_CLASSES = (
    *SOURCE_ERROR_REASONS,
    "payload_error",
    "unexpected_status",
    "contract",
    "too_few_rows",
    "entity_missing",
    "unstructured",
    "space_unreachable",
)


@dataclass(frozen=True)
class Probe:
    """One fixed question. ``cassette`` names the recorded scenario the offline test replays."""

    id: str
    endpoint: str
    args: tuple[str, ...]
    sources: tuple[str, ...]
    cassette: str
    expect_status: str = "ok"
    min_rows: int = 1
    entities: tuple[str, ...] = ()
    unresolved: tuple[str, ...] = ()
    why: str = ""


# Stable facts only: an entity named here must survive a normal week of source updates. Trials
# expect no row count because the recruiting list itself moves.
PROBES: tuple[Probe, ...] = (
    Probe(
        "drugs_kras",
        "drugs_for_targets",
        ("KRAS",),
        ("hgnc", "opentargets", "chembl"),
        "q1_q2_drugs_kras",
        entities=("SOTORASIB", "ADAGRASIB"),
        why="Open Targets drugs with ChEMBL mechanisms (Phase 0 Q1/Q2)",
    ),
    Probe(
        "annotate_elavl1",
        "annotate_genes",
        ("ELAVL1",),
        ("hgnc", "uniprot"),
        "q3_annotate_elavl1",
        entities=("Q15717",),
        why="a reviewed UniProt entry (Phase 0 Q3)",
    ),
    Probe(
        "trials_cd40",
        "trials_for_genes",
        ("CD40", "", ""),
        ("clinicaltrials",),
        "q5_trials_cd40",
        min_rows=0,
        why="ClinicalTrials.gov answers; the recruiting set itself moves (Phase 0 Q5)",
    ),
    Probe(
        "pathways_isg15",
        "pathways_for_genes",
        ("ISG15,RSAD2,OAS1",),
        ("hgnc", "reactome"),
        "q9_pathways_isg15_rsad2_oas1",
        entities=("R-HSA-909733",),
        why="Reactome membership: Interferon alpha/beta signaling (Phase 0 Q9)",
    ),
    Probe(
        "targets_pdac",
        "targets_for_disease",
        ("pancreatic ductal adenocarcinoma", ""),
        ("opentargets",),
        "q8_targets_pdac",
        entities=("KRAS",),
        why="Open Targets disease search and associations (Phase 0 Q8)",
    ),
    Probe(
        "unknown_symbol",
        "drugs_for_targets",
        ("NOTAGENE1",),
        ("hgnc",),
        "edge_unknown_symbol",
        min_rows=0,
        unresolved=("NOTAGENE1",),
        why="an unknown symbol is reported unresolved, as data (HGNC answers 'not found')",
    ),
)

Caller = Callable[[Probe], Any]


# --------------------------------------------------------------------------- #
# One probe
# --------------------------------------------------------------------------- #
def _failure(cls: str, detail: str, source: str = "") -> dict[str, str]:
    assert cls in FAILURE_CLASSES, cls
    return {"class": cls, "source": source, "detail": detail[:300]}


def evaluate(probe: Probe, payload: Any) -> tuple[dict[str, Any], dict[str, str] | None]:
    """(facts about the payload, the first failure or None). Pure: no I/O."""
    if not isinstance(payload, dict) or payload.get("status") not in STATUSES:
        shown = payload.get("status") if isinstance(payload, dict) else type(payload).__name__
        return {"status": None}, _failure("unstructured", f"not a contract payload: {shown!r}")

    rows = payload.get("rows") or []
    errors = payload.get("source_errors") or []
    cache = Counter((r.get("provenance") or {}).get("cache", "?") for r in rows)
    facts: dict[str, Any] = {
        "status": payload["status"],
        "payload_run_id": payload.get("run_id"),
        "n_rows": len(rows),
        "n_unresolved": len(payload.get("unresolved") or []),
        "cache": dict(cache),
        "sources_seen": sorted(
            {(r.get("provenance") or {}).get("source", "?") for r in rows}
            | {a.get("source", "?") for a in payload.get("attribution") or []}
        ),
        "source_errors": [
            {"source": e.get("source", ""), "reason": e.get("reason", "")} for e in errors
        ],
    }

    # A source error is the most specific cause, so it wins even on an `error` payload.
    for err in errors:
        reason = err.get("reason", "")
        if reason not in SOURCE_ERROR_REASONS:
            return facts, _failure("unstructured", f"unknown source-error reason {reason!r}")
    if errors:
        first = errors[0]
        return facts, _failure(first["reason"], first.get("detail", ""), first.get("source", ""))

    status = payload["status"]
    if status != probe.expect_status:
        cls = "payload_error" if status == "error" else "unexpected_status"
        why = payload.get("reason") or payload.get("message") or ""
        return facts, _failure(cls, f"status {status!r}, expected {probe.expect_status!r}: {why}")

    if status == "ok":
        if payload.get("kind") != "annotation":
            return facts, _failure("contract", f"kind {payload.get('kind')!r}")
        for i, row in enumerate(rows):
            gaps = missing_fields(row)
            if gaps:
                return facts, _failure("contract", f"row {i} lacks provenance {gaps}")

    if len(rows) < probe.min_rows:
        return facts, _failure("too_few_rows", f"{len(rows)} rows, expected ≥ {probe.min_rows}")

    text = json.dumps(rows, ensure_ascii=False).upper()
    absent = [e for e in probe.entities if e.upper() not in text]
    if absent:
        return facts, _failure("entity_missing", f"absent: {', '.join(absent)}")
    unresolved = {u.get("input", "").upper() for u in payload.get("unresolved") or []}
    missing = [u for u in probe.unresolved if u.upper() not in unresolved]
    if missing:
        return facts, _failure("entity_missing", f"not listed unresolved: {', '.join(missing)}")
    return facts, None


def run_probe(probe: Probe, call: Caller, clock: Callable[[], float] = time.monotonic) -> dict:
    """Ask one probe through ``call`` and grade it. An exception is a finding, never a crash."""
    started = clock()
    try:
        payload = call(probe)
    except SpaceUnreachable as exc:
        facts, failure = {"status": None}, _failure("space_unreachable", str(exc))
    except Exception as exc:  # noqa: BLE001 — anything escaping an endpoint is the finding itself
        facts, failure = {"status": None}, _failure("unstructured", f"{type(exc).__name__}: {exc}")
    else:
        facts, failure = evaluate(probe, payload)
    return {
        "id": probe.id,
        "endpoint": probe.endpoint,
        "sources": list(probe.sources),
        "passed": failure is None,
        "failure": failure,
        "seconds": round(clock() - started, 2),
        **facts,
    }


# --------------------------------------------------------------------------- #
# Targets
# --------------------------------------------------------------------------- #
class SpaceUnreachable(RuntimeError):
    """The Space did not answer at all (asleep, building, network): not a source failure."""


class _NoCache:
    """A cache that never hits: every local probe call reaches its source."""

    def get(self, source: str, tool: str, args: dict) -> None:
        return None

    def put(self, *args: Any, **kwargs: Any) -> None:
        return None


def local_caller() -> tuple[Caller, dict[str, str]]:
    """In-process endpoints on a fresh, cache-less client. Builds ToolUniverse once, up front."""
    from src.endpoints import ENDPOINTS, _common
    from src.tu_client import ToolClient, build_tooluniverse

    tu = build_tooluniverse()  # refuses with a model key present, as the Space does
    client = ToolClient(cache=_NoCache(), tooluniverse=tu)
    _common.set_client(client)

    def call(probe: Probe) -> Any:
        return ENDPOINTS[probe.endpoint][0](*probe.args)

    return call, {"target": "in-process", "tooluniverse_version": client.tooluniverse_version}


def space_caller(space: str, token: str) -> tuple[Caller, dict[str, str]]:
    """The deployed Space's machine endpoints, signed in as the orchestrator signs in."""
    from gradio_client import Client

    try:
        client = Client(space, token=token, headers={"x-orchestrator-token": token}, verbose=False)
    except Exception as exc:  # any failure to open it means it could not be reached
        raise SpaceUnreachable(f"{type(exc).__name__}: {exc}") from exc

    def call(probe: Probe) -> Any:
        try:
            raw = client.predict(*probe.args, api_name=f"/{probe.endpoint}")
        except Exception as exc:  # transport, not a payload
            raise SpaceUnreachable(f"{type(exc).__name__}: {exc}") from exc
        return json.loads(raw)  # a non-JSON reply raises here and is graded `unstructured`

    return call, {"target": space}


def _hf_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 _code_sha() -> str:
    if os.environ.get("SMOKE_CODE_SHA"):
        return os.environ["SMOKE_CODE_SHA"]
    try:
        out = subprocess.run(
            ["git", "-C", str(ROOT), "rev-parse", "--short=7", "HEAD"],
            capture_output=True,
            text=True,
            timeout=5,
            check=False,
        )
        return out.stdout.strip() or "unknown"
    except Exception:  # noqa: BLE001 — a Job image may have no git
        return "unknown"


def run(call: Caller, meta: dict[str, str], probes: Iterable[Probe] = PROBES) -> dict:
    """One smoke run: every probe, in order, graded. Returns the record (not yet written)."""
    started = datetime.datetime.now(datetime.UTC)
    results = [run_probe(p, call) for p in probes]
    passed = sum(r["passed"] for r in results)
    return {
        "schema": RECORD_SCHEMA,
        "started_at": started.strftime("%Y-%m-%dT%H:%M:%SZ"),
        "runner": os.environ.get("SMOKE_RUNNER", "manual"),
        "code_sha": _code_sha(),
        **meta,
        "passed": passed,
        "total": len(results),
        "probes": results,
    }


def append_record(record: dict, out: Path) -> None:
    out.parent.mkdir(parents=True, exist_ok=True)
    with out.open("a", encoding="utf-8") as fh:
        fh.write(json.dumps(record, ensure_ascii=False, sort_keys=True) + "\n")


def print_run(record: dict) -> None:
    print(f"smoke: {record['target']} @ {record['code_sha']} ({record['started_at']})")
    for r in record["probes"]:
        mark = "PASS" if r["passed"] else "FAIL"
        cache = ",".join(f"{k}={v}" for k, v in sorted(r.get("cache", {}).items())) or "-"
        line = f"{mark}  {r['id']:<16} {r['seconds']:>6.1f} s  rows={r.get('n_rows', '-')} cache={cache}"
        if r["failure"]:
            f = r["failure"]
            line += f"  [{f['class']}{'/' + f['source'] if f['source'] else ''}] {f['detail']}"
        print(line)
    print(f"{record['passed']}/{record['total']} probes passed")


# --------------------------------------------------------------------------- #
# The week
# --------------------------------------------------------------------------- #
def _parse_time(stamp: str) -> datetime.datetime:
    return datetime.datetime.strptime(stamp, "%Y-%m-%dT%H:%M:%SZ").replace(tzinfo=datetime.UTC)


def _nearest_rank(values: list[float], q: float) -> float:
    ordered = sorted(values)
    return ordered[max(0, math.ceil(q * len(ordered)) - 1)]


def load_records(path: Path) -> list[dict]:
    records = []
    for line in path.read_text(encoding="utf-8").splitlines():
        line = line.strip()
        if line:
            rec = json.loads(line)
            if rec.get("schema") == RECORD_SCHEMA:
                records.append(rec)
    return records


def summarise(
    records: list[dict], *, days: int = WEEK_DAYS, target: str = "in-process"
) -> dict[str, Any]:
    """The reliability measure over the ``days`` ending at the newest record of ``target``."""
    mine = [r for r in records if r.get("target") == target]
    if not mine:
        return {"target": target, "runs": 0, "verdict": "NO DATA"}
    end = max(_parse_time(r["started_at"]) for r in mine)
    start = end - datetime.timedelta(days=days)
    window = [r for r in mine if _parse_time(r["started_at"]) > start]
    probes = [p for r in window for p in r["probes"]]
    failures = [p for p in probes if not p["passed"]]
    by_class = Counter(p["failure"]["class"] for p in failures)
    unstructured = by_class.get("unstructured", 0)

    # A source is charged only with its own source errors; a missing entity or a contract break
    # counts against the overall measure, not against whichever source the probe touched.
    per_source: dict[str, list[int]] = defaultdict(lambda: [0, 0])  # [ok, total]
    for p in probes:
        failed = {e["source"] for e in p.get("source_errors") or []}
        for source in p["sources"]:
            per_source[source][1] += 1
            per_source[source][0] += source not in failed
    latency: dict[str, list[float]] = defaultdict(list)
    for p in probes:
        latency[p["id"]].append(p["seconds"])

    first = min(_parse_time(r["started_at"]) for r in window)
    span_days = (end - first).total_seconds() / 86400
    rate = (len(probes) - len(failures)) / len(probes) if probes else 0.0
    if rate >= TARGET_SUCCESS and unstructured == 0:
        verdict = "PASS" if span_days >= days - 0.5 else "ON TRACK (window not yet complete)"
    else:
        verdict = "FAIL"
    return {
        "target": target,
        "window": f"{first:%Y-%m-%d %H:%M} → {end:%Y-%m-%d %H:%M} UTC",
        "span_days": round(span_days, 2),
        "runs": len(window),
        "probes": len(probes),
        "passed": len(probes) - len(failures),
        "success": round(rate, 4),
        "failures_by_class": dict(by_class),
        "unstructured": unstructured,
        "per_source": {s: round(ok / n, 4) for s, (ok, n) in sorted(per_source.items())},
        "latency_s": {
            pid: {"p50": _nearest_rank(v, 0.5), "p95": _nearest_rank(v, 0.95), "n": len(v)}
            for pid, v in sorted(latency.items())
        },
        "code_shas": sorted({r["code_sha"] for r in window}),
        "verdict": verdict,
    }


def print_summary(s: dict) -> None:
    if not s["runs"]:
        print(f"report: no {s['target']} runs recorded")
        return
    print(f"report: {s['target']}, {s['window']} ({s['span_days']} days, {s['runs']} runs)")
    print(
        f"success {s['passed']}/{s['probes']} = {s['success']:.1%} "
        f"(target ≥ {TARGET_SUCCESS:.0%}); unstructured failures: {s['unstructured']}"
    )
    if s["failures_by_class"]:
        print("failures: " + ", ".join(f"{k} {v}" for k, v in s["failures_by_class"].items()))
    print("per source: " + ", ".join(f"{k} {v:.1%}" for k, v in s["per_source"].items()))
    for pid, lat in s["latency_s"].items():
        print(f"  {pid:<16} p50 {lat['p50']:.1f} s  p95 {lat['p95']:.1f} s  (n={lat['n']})")
    print(f"code: {', '.join(s['code_shas'])}")
    print(f"verdict: {s['verdict']}")


def main(argv: list[str] | None = None) -> int:
    ap = argparse.ArgumentParser(description=__doc__.split("\n")[0])
    ap.add_argument("--target", choices=("local", "space"), default="local")
    ap.add_argument("--space", default=DEV_SPACE, help=f"Space id for --target space ({DEV_SPACE})")
    ap.add_argument(
        "--out", type=Path, default=DEFAULT_OUT, help=f"JSONL to append ({DEFAULT_OUT})"
    )
    ap.add_argument("--no-write", action="store_true", help="print the run, append nothing")
    ap.add_argument("--report", type=Path, help="summarise this JSONL instead of running")
    ap.add_argument("--days", type=int, default=WEEK_DAYS)
    ap.add_argument("--report-target", default="in-process", help="which target to summarise")
    ap.add_argument("--json", action="store_true", help="print the record or summary as JSON")
    ap.add_argument(
        "--print-record",
        action="store_true",
        help=f"also print the record as one {RECORD_MARKER!r} line",
    )
    args = ap.parse_args(argv)

    if args.report:
        summary = summarise(load_records(args.report), days=args.days, target=args.report_target)
        print(json.dumps(summary, indent=1, ensure_ascii=False)) if args.json else print_summary(
            summary
        )
        return 0 if summary["verdict"].startswith(("PASS", "ON TRACK")) else 1

    try:
        if args.target == "local":
            call, meta = local_caller()
        else:
            token = _hf_token()
            if not token:
                print("smoke: no HF token (set HF_TOKEN or `hf auth login`)", file=sys.stderr)
                return 2
            call, meta = space_caller(args.space, token)
    except Exception as exc:  # noqa: BLE001 — a refused startup or an unreachable Space
        print(f"smoke: could not start ({type(exc).__name__}: {exc})", file=sys.stderr)
        return 2

    record = run(call, meta)
    if args.json:
        print(json.dumps(record, indent=1, ensure_ascii=False))
    else:
        print_run(record)
    if args.print_record:
        print(RECORD_MARKER + json.dumps(record, ensure_ascii=False, sort_keys=True), flush=True)
    if not args.no_write:
        append_record(record, args.out)
        print(f"appended to {args.out}")
    return 0 if record["passed"] == record["total"] else 1


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