"""Dryrun: replay sample frames through Detector + StateMachine; compute metrics."""
from __future__ import annotations

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

from .config import Config
from .state_machine import StateMachine


@dataclass
class ConfusionMatrix:
    counts: dict[str, dict[str, int]] = field(default_factory=dict)

    def add(self, label: str, predicted: str) -> None:
        if label not in self.counts:
            self.counts[label] = {}
        self.counts[label][predicted] = self.counts[label].get(predicted, 0) + 1

    def per_label(self) -> dict[str, dict[str, float]]:
        # Column sums: total times each class was predicted (across all true labels)
        col_sums: dict[str, int] = {}
        for preds in self.counts.values():
            for pred, cnt in preds.items():
                col_sums[pred] = col_sums.get(pred, 0) + cnt

        result: dict[str, dict[str, float]] = {}
        for label, preds in self.counts.items():
            tp = preds.get(label, 0)
            support = sum(preds.values())
            total_predicted_as_label = col_sums.get(label, 0)
            fp = total_predicted_as_label - tp
            precision = tp / (tp + fp) if (tp + fp) > 0 else 0.0
            recall = tp / support if support > 0 else 0.0
            f1 = (
                2.0 * precision * recall / (precision + recall)
                if (precision + recall) > 0.0
                else 0.0
            )
            result[label] = {
                "precision": precision,
                "recall": recall,
                "f1": f1,
                "support": float(support),
            }
        return result

    def overall_accuracy(self) -> float:
        tp_total = sum(preds.get(lbl, 0) for lbl, preds in self.counts.items())
        total = sum(sum(preds.values()) for preds in self.counts.values())
        return tp_total / total if total > 0 else 0.0


@dataclass
class DryrunResult:
    n_samples: int
    n_labeled: int
    confusion: ConfusionMatrix
    fire_events: list[dict]
    precision_overall: float
    recall_overall: float
    acceptance_pass: bool


def dryrun(
    samples_dir: Path,
    labels_path: Path,
    cfg: Config,
    frame_loader: Callable | None = None,
) -> DryrunResult:
    from .detector import Detector

    _loader: Callable
    if frame_loader is None:
        import cv2 as _cv2

        def _default_loader(p: Path):  # type: ignore[misc]
            return _cv2.imread(str(p))

        _loader = _default_loader
    else:
        _loader = frame_loader

    labels: dict[str, str] = json.loads(labels_path.read_text(encoding="utf-8"))

    # Use "background" ColorSpec for bg detection if provided; else sensible default
    bg_spec = cfg.colors.get("background")
    bg_rgb: tuple[int, int, int] = bg_spec.rgb if bg_spec else (18, 18, 18)
    bg_tol: float = float(bg_spec.tolerance) if bg_spec else 15.0

    samples = sorted(samples_dir.glob("*.png"))

    current_frame: list = [None]

    def capture():
        return current_frame[0]

    detector = Detector(cfg, capture, bg_rgb=bg_rgb, bg_tol=bg_tol)
    sm = StateMachine(lockout_s=cfg.lockout_s)
    cm = ConfusionMatrix()
    fire_events: list[dict] = []
    n_labeled = 0

    for i, png_path in enumerate(samples):
        stem = png_path.stem
        label = labels.get(stem)
        if label is None:
            continue  # unlabeled → skip

        n_labeled += 1
        current_frame[0] = _loader(png_path)
        result = detector.step(ts=float(i) * 5.0)
        predicted: str = result.color if result.color is not None else "UNKNOWN"
        cm.add(label, predicted)

        if predicted != "UNKNOWN":
            t = sm.feed(predicted, float(i) * 5.0)  # type: ignore[arg-type]
            if t.trigger is not None and not t.locked:
                fire_events.append(
                    {"ts": float(i) * 5.0, "direction": t.trigger, "sample": stem}
                )

    per = cm.per_label()
    total_support = sum(v["support"] for v in per.values())
    if total_support > 0:
        precision_overall = (
            sum(v["precision"] * v["support"] for v in per.values()) / total_support
        )
        recall_overall = (
            sum(v["recall"] * v["support"] for v in per.values()) / total_support
        )
    else:
        precision_overall = 0.0
        recall_overall = 0.0

    acceptance_pass = precision_overall >= 1.0 and recall_overall >= 0.95

    return DryrunResult(
        n_samples=len(samples),
        n_labeled=n_labeled,
        confusion=cm,
        fire_events=fire_events,
        precision_overall=precision_overall,
        recall_overall=recall_overall,
        acceptance_pass=acceptance_pass,
    )


def print_report(result: DryrunResult) -> None:
    print("\n=== Dryrun Report ===")
    print(f"Samples: {result.n_samples}  Labeled: {result.n_labeled}")
    print(
        f"Precision: {result.precision_overall:.3f}  "
        f"Recall: {result.recall_overall:.3f}  "
        f"Gate: {'PASS' if result.acceptance_pass else 'FAIL'}"
    )
    print(f"Fire events: {len(result.fire_events)}")
    for ev in result.fire_events:
        print(f"  {ev}")
    per = result.confusion.per_label()
    print(f"\n{'Label':<15} {'Prec':>6} {'Rec':>6} {'F1':>6} {'Sup':>5}")
    for lbl, m in sorted(per.items()):
        print(
            f"{lbl:<15} {m['precision']:>6.3f} {m['recall']:>6.3f} "
            f"{m['f1']:>6.3f} {int(m['support']):>5}"
        )
