"""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 json
import statistics
import time
from dataclasses import asdict, dataclass, field
from datetime import datetime
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.inputs import Constraints
from src.agents.planner.registry import default_registry
from src.agents.planner.service import PlannerService

from .catalog_fixture import build_pa_catalog, name_to_id

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

_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.",
)


# --------------------------------------------------------------------------- #
# 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 == "infeasible":
            res.append(("infeasible" if want else "feasible", f["infeasible"] == want, f"infeasible={f['infeasible']}"))
        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

    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)


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

@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


async def _run_case(planner: PlannerService, catalog: Any, tools: Any, n2id: dict, case: dict) -> CaseResult:
    tok = _TokenCounter()
    started = time.perf_counter()
    facts: dict = {}
    try:
        task_list = await planner.plan(
            _CONTEXT, catalog, tools, case["question"], Constraints(), callbacks=[tok]
        )
        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,
        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("--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

    data = json.loads(DATASET.read_text(encoding="utf-8"))
    cases = data["cases"][: args.limit] if args.limit else data["cases"]
    catalog, tools, n2id = build_pa_catalog(), default_registry(), name_to_id()
    planner = PlannerService()

    results: list[CaseResult] = []
    for case in cases:
        r = await _run_case(planner, catalog, tools, n2id, case)
        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)


def _summarize(results: list[CaseResult]) -> None:
    total = len(results)
    passed = sum(r.passed for r in results)
    print("\n" + "=" * 60)
    print(f"OVERALL: {passed}/{total} passed ({passed / total:.0%})" if total else "no cases")

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

    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)}")


def _write(results: list[CaseResult], dataset: dict) -> None:
    RESULTS_DIR.mkdir(exist_ok=True)
    ts = datetime.now().strftime("%Y-%m-%d_%H%M%S")
    out = RESULTS_DIR / f"planner_result_{ts}.json"
    payload = {
        "timestamp": ts,
        "total": len(results),
        "passed": sum(r.passed for r in results),
        "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 = 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("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"

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


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