"""Persisted "repair library" — the model's accumulated knowledge of
known breakage -> repair pairs. Curated from successful rollouts during
training. Loaded at inference time as a few-shot prefix when the agent
recognises a familiar error class.
"""
from __future__ import annotations

import json
from dataclasses import asdict, dataclass, field
from pathlib import Path
from typing import Any, Optional


@dataclass
class RepairExample:
    primitive_type: str
    breakage_params: dict[str, Any]
    error_signature: str
    repair_diff: str
    visible_reward: float
    held_out: dict[str, float]
    task_id: str = ""

    def signature_key(self) -> str:
        return f"{self.primitive_type}::{self.error_signature[:80]}"


@dataclass
class RepairLibrary:
    examples: list[RepairExample] = field(default_factory=list)

    def add(self, example: RepairExample) -> None:
        self.examples.append(example)

    def best_match(self, primitive_type: str, error_text: str) -> Optional[RepairExample]:
        """Return the highest-reward example whose primitive_type matches and
        whose error text overlaps."""
        candidates = [
            e for e in self.examples if e.primitive_type == primitive_type
        ]
        if not candidates:
            return None
        scored = sorted(
            candidates,
            key=lambda e: (
                _ngram_overlap(e.error_signature, error_text),
                e.visible_reward,
            ),
            reverse=True,
        )
        return scored[0] if scored else None

    def to_dict(self) -> dict:
        return {
            "version": "1",
            "examples": [asdict(e) for e in self.examples],
            "size": len(self.examples),
            "by_primitive": _count_by_primitive(self.examples),
        }

    def save(self, path: str | Path) -> None:
        path = Path(path)
        path.parent.mkdir(parents=True, exist_ok=True)
        path.write_text(json.dumps(self.to_dict(), indent=2), encoding="utf-8")

    @classmethod
    def load(cls, path: str | Path) -> "RepairLibrary":
        data = json.loads(Path(path).read_text(encoding="utf-8"))
        examples = [RepairExample(**e) for e in data.get("examples", [])]
        return cls(examples=examples)


def _ngram_overlap(a: str, b: str, n: int = 3) -> float:
    if not a or not b:
        return 0.0

    def grams(text: str) -> set[str]:
        text = text.lower()
        return {text[i : i + n] for i in range(len(text) - n + 1)}

    ga, gb = grams(a), grams(b)
    if not ga or not gb:
        return 0.0
    return len(ga & gb) / max(1, len(ga | gb))


def _count_by_primitive(examples: list[RepairExample]) -> dict[str, int]:
    counts: dict[str, int] = {}
    for e in examples:
        counts[e.primitive_type] = counts.get(e.primitive_type, 0) + 1
    return counts


def curate_from_rollouts(
    rollout_results: list,
    min_reward: float = 0.6,
    min_held_out_clean: float = 0.5,
) -> RepairLibrary:
    """Build a RepairLibrary from a list of rollout dicts/RolloutResults."""
    lib = RepairLibrary()
    for r in rollout_results:
        get = r.get if isinstance(r, dict) else lambda k, default=None: getattr(r, k, default)
        if float(get("visible_reward", 0.0) or 0.0) < min_reward:
            continue
        if float(get("held_out_breakdown", {}).get("executed_cleanly", 0.0)) < min_held_out_clean:
            continue
        lib.add(
            RepairExample(
                primitive_type=str(get("primitive_type", "unknown")),
                breakage_params=dict(get("info", {}).get("breakage_spec", {}).get("params", {}))
                if isinstance(get("info", {}), dict)
                else {},
                error_signature=str(get("error_trace", "") or "")[:160],
                repair_diff=str(get("repair_completion", "") or get("info", {}).get("repair_diff", ""))[:2000],
                visible_reward=float(get("visible_reward", 0.0) or 0.0),
                held_out=dict(get("held_out_breakdown", {}) or {}),
                task_id=str(get("task_id", "")),
            )
        )
    return lib