#!/usr/bin/env python3
"""Fiction-wing retrieval probe harness — the recipe R42 (brokkr-smithy-dev) builds against
and the re-runnable #389 acceptance gate.

Two retrieval paths, kept SEPARATE because they have different noise properties:

  * search_library (raw, RANKING-clean): drive the `mimir` foundational agent (all-wing
    librarian) with a FIXED query string. Deterministic against a fixed index — use it for
    ranking baselines (R42 arm-2). No LLM in the query loop.
  * reference_knowledge (the Tier-3 consumer path): drive `ratatoskr:donut`; her reasoning
    REFORMULATES the query each turn, so this path carries QUERY-FORMULATION variance
    (the arm-4 signal), attributable via the captured tool_start query. NOT for ranking numbers.

Scoring conventions (identical across both paths and all R42 arms):
  * confidence BUCKET vs WT's RRF thresholds: high >= 0.030, medium >= 0.016, low < 0.016.
  * ON-TARGET (load-bearing): a returned row is on-target iff its excerpt actually NAMES or
    describes the queried subject (keyword match on the subject's distinctive tokens). The
    failure signature "10 hits / MEDIUM / 0 on-target" = present-by-topic, subject absent —
    the split that separated #384 (packaging) / #387 (coverage) / #389 (ranking).
  * MISS = no on-target row in the returned top-k.

Noise floor: freeze the generation (pin the b-tag) to remove extraction variance; fixed-string
search_library is deterministic (no CI needed); reference_knowledge variance is query-
formulation, not floor noise. Residual = bucket-boundary sensitivity at 0.016/0.030 — so probe
N>=3-5 times per term and report the bucket DISTRIBUTION, never a single-run point label.

Session protocol: ONE fresh session per query. A reused mimir session returns EMPTY
search_library results after its first turn (Worldtree #391), silently scoring later terms
false-MISS; the retrieval helpers self-session to enforce it. Never hoist the session out.

Config from env (source ratatoskr's env.sh): WORLDTREE_API_URL, WORLDTREE_API_KEY,
RATATOSKR_END_USER_ID. No secrets are stored here.

Usage:
  uv run python docs/diagnostics/fiction_wing_probe.py                 # run the artifact yardstick
  uv run python docs/diagnostics/fiction_wing_probe.py --runs 5        # N repeats -> bucket distribution
  uv run python docs/diagnostics/fiction_wing_probe.py --term "Enhanced Pet Biscuit" --keywords biscuit
"""

from __future__ import annotations

import argparse
import json
import os
import unicodedata
from collections import Counter

import httpx

HIGH, MEDIUM = 0.030, 0.016  # WT RRF confidence thresholds


def _cfg() -> tuple[str, dict, str]:
    base = os.environ.get("WORLDTREE_API_URL", "http://10.250.50.152:8081")
    key = os.environ.get("WORLDTREE_API_KEY")
    if not key:
        raise SystemExit("WORLDTREE_API_KEY unset — source env.sh first.")
    end_user = os.environ.get("RATATOSKR_END_USER_ID", "ratatoskr-tui")
    return base, {"Authorization": f"Bearer {key}"}, end_user


def _bucket(score: float | None) -> str:
    if score is None:
        return "none"
    return "high" if score >= HIGH else "medium" if score >= MEDIUM else "low"


# Curly punctuation the b170 extraction emits (U+2019 etc.) folded to ASCII so a
# possessive-named subject ("Darcy's letter") matches regardless of quote style.
_QUOTE_FOLD = str.maketrans({
    "‘": "'", "’": "'",  # noqa: RUF001 - single curly quotes / apostrophe
    "“": '"', "”": '"',  # double curly quotes
    "′": "'", "″": '"',  # noqa: RUF001 - primes
})


def _fold(s: str) -> str:
    """NFKC-normalize, fold curly quotes/apostrophes to ASCII, lowercase.
    NFKC alone does NOT fold U+2019, so the explicit quote-fold is load-bearing."""
    return unicodedata.normalize("NFKC", s or "").translate(_QUOTE_FOLD).lower()


def _on_target(excerpt: str, keywords: list[str]) -> bool:
    ex = _fold(excerpt)
    return any(_fold(k) in ex for k in keywords)


def _session(base: str, headers: dict, agent_id: str, end_user: str) -> str:
    r = httpx.post(f"{base}/sessions", json={"agent_id": agent_id, "end_user_id": end_user},
                   headers=headers, timeout=30)
    return r.json()["session_id"]


def _drive(base: str, headers: dict, sid: str, content: str) -> tuple[str | None, dict]:
    """POST a turn, return (tool_query, tool_result_dict). tool_result is the first tool packet."""
    tool_query, result = None, {}
    with httpx.stream("POST", f"{base}/sessions/{sid}/messages", json={"content": content},
                      headers=headers, timeout=120) as r:
        for line in r.iter_lines():
            if not line.startswith("data: "):
                continue
            ev = json.loads(line[6:])
            t = ev.get("type")
            if t == "tool_start" and tool_query is None:
                tool_query = (ev.get("arguments") or {}).get("query")
            elif t == "tool_result" and not result:
                result = ev.get("result") if isinstance(ev.get("result"), dict) else {}
            elif t == "done":
                break
    return tool_query, result


def search_library(base, headers, end_user, term, keywords):
    """RANKING-clean path: fixed-string search over the mimir librarian. Deterministic.

    Opens a FRESH mimir session per call — REQUIRED. mimir stops returning
    search_library results after the first turn on a reused session (Worldtree #391),
    silently scoring every later term a false-MISS; do not hoist the session to the caller.
    """
    msid = _session(base, headers, "mimir", end_user)
    _, res = _drive(base, headers, msid, f"Use search_library to find: {term}")
    rows = res.get("results", []) if isinstance(res, dict) else []
    on = [h for h in rows if isinstance(h, dict) and _on_target(h.get("excerpt", ""), keywords)]
    top = on[0] if on else None
    score = round(top["score"], 4) if top else None
    return {"n": len(rows), "on_target": len(on), "hit": bool(on),
            "score": score, "bucket": _bucket(score),
            "excerpt": (top.get("excerpt", "")[:140] if top else None)}


def reference_knowledge(base, headers, end_user, question, keywords):
    """Consumer path: Donut reformulates -> capture her tool_query. NOT for ranking numbers.

    Fresh donut session per call (same reuse-degradation guard as search_library, and
    it keeps each run an independent first-turn sample rather than a growing conversation).
    """
    dsid = _session(base, headers, "ratatoskr:donut", end_user)
    q, res = _drive(base, headers, dsid, question)
    hits = res.get("hits", []) if isinstance(res, dict) else []
    on = [h for h in hits if isinstance(h, dict) and _on_target(h.get("excerpt", ""), keywords)]
    return {"tool_query": q, "n": len(hits), "on_target": len(on),
            "confidence": res.get("confidence") if isinstance(res, dict) else None}


# Artifact yardstick — worldtree-dev grep-confirmed in DCC book-1. Frozen arm-2 baseline.
YARDSTICK = [
    ("Enchanted Crown of the Sepsis Whore", "Crown of the Sepsis Whore", ["sepsis", "crown"]),
    ("Enhanced Pet Biscuit", "Pet Biscuit", ["biscuit"]),
    ("Enchanted BigBoi Boxers", "BigBoi Boxers", ["boxers", "bigboi"]),
    ("Enchanted Toe Ring of the Splatter Skunk", "Toe Ring of the Splatter Skunk",
     ["toe ring", "splatter", "skunk"]),
]


def run_yardstick(runs: int) -> None:
    base, headers, end_user = _cfg()
    print(f"# Fiction-wing ranking yardstick  (search_library, {runs} run(s) per name)\n")
    for full, partial, kw in YARDSTICK:
        for label, term in (("full ", full), ("part ", partial)):
            buckets, hits = Counter(), 0
            for _ in range(runs):
                r = search_library(base, headers, end_user, term, kw)
                buckets[r["bucket"]] += 1
                hits += r["hit"]
            dist = " ".join(f"{b}:{c}" for b, c in buckets.most_common())
            print(f"  [{label}] {term:<42} hit {hits}/{runs}  buckets({dist})")
        print()


def run_term(term: str, keywords: list[str], runs: int) -> None:
    base, headers, end_user = _cfg()
    print(f"# Probe: {term!r}  ({runs} run(s))\n")
    sl_buckets, sl_hits = Counter(), 0
    for _ in range(runs):
        r = search_library(base, headers, end_user, term, keywords)
        sl_buckets[r["bucket"]] += 1
        sl_hits += r["hit"]
    print(f"  search_library : hit {sl_hits}/{runs}  buckets({dict(sl_buckets)})")
    for _ in range(runs):
        rk = reference_knowledge(base, headers, end_user, f"Tell me about the {term}.", keywords)
        print(f"  reference_knowledge: conf={rk['confidence']} on_target={rk['on_target']}"
              f"  (donut query: {rk['tool_query']!r})")


def main() -> None:
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("--runs", type=int, default=1, help="repeats per term (>=3-5 near the floor)")
    ap.add_argument("--term", help="probe a single term instead of the yardstick")
    ap.add_argument("--keywords", nargs="*", default=[], help="on-target keywords for --term")
    ns = ap.parse_args()
    if ns.term:
        run_term(ns.term, ns.keywords or [ns.term.split()[-1]], ns.runs)
    else:
        run_yardstick(ns.runs)


if __name__ == "__main__":
    main()