"""Planner eval runner (E-planner).

Feeds each golden case in `planner_dataset.json` to the LIVE planner
(`PlannerService.plan`) against the `PA Data Dummy` fixture catalog, then scores
each case by RULE-COMPLIANCE assertions on the emitted plan/IR (not exact IR
match — many IRs are valid; we only pin the properties that matter). Records
latency + token usage, prints a per-case + aggregate summary, and writes a
timestamped JSON report under `results/` (never overwritten — diff runs over
time).

Run before any deploy that touches planner.md or examples.py:

    uv run python -m eval.planner.run_eval
    uv run python -m eval.planner.run_eval --limit 6     # quick smoke test

Needs Azure OpenAI creds in the env (same as the live planner). The scoring
layer is deterministic and unit-tested via `--selfcheck` (no LLM call).
"""

from __future__ import annotations

import argparse
import asyncio
import inspect
import json
import statistics
import time
from dataclasses import asdict, dataclass, field
from datetime import date, datetime
from functools import cache
from pathlib import Path
from typing import Any

from langchain_core.callbacks import BaseCallbackHandler
from langchain_core.outputs import LLMResult

from src.agents.planner.contracts import BusinessContext
from src.agents.planner.groups import scope
from src.agents.planner.inputs import Constraints
from src.agents.planner.registry import default_registry
from src.agents.planner.service import PlannerService

from ..run_log import run_stamp, start_run_log
from .catalog_fixture import (
    build_db_catalog,
    build_db_dict_catalog,
    build_ema_catalog,
    build_pa_catalog,
    db_name_to_id,
    dict_name_to_id,
    id_to_name,
    name_to_id,
    with_documents,
)

_HERE = Path(__file__).resolve().parent
DATASET = _HERE / "planner_dataset.json"
RESULTS_DIR = _HERE / "results"

# A3 (2026-09-23): cases carrying `document_context` exercise retrieve-then-plan
# (PR23 part 2, E1). Until `PlannerService.plan` accepts it they are SKIPPED, not
# failed — the same rule as a withheld tool.
_PLANNER_TAKES_DOCUMENT_CONTEXT = "document_context" in inspect.signature(
    PlannerService.plan
).parameters


def _all_name_to_id() -> dict[str, str]:
    """Column name -> id across ALL fixtures (their column names do not collide).

    The EMA fixture (2026-10-22) is left out on purpose: its tables share column
    names (`site_id`, `pa`) with each other and with the DB fixture. Its cases use the
    name-based assertions (`filter_col`, `group_by_cols`, `reads_table`, ...) instead.
    """
    return {**name_to_id(), **db_name_to_id(), **dict_name_to_id()}


@cache
def _id2name() -> dict[str, str]:
    """Column/table id -> name across every fixture, for the name-based assertions."""
    return id_to_name(build_pa_catalog(), build_db_catalog(), build_db_dict_catalog(),
                      build_ema_catalog())


_CONTEXT = BusinessContext(
    project_id="eval-planner",
    industry="mining",
    completeness="partial",
    business_description=(
        "Physical Availability (PA) analysis of heavy mining equipment — haulers "
        "and loaders — from daily operational records."
    ),
    scale_and_scope="9,729 daily equipment records, single site, April 2026.",
)

#: For `catalog: ema` cases (2026-10-22). `_CONTEXT` describes the PA fixture —
#: "single site, April 2026" — and EMA cases planned against it filtered April 2026
#: from it. Their own context says only what the EMA fixture is.
_EMA_CONTEXT = BusinessContext(
    project_id="eval-planner",
    industry="mining",
    completeness="partial",
    business_description=(
        "Maintenance performance of heavy mining equipment: monthly PA, bad-actor "
        "flags and unscheduled-breakdown causes per unit."
    ),
    scale_and_scope="Several sites, monthly records, Oct 2025 - Sep 2026.",
)


# --------------------------------------------------------------------------- #
# Plan introspection + assertion scoring (deterministic — no LLM)
# --------------------------------------------------------------------------- #

# Grouping/agg can be expressed EITHER in the retrieve_data IR (group_by + agg
# select) OR via the analyze_aggregate tool (group_by + aggregations args) — both
# are valid (planner.md R2). Checks below look at both. IR uses fn "avg"; the
# aggregate tool uses "mean" — treat as synonyms.
_AGG_SYN = {
    "avg": {"avg", "mean"}, "mean": {"avg", "mean"}, "sum": {"sum"},
    "count": {"count"}, "min": {"min"}, "max": {"max"},
    "count_distinct": {"count_distinct", "nunique"},
}


def extract_facts(task_list: Any) -> dict:
    """Flatten a TaskList into the plain facts the scorer needs (JSON-safe, so a
    run can be re-scored offline via --rescore without re-calling the LLM)."""
    irs: list[dict] = []
    agg_args: list[dict] = []
    tools: set[str] = set()
    for t in task_list.tasks:
        for c in t.tool_calls:
            tools.add(c.tool)
            if c.tool == "retrieve_data" and isinstance(c.args.get("ir"), dict):
                irs.append(c.args["ir"])
            elif c.tool == "analyze_aggregate":
                agg_args.append(c.args)
    return {
        "tools": sorted(tools),
        "irs": irs,
        "agg_args": agg_args,
        "infeasible": (not task_list.tasks) or bool(getattr(task_list, "infeasible_reason", None)),
    }


def _ir_agg_fns(ir: dict) -> list[str]:
    return [s.get("fn") for s in ir.get("select", []) if isinstance(s, dict) and s.get("kind") == "agg"]


def _all_agg_fns(f: dict) -> list[str]:
    fns = [fn for ir in f["irs"] for fn in _ir_agg_fns(ir)]
    for a in f["agg_args"]:
        for lst in (a.get("aggregations") or {}).values():
            fns += lst if isinstance(lst, list) else [lst]
    return fns


def _group_by_ids(f: dict) -> list[str]:
    return [g for ir in f["irs"] for g in (ir.get("group_by") or [])]


def _group_by_aliases(f: dict) -> list[str]:
    return [str(g) for a in f["agg_args"] for g in (a.get("group_by") or [])]


def _alias_to_id(f: dict) -> dict[str, str]:
    """Map each retrieve_data SELECT alias -> its column_id, so an
    analyze_aggregate group_by (which references aliases) resolves back to a
    real column even when the planner aliases it differently from the name."""
    m: dict[str, str] = {}
    for ir in f["irs"]:
        for s in ir.get("select", []):
            if isinstance(s, dict) and s.get("alias") and s.get("column_id"):
                m[s["alias"]] = s["column_id"]
    return m


def _filter_ops(f: dict) -> list[str]:
    return [flt.get("op") for ir in f["irs"] for flt in ir.get("filters", []) if isinstance(flt, dict)]


def _selected_col_ids(f: dict) -> list[str]:
    """Every column_id referenced in any retrieve_data select (column or agg).
    Used to catch the wrong-column bug where the planner selects Plan_PA_Percent
    but aliases it 'pa_percent' — the alias hides it, the column_id doesn't."""
    return [s["column_id"] for ir in f["irs"] for s in ir.get("select", [])
            if isinstance(s, dict) and s.get("column_id")]


def evaluate_facts(f: dict, expect: dict, n2id: dict[str, str]) -> list[tuple[str, bool, str]]:
    """Return [(check, passed, detail)] for every assertion. Grouping/agg checks
    honor BOTH the IR and the analyze_aggregate tool."""
    irs, tools = f["irs"], set(f["tools"])
    grouped = bool(_group_by_ids(f) or _group_by_aliases(f))
    res: list[tuple[str, bool, str]] = []

    for key, want in expect.items():
        if key == "has_tool":
            res.append((f"has_tool={want}", want in tools, f"tools={sorted(tools)}"))
        elif key == "no_tool":
            res.append((f"no_tool={want}", want not in tools, f"tools={sorted(tools)}"))
        elif key == "any_tool":  # at least one of these tools present
            res.append((f"any_tool={want}", any(t in tools for t in want), f"tools={sorted(tools)}"))
        elif key == "selects_col":  # a select references this column (by id, not the alias)
            col_id = n2id.get(want, want)
            ids = _selected_col_ids(f)
            res.append((f"selects_col={want}", col_id in ids, f"selected={ids}"))
        elif key == "not_selects_col":  # this column must NOT be selected (wrong-column guard)
            col_id = n2id.get(want, want)
            ids = _selected_col_ids(f)
            res.append((f"not_selects_col={want}", col_id not in ids, f"selected={ids}"))
        elif key == "select_agg":
            syn = _AGG_SYN.get(want, {want})
            got = _all_agg_fns(f)
            res.append((f"select_agg={want}", any(g in syn for g in got), f"aggs={got}"))
        elif key == "group_by":
            res.append(("group_by" if want else "no_group_by", grouped == want, f"grouped={grouped}"))
        elif key == "chart_aggregated":  # chart data is aggregated somehow (group_by IR OR any analyze_* step), not raw rows
            analyze = [t for t in tools if t.startswith("analyze_")]
            ok = grouped or bool(analyze)
            res.append(("chart_aggregated" if want else "chart_raw", ok == want, f"grouped={grouped} analyze={analyze}"))
        elif key == "group_by_col":
            col_id = n2id.get(want, want)
            a2id = _alias_to_id(f)
            resolved = _group_by_ids(f) + [a2id.get(a) for a in _group_by_aliases(f)]
            hit = col_id in resolved or want.lower() in [a.lower() for a in _group_by_aliases(f)]
            res.append((f"group_by_col={want}", hit, f"ids={_group_by_ids(f)} aliases={_group_by_aliases(f)} resolved={[r for r in resolved if r]}"))
        elif key == "filter_op":
            got = _filter_ops(f)
            res.append((f"filter_op={want}", want in got, f"ops={got}"))
        elif key == "no_filter_op":
            got = _filter_ops(f)
            res.append((f"no_filter_op={want}", want not in got, f"ops={got}"))
        elif key == "has_filter":
            has = any(ir.get("filters") for ir in irs)
            res.append(("has_filter", has == want, f"filter_present={has}"))
        elif key == "order_dir":
            got = [o.get("dir", "asc") for ir in irs for o in ir.get("order_by", []) if isinstance(o, dict)]
            res.append((f"order_dir={want}", want in got, f"dirs={got}"))
        elif key == "limit":
            got = [ir.get("limit") for ir in irs]
            res.append((f"limit={want}", want in got, f"limits={got}"))
        elif key == "has_join":  # A2: some IR joins another table (or none does)
            has = any(ir.get("joins") for ir in irs)
            res.append(("has_join" if want else "no_join", has == want, f"join_present={has}"))
        elif key == "filter_value":  # A3: some IR filter compares against this value
            got = [flt.get("value") for ir in irs
                   for flt in ir.get("filters", []) if isinstance(flt, dict)]
            numeric = isinstance(want, int | float)
            hit = any(g == want or (numeric and isinstance(g, int | float) and abs(g - want) < 1e-9)
                      for g in got)
            res.append((f"filter_value={want}", hit, f"values={got}"))
        elif key == "infeasible":
            res.append(("infeasible" if want else "feasible", f["infeasible"] == want, f"infeasible={f['infeasible']}"))
        # ---- name-based (2026-10-22, EMA fixture): resolve ids to names first ----
        elif key in ("filter_col", "no_filter_col"):
            got = [_id2name().get(flt.get("column_id"), flt.get("column_id")) for ir in irs
                   for flt in ir.get("filters", []) if isinstance(flt, dict)]
            ok = (want in got) == (key == "filter_col")
            res.append((f"{key}={want}", ok, f"filter_cols={got}"))
        elif key == "no_literal_filter_col":  # never filtered by a TYPED value; a
            # value-handoff placeholder ("${t1.site_id}") is the correct lookup
            got = [(_id2name().get(flt.get("column_id"), flt.get("column_id")), flt.get("value"))
                   for ir in irs for flt in ir.get("filters", []) if isinstance(flt, dict)]
            def _handoff(v: Any) -> bool:  # "${t1.x}", or a list holding one
                vs = v if isinstance(v, list) else [v]
                return any(isinstance(x, str) and x.startswith("${") for x in vs)
            typed = [v for n, v in got if n == want and not _handoff(v)]
            res.append((f"no_literal_filter_col={want}", not typed, f"typed_values={typed}"))
        elif key == "group_by_cols":  # every listed name is grouped somewhere (IR or tool)
            got = {_id2name().get(g, g) for g in _group_by_ids(f)}
            got |= {_id2name().get(_alias_to_id(f).get(a, a), a) for a in _group_by_aliases(f)}
            res.append((f"group_by_cols={want}", set(want) <= got, f"grouped={sorted(got)}"))
        elif key == "reads_table":  # some IR reads (or joins) this table, by name
            got = [_id2name().get(t, t) for ir in irs
                   for t in [ir.get("table_id"), *[j.get("target_table_id") for j in ir.get("joins") or []]]]
            res.append((f"reads_table={want}", want in got, f"tables={got}"))
        elif key == "ir_group_by":  # grouping done IN the query, not by an analyze_* step
            has = any(ir.get("group_by") for ir in irs)
            res.append(("ir_group_by" if want else "no_ir_group_by", has == want, f"ir_grouped={has}"))
        elif key == "no_limit":  # no IR uses this limit (e.g. 1 hides ties)
            got = [ir.get("limit") for ir in irs]
            res.append((f"no_limit={want}", want not in got, f"limits={got}"))
        elif key == "no_filter_value":
            got = [flt.get("value") for ir in irs
                   for flt in ir.get("filters", []) if isinstance(flt, dict)]
            res.append((f"no_filter_value={want}", want not in got, f"values={got}"))
        elif key == "selects_name":  # a select references a column with this NAME
            got = [_id2name().get(c, c) for c in _selected_col_ids(f)]
            res.append((f"selects_name={want}", want in got, f"selected={got}"))
        else:
            res.append((f"UNKNOWN:{key}", False, "unknown expect key"))
    return res


# --------------------------------------------------------------------------- #
# Token callback (parity with intent eval)
# --------------------------------------------------------------------------- #

class _TokenCounter(BaseCallbackHandler):
    def __init__(self) -> None:
        self.total = 0
        self.input = 0
        # A6 (2026-09-23): prompt tokens served from the provider's prefix cache.
        # langchain-openai maps prompt_tokens_details.cached_tokens here.
        self.cache_read = 0

    def on_llm_end(self, response: LLMResult, **kwargs: Any) -> None:
        for gen_list in response.generations:
            for gen in gen_list:
                msg = getattr(gen, "message", None)
                usage = getattr(msg, "usage_metadata", None) if msg else None
                if usage:
                    self.total += usage.get("total_tokens", 0)
                    self.input += usage.get("input_tokens", 0)
                    details = usage.get("input_token_details") or {}
                    self.cache_read += details.get("cache_read", 0) or 0


# --------------------------------------------------------------------------- #
# Runner
# --------------------------------------------------------------------------- #

def skip_reason(case: dict, active: set[str]) -> str | None:
    """Why this case can't run against the CURRENT registry, or None if it can.

    A case's tool requirement is DERIVED from its own assertions rather than a
    hand-kept field, so the dataset stays self-maintaining:
      - `has_tool: X`      -> X must be planner-visible, else the case is unrunnable.
      - `any_tool: [...]`  -> only unrunnable when EVERY listed tool is withheld.
      - `no_tool`          -> asserts ABSENCE; never a reason to skip.

    A withheld tool is a registry state (`DEFERRED_ANALYTICS_TOOLS`), not a planner
    regression, so the case is reported `skipped` and dropped from the denominator —
    otherwise deferring a tool turns its cases red and masks a real regression in the
    same run. The case stays in the dataset and comes back automatically the moment
    the tool is moved to ACTIVE.
    """
    if case.get("document_context") and not _PLANNER_TAKES_DOCUMENT_CONTEXT:
        return "needs retrieve-then-plan (PlannerService.plan has no document_context yet)"
    expect = case.get("expect", {})
    required = expect.get("has_tool")
    if required and required not in active:
        return f"requires {required}, not in the active registry"
    alternatives = expect.get("any_tool") or []
    if alternatives and not (set(alternatives) & active):
        return f"requires one of {sorted(alternatives)}, none in the active registry"
    return None


@dataclass
class CaseResult:
    id: str
    category: str
    lang: str
    carried_over: bool
    question: str
    passed: bool
    checks: list[dict] = field(default_factory=list)
    facts: dict = field(default_factory=dict)  # extracted plan (for offline --rescore)
    error: str | None = None
    latency_ms: int = 0
    tokens: int = 0
    input_tokens: int = 0
    cache_read: int = 0
    # Set when the case's required tool is withheld from the Planner. A skipped case
    # is neither passed nor failed — it leaves the denominator entirely.
    skipped: bool = False
    skip_reason: str | None = None
    # Added 2026-09-23 — a run is only reviewable if the plans themselves survive.
    # `plan` is the TaskList the case was scored on (the infeasible sentinel when the
    # planner exhausted its retries); `attempts` is every TaskList the model returned,
    # in order, with the error that rejected it (null on the accepted one). `facts`
    # stays: --rescore reads it.
    plan: dict | None = None
    attempts: list[dict] = field(default_factory=list)


class _AttemptRecorder:
    """Stands in as the planner's chain AND validator, so each attempt is kept.

    Wraps the live ones unchanged: the planner under test is the production
    `PlannerService` — only what it produced and why it retried are written down.
    """

    def __init__(self) -> None:
        from src.agents.planner.service import _get_default_chain
        from src.agents.planner.validator import PlannerValidator

        self._chain = _get_default_chain()
        self._validator = PlannerValidator()
        self.attempts: list[dict] = []

    async def ainvoke(self, inp: Any, config: Any = None) -> Any:
        try:
            out = await (
                self._chain.ainvoke(inp, config=config) if config else self._chain.ainvoke(inp)
            )
        except Exception as e:  # schema failure — the service retries it
            self.attempts.append({"plan": None, "error": f"schema: {type(e).__name__}: {e}"})
            raise
        self.attempts.append({"plan": out.model_dump(mode="json"), "error": None})
        return out

    def validate(self, task_list: Any, tools: Any, catalog: Any, constraints: Any) -> None:
        try:
            self._validator.validate(task_list, tools, catalog, constraints)
        except Exception as e:
            self.attempts[-1]["error"] = str(e)
            raise


async def _run_case(catalog: Any, tools: Any, n2id: dict, case: dict, **plan_kw: Any) -> CaseResult:
    # A case may carry the analysis objective (2026-10-22) and a fixed `today`; every
    # other case keeps the plain context. `today` defaults to the real date, which the
    # live planner also sees.
    context = _EMA_CONTEXT if case.get("catalog") == "ema" else _CONTEXT
    if case.get("objective"):
        context = context.model_copy(update={"analysis_objective": case["objective"]})
    if case.get("today"):
        plan_kw["today"] = date.fromisoformat(case["today"])
    tok = _TokenCounter()
    started = time.perf_counter()
    facts: dict = {}
    plan: dict | None = None
    rec = _AttemptRecorder()
    planner = PlannerService(structured_chain=rec, validator=rec)
    try:
        task_list = await planner.plan(
            context, catalog, tools, case["question"], Constraints(), callbacks=[tok], **plan_kw
        )
        plan = task_list.model_dump(mode="json")
        facts = extract_facts(task_list)
        checks = evaluate_facts(facts, case["expect"], n2id)
        passed = all(ok for _, ok, _ in checks)
        err = None
    except Exception as e:  # planner failure = case fails (record why)
        checks, passed, err = [], False, f"{type(e).__name__}: {e}"
    latency = int((time.perf_counter() - started) * 1000)
    return CaseResult(
        id=case["id"], category=case["category"], lang=case["lang"],
        carried_over=case.get("carried_over", False), question=case["question"],
        passed=passed, error=err, latency_ms=latency, tokens=tok.total, facts=facts,
        input_tokens=tok.input, cache_read=tok.cache_read, plan=plan, attempts=rec.attempts,
        checks=[{"check": c, "ok": ok, "detail": d} for c, ok, d in checks],
    )


async def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--limit", type=int, default=None, help="run only the first N cases")
    ap.add_argument("--only", default=None, help="comma-separated case ids to run (smoke a subset)")
    ap.add_argument("--selfcheck", action="store_true", help="test the scorer on a synthetic plan (no LLM)")
    ap.add_argument("--rescore", metavar="RESULTS.json", help="re-score a saved run's facts with the current assertions (no LLM)")
    args = ap.parse_args()

    if args.selfcheck:
        _selfcheck()
        return
    if args.rescore:
        _rescore(Path(args.rescore))
        return

    # A paid run: its console output (rejected attempts, validator errors) is kept
    # beside the result file, same stamp.
    stamp = run_stamp()
    start_run_log(RESULTS_DIR, "planner_run", stamp)
    data = json.loads(DATASET.read_text(encoding="utf-8"))
    cases = data["cases"][: args.limit] if args.limit else data["cases"]
    if args.only:
        wanted = {s.strip() for s in args.only.split(",") if s.strip()}
        cases = [c for c in cases if c["id"] in wanted]
    catalogs = {
        "pa": build_pa_catalog(),
        "db": build_db_catalog(),
        # KM-717 (2026-09-30): the DB fixture with data-dictionary entries.
        "db_dict": build_db_dict_catalog(),
        "db_dict_injected": build_db_dict_catalog(injected=True),
        # 2026-10-22: EMA-shaped DB (large tables, site lookup, pareto causes).
        "ema": build_ema_catalog(),
    }
    tools, n2id = default_registry(), _all_name_to_id()

    active = tools.names()

    results: list[CaseResult] = []
    for case in cases:
        # Withheld tool -> skip before spending an LLM call on a plan that cannot
        # satisfy the assertion anyway.
        reason = skip_reason(case, active)
        if reason is not None:
            r = CaseResult(
                id=case["id"], category=case["category"], lang=case["lang"],
                carried_over=case.get("carried_over", False), question=case["question"],
                passed=False, skipped=True, skip_reason=reason,
            )
            print(f"[SKIP] {r.id:<28} {r.lang}  {reason}")
            results.append(r)
            continue
        # A case may bind documents (UNSTRUCTURED_V2_PLAN U7); every other case plans
        # against the unchanged fixture, so its prompt is byte-identical to before.
        # Per case (2026-09-23, PR23 part 2): which fixture (A2), bound documents,
        # the router's group hint applied exactly as live (A1: `groups.scope`, fail-open
        # on null), and document context (A3, once the planner accepts it). A case
        # without these fields plans exactly as before.
        case_catalog = catalogs[case.get("catalog", "pa")]
        docs = case.get("documents")
        if docs:
            case_catalog = with_documents(case_catalog, docs)
        case_tools = scope(tools, case["group_hint"]) if case.get("group_hint") else tools
        plan_kw = {}
        if case.get("document_context") and _PLANNER_TAKES_DOCUMENT_CONTEXT:
            plan_kw["document_context"] = case["document_context"]
        r = await _run_case(case_catalog, case_tools, n2id, case, **plan_kw)
        mark = "PASS" if r.passed else ("ERR " if r.error else "FAIL")
        print(f"[{mark}] {r.id:<28} {r.lang}  {r.latency_ms:>5}ms  {r.tokens:>5}tok")
        if not r.passed:
            if r.error:
                print(f"        error: {r.error}")
            for c in r.checks:
                if not c["ok"]:
                    print(f"        ✗ {c['check']}  ({c['detail']})")
        results.append(r)

    _summarize(results)
    _write(results, data, stamp)


def _summarize(results: list[CaseResult]) -> None:
    # Skipped cases leave the denominator, so the score stays comparable to a run
    # taken when a different set of tools was active.
    skipped = [r for r in results if r.skipped]
    scored = [r for r in results if not r.skipped]
    total = len(scored)
    passed = sum(r.passed for r in scored)
    print("\n" + "=" * 60)
    if total:
        print(f"OVERALL: {passed}/{total} passed ({passed / total:.0%})")
    else:
        print("no scored cases")
    if skipped:
        print(f"SKIPPED: {len(skipped)} of {len(results)} (tool not in the active registry)")
        for r in skipped:
            print(f"  - {r.id:<28} {r.skip_reason}")

    def rate(subset: list[CaseResult]) -> str:
        subset = [r for r in subset if not r.skipped]
        return f"{sum(r.passed for r in subset)}/{len(subset)}" if subset else "0/0"

    results = scored
    cats = sorted({r.category for r in results})
    print("\nby category:")
    for c in cats:
        print(f"  {c:<22} {rate([r for r in results if r.category == c])}")
    print("\nregression guard:")
    print(f"  carried_over (must stay green) {rate([r for r in results if r.carried_over])}")
    print(f"  new (target bugs)              {rate([r for r in results if not r.carried_over])}")
    lat = [r.latency_ms for r in results if r.latency_ms]
    if lat:
        print(f"\nlatency ms: median={statistics.median(lat):.0f} max={max(lat)}")
    print(f"tokens total: {sum(r.tokens for r in results)}")
    prompt_in = sum(r.input_tokens for r in results)
    cached = sum(r.cache_read for r in results)
    if prompt_in:
        print(f"prompt cache: {cached}/{prompt_in} input tokens served from cache "
              f"({cached / prompt_in:.0%})")


def _write(results: list[CaseResult], dataset: dict, stamp: str | None = None) -> None:
    RESULTS_DIR.mkdir(exist_ok=True)
    ts = stamp or datetime.now().strftime("%Y-%m-%d_%H%M%S")
    out = RESULTS_DIR / f"planner_result_{ts}.json"
    scored = [r for r in results if not r.skipped]
    payload = {
        "timestamp": ts,
        # `total` counts SCORED cases only, so `passed/total` is comparable across
        # runs taken with different tools active; `skipped` records what was withheld.
        "total": len(scored),
        "passed": sum(r.passed for r in scored),
        "skipped": len(results) - len(scored),
        "cases": [asdict(r) for r in results],
    }
    out.write_text(json.dumps(payload, indent=2, ensure_ascii=False), encoding="utf-8")
    print(f"\nwrote {out}")


# --------------------------------------------------------------------------- #
# Selfcheck — verifies the scorer without an LLM call
# --------------------------------------------------------------------------- #

def _rescore(path: Path) -> None:
    """Re-score a saved run's persisted `facts` with the CURRENT assertions —
    iterate on assertions/dataset without spending another LLM run."""
    saved = json.loads(path.read_text(encoding="utf-8"))
    expect_by_id = {c["id"]: c["expect"] for c in json.loads(DATASET.read_text(encoding="utf-8"))["cases"]}
    n2id = _all_name_to_id()
    results: list[CaseResult] = []
    for c in saved["cases"]:
        facts, exp = c.get("facts") or {}, expect_by_id.get(c["id"], {})
        if c.get("skipped"):
            # Carried through as-is: there is no plan to re-score, and counting it
            # as a failure would defeat the point of having skipped it.
            r = CaseResult(
                id=c["id"], category=c["category"], lang=c["lang"],
                carried_over=c["carried_over"], question=c["question"], passed=False,
                skipped=True, skip_reason=c.get("skip_reason"),
            )
            print(f"[SKIP] {r.id:<28} {r.lang}  {r.skip_reason}")
            results.append(r)
            continue
        if c.get("error") or not facts:
            checks, passed = [], False
        else:
            t = evaluate_facts(facts, exp, n2id)
            checks = [{"check": ck, "ok": ok, "detail": d} for ck, ok, d in t]
            passed = all(x["ok"] for x in checks)
        r = CaseResult(
            id=c["id"], category=c["category"], lang=c["lang"], carried_over=c["carried_over"],
            question=c["question"], passed=passed, checks=checks, facts=facts, error=c.get("error"),
        )
        mark = "PASS" if r.passed else ("ERR " if r.error else "FAIL")
        print(f"[{mark}] {r.id:<28} {r.lang}")
        if not r.passed:
            if r.error:
                print(f"        error: {r.error}")
            for x in r.checks:
                if not x["ok"]:
                    print(f"        ✗ {x['check']}  ({x['detail']})")
        results.append(r)
    _summarize(results)
    print(f"\n(re-scored {path.name} — no LLM calls)")


def _selfcheck() -> None:
    from types import SimpleNamespace as NS
    n2id = name_to_id()

    def plan(tasks_tcs: list[list[tuple[str, dict]]], infeasible: str | None = None):
        tasks = [NS(tool_calls=[NS(tool=t, args=a) for t, a in tcs]) for tcs in tasks_tcs]
        return NS(tasks=tasks, infeasible_reason=infeasible)

    def ev(tl: Any, expect: dict) -> bool:
        return all(ok for _, ok, _ in evaluate_facts(extract_facts(tl), expect, n2id))

    exp_rank = {"group_by": True, "group_by_col": "Equipment_Number", "select_agg": "avg", "order_dir": "asc", "limit": 5}

    # ranking via IR group_by — passes
    good_ir = plan([[("retrieve_data", {"ir": {
        "select": [{"kind": "agg", "fn": "avg", "column_id": "c_pa_percent"}],
        "group_by": ["c_equipment_number"],
        "order_by": [{"column_id": "avg_pa", "dir": "asc"}], "limit": 5}})]])
    assert ev(good_ir, exp_rank), "IR-group ranking should pass"

    # grouping via analyze_aggregate tool (aliases + 'mean') must ALSO count
    agg_tool = plan([
        [("retrieve_data", {"ir": {"select": [
            {"kind": "column", "column_id": "c_section", "alias": "section"},
            {"kind": "column", "column_id": "c_pa_percent", "alias": "pa"}]}})],
        [("analyze_aggregate", {"group_by": ["section"], "aggregations": {"pa": ["mean"]}})],
    ])
    assert ev(agg_tool, {"group_by": True, "group_by_col": "Section", "select_agg": "avg"}), \
        "analyze_aggregate grouping (mean==avg) should pass"

    # buggy: raw rows, no grouping anywhere — must FAIL
    bad = plan([[("retrieve_data", {"ir": {"select": [{"kind": "column", "column_id": "c_equipment_number"}],
                 "order_by": [{"column_id": "c_pa_percent", "dir": "asc"}], "limit": 5}})]])
    assert not ev(bad, exp_rank), "raw-row ranking should fail"

    # fuzzy: enumerated 'in' fails; non-enumerated (like/=) passes
    in_ir = plan([[("retrieve_data", {"ir": {"select": [{"kind": "agg", "fn": "count"}],
                   "filters": [{"column_id": "c_model_unit", "op": "in", "value": ["777E", "777D"]}]}})]])
    assert not ev(in_ir, {"no_filter_op": "in"}), "enumerated 'in' should fail"
    ok_ir = plan([[("retrieve_data", {"ir": {"select": [{"kind": "agg", "fn": "count"}],
                   "filters": [{"column_id": "c_model_unit", "op": "like", "value": "777%"}]}})]])
    assert ev(ok_ir, {"no_filter_op": "in"}), "non-enumerated filter should pass"

    assert ev(NS(tasks=[], infeasible_reason="no churn data"), {"infeasible": True})

    # column disambiguation: selecting Plan_PA_Percent aliased "pa_percent" must FAIL
    disambig = {"selects_col": "PA_Percent", "not_selects_col": "Plan_PA_Percent"}
    wrong_col = plan([[("retrieve_data", {"ir": {"select": [
        {"kind": "column", "column_id": "c_plan_pa_percent", "alias": "pa_percent"}]}})]])
    assert not ev(wrong_col, disambig), "wrong column (Plan_PA_Percent aliased pa_percent) should fail"
    right_col = plan([[("retrieve_data", {"ir": {"select": [
        {"kind": "column", "column_id": "c_pa_percent", "alias": "pa_percent"}]}})]])
    assert ev(right_col, disambig), "right column (PA_Percent) should pass"

    # trend chart must be aggregated (group_by in IR OR an analyze_* step), not raw
    chart_exp = {"has_tool": "render_chart", "chart_aggregated": True}
    raw_chart = plan([
        [("retrieve_data", {"ir": {"select": [{"kind": "column", "column_id": "c_pa_percent"}]}})],
        [("render_chart", {})]])
    assert not ev(raw_chart, chart_exp), "raw retrieve -> chart (no aggregation) should fail"
    # aggregate in the retrieve IR (group_by date) -> chart : valid, no analyze_* needed
    grp_chart = plan([
        [("retrieve_data", {"ir": {"select": [{"kind": "agg", "fn": "avg", "column_id": "c_pa_percent"}],
                                    "group_by": ["c_from_date"]}})],
        [("render_chart", {})]])
    assert ev(grp_chart, chart_exp), "retrieve(group_by date) -> chart should pass"
    # or analyze_trend -> chart : also valid
    trend_chart = plan([
        [("retrieve_data", {"ir": {"select": [{"kind": "column", "column_id": "c_pa_percent"}]}})],
        [("analyze_trend", {})], [("render_chart", {})]])
    assert ev(trend_chart, chart_exp), "retrieve -> trend -> chart should pass"

    # name-based EMA checks (2026-10-22): a site filtered by its code through the join
    site_ok = plan([[("retrieve_data", {"ir": {"table_id": "t_ema_lag",
        "joins": [{"target_table_id": "t_ema_site"}],
        "select": [{"kind": "agg", "fn": "avg", "column_id": "c_ema_lag_pa"}],
        "filters": [{"column_id": "c_ema_site_site_code", "op": "=", "value": "SLT"}]}})]])
    site_exp = {"filter_col": "site_code", "no_literal_filter_col": "site_id", "reads_table": "ema_site"}
    assert ev(site_ok, site_exp), "site filtered by code through the join should pass"
    site_guess = plan([[("retrieve_data", {"ir": {"table_id": "t_ema_lag",
        "select": [{"kind": "agg", "fn": "avg", "column_id": "c_ema_lag_pa"}],
        "filters": [{"column_id": "c_ema_lag_site_id", "op": "=", "value": 3002}]}})]])
    assert not ev(site_guess, site_exp), "a typed site id should fail"
    assert ev(site_ok, {"no_limit": 1, "ir_group_by": False})
    assert not ev(plan([[("retrieve_data", {"ir": {"limit": 1, "select": []}})]]), {"no_limit": 1})

    print("selfcheck OK — scorer distinguishes good vs buggy plans (agg paths + column + chart)")


if __name__ == "__main__":
    asyncio.run(main())