"""the ablation loop: run each task bare vs. with-zigman, summarize the lift."""

from __future__ import annotations

import json
from dataclasses import dataclass
from pathlib import Path
from typing import Callable

from .config import Config
from .pi import ArmResult, run_task
from .tasks import TASKS, Task


@dataclass
class Report:
    model: str
    arms: list[str]
    tasks: list[Task]
    results: dict[str, dict[str, ArmResult]]  # task.id -> arm -> result

    def total(self, arm: str) -> int:
        return sum(self.results[t.id][arm].oracle_ok for t in self.tasks)

    def lift(self) -> int | None:
        if "bare" in self.arms and "zigman" in self.arms:
            return self.total("zigman") - self.total("bare")
        return None

    def zigman_use_count(self) -> int:
        if "zigman" not in self.arms:
            return 0
        return sum(self.results[t.id]["zigman"].used_zigman for t in self.tasks)


def ablate(
    cfg: Config,
    task_filter: str = "",
    arms: tuple[str, ...] = ("bare", "zigman"),
    on_result: Callable[[ArmResult], None] | None = None,
) -> Report:
    tasks = [t for t in TASKS if not task_filter or task_filter in t.id]
    results: dict[str, dict[str, ArmResult]] = {}
    for t in tasks:
        results[t.id] = {}
        for arm in arms:
            r = run_task(cfg, t, with_zigman=(arm == "zigman"))
            results[t.id][arm] = r
            if on_result:
                on_result(r)
    return Report(model=cfg.model, arms=list(arms), tasks=tasks, results=results)


def render(report: Report) -> str:
    lines = [f"\nagent: Pi  model: {report.model}"]
    header = "task".ljust(16) + "".join(a.ljust(8) for a in report.arms)
    lines.append(header)
    for t in report.tasks:
        row = t.id.ljust(16)
        for a in report.arms:
            row += ("pass" if report.results[t.id][a].oracle_ok else "·").ljust(8)
        lines.append(row)
    lines.append("-" * len(header))
    lines.append("total".ljust(16) + "".join(f"{report.total(a)}/{len(report.tasks)}".ljust(8) for a in report.arms))
    if "zigman" in report.arms:
        lines.append(f"zigman used in {report.zigman_use_count()}/{len(report.tasks)} zigman-arm runs")
    if (lift := report.lift()) is not None:
        lines.append(f"zigman lift: {lift:+d}")
    return "\n".join(lines)


def result_line(r: ArmResult) -> str:
    to = " TIMEOUT" if r.timed_out else ""
    uz = (" zigman:" + (str(r.zigman_queries) if r.used_zigman else "none")) if r.arm == "zigman" else ""
    return (f"  {r.arm:<7} {r.task:<14} {'PASS' if r.oracle_ok else 'fail'} "
            f"[{','.join(r.tools) or 'no-tools'}{to}]{uz}")


def save(report: Report, out_dir: Path) -> Path:
    out_dir.mkdir(parents=True, exist_ok=True)
    payload = {
        "model": report.model,
        "arms": report.arms,
        "lift": report.lift(),
        "results": {
            tid: {arm: vars(r) | {"used_zigman": r.used_zigman} for arm, r in by_arm.items()}
            for tid, by_arm in report.results.items()
        },
    }
    out = out_dir / f"pi_{Path(report.model).name}.json"
    out.write_text(json.dumps(payload, indent=2))
    return out
