"""Score a predictor (anything with `.ask(question, history)`) on generated eval examples."""

from __future__ import annotations

import json
from collections import defaultdict
from pathlib import Path

from zenith.engines.base import Result
from zenith.evals.metrics import caveats_recalled, faithful, params_correct, routed, sim_op_correct
from zenith.sft.generate import Example

SHOWCASE = Path(__file__).with_name("showcase.txt")
TARGETS = {"sim_op_accuracy": 0.90, "param_accuracy": 0.75, "faithfulness": 0.95, "caveat_recall": 0.70, "routing": 0.85}


def _mean(xs: list[bool]) -> float | None:
    return sum(xs) / len(xs) if xs else None


def score(example: Example, response) -> dict[str, bool]:
    """Metric -> pass/fail for one example; metrics that don't apply are absent."""
    gold = example.turns[-1]
    question = gold.question
    if example.kind in ("none", "unsupported"):
        return {"routing": routed(example.kind, response.spec)}
    out = {"sim_op_accuracy": sim_op_correct(gold.spec, response.spec),
           "param_accuracy": params_correct(gold.spec, response.spec)}
    if response.result is not None and response.answer is not None:
        result = Result(**response.result)
        out["faithfulness"] = faithful(response.answer, result, response.spec or {}, question)
        if result.notes:
            out["caveat_recall"] = caveats_recalled(response.answer, result)
    return out


def _group(example: Example) -> str:
    spec = example.turns[-1].spec
    return f"{spec['sim']}/{spec['op']}" if "op" in spec else example.kind


def run_eval(predictor, examples: list[Example], out_dir: Path, ppl: dict | None = None) -> dict:
    out_dir = Path(out_dir)
    out_dir.mkdir(parents=True, exist_ok=True)
    by_group: dict[str, dict[str, list[bool]]] = defaultdict(lambda: defaultdict(list))
    overall: dict[str, list[bool]] = defaultdict(list)
    for e in examples:
        response = predictor.ask(e.turns[-1].question, e.turns[:-1])
        for metric, ok in score(e, response).items():
            by_group[_group(e)][metric].append(ok)
            overall[metric].append(ok)
    report = {
        "examples": len(examples),
        "overall": {m: _mean(overall.get(m, [])) for m in TARGETS},
        "groups": {g: {m: _mean(v.get(m, [])) for m in TARGETS} for g, v in sorted(by_group.items())},
        "targets": TARGETS,
        "perplexity": ppl,
    }
    (out_dir / "report.json").write_text(json.dumps(report, indent=2))
    (out_dir / "report.md").write_text(_markdown(report))
    return report


def _cell(v: float | None) -> str:
    return "—" if v is None else f"{v:.1%}"


def _markdown(report: dict) -> str:
    metrics = list(TARGETS)
    lines = [f"# Evaluation ({report['examples']} examples)", "",
             "| group | " + " | ".join(metrics) + " |", "|---" * (len(metrics) + 1) + "|"]
    for g, row in report["groups"].items():
        lines.append(f"| {g} | " + " | ".join(_cell(row[m]) for m in metrics) + " |")
    lines.append("| **overall** | " + " | ".join(_cell(report["overall"][m]) for m in metrics) + " |")
    lines.append("| *target* | " + " | ".join(f"≥ {TARGETS[m]:.0%}" for m in metrics) + " |")
    if report.get("perplexity"):
        lines += ["", "Validation bits per byte: " + ", ".join(f"{k} {v:.3f}" for k, v in report["perplexity"].items())]
    return "\n".join(lines) + "\n"


def run_showcase(predictor, out_path: Path) -> None:
    sections = ["# Showcase transcripts", ""]
    for i, q in enumerate([l for l in SHOWCASE.read_text().splitlines() if l.strip()], 1):
        r = predictor.ask(q, [])
        sections += [f"### {i}. {q}", "", f"- **spec:** `{json.dumps(r.spec)}`"]
        if r.error:
            sections.append(f"- **error:** {r.error}")
        else:
            sections.append(f"- **answer:** {r.answer}" + (" *(truncated)*" if getattr(r, "truncated", False) else ""))
        sections.append("")
    Path(out_path).write_text("\n".join(sections))