from __future__ import annotations

import json
import os
from collections import Counter
from dataclasses import dataclass, field

from .schema import Benchmark, BenchmarkItem, SourceRef


def load_benchmark(path: str) -> Benchmark:
    with open(path, "r", encoding="utf-8") as f:
        data = json.load(f)
    return Benchmark.model_validate(data)


def save_benchmark(bench: Benchmark, path: str) -> None:
    os.makedirs(os.path.dirname(path) or ".", exist_ok=True)
    with open(path, "w", encoding="utf-8") as f:
        json.dump(bench.model_dump(), f, ensure_ascii=False, indent=2)


# --------------------------------------------------------------------------- #
# Source-ref validity -- the core guarantee of the benchmark.
# --------------------------------------------------------------------------- #

def ref_is_valid(ref: SourceRef, *, runs_root: str, durations: dict[str, float]) -> tuple[bool, str]:
    if ref.type == "video":
        if not ref.clip_id:
            return False, "video ref missing clip_id"
        if ref.start is None or ref.end is None:
            return False, "video ref missing start/end"
        if ref.end < ref.start:
            return False, "video ref end < start"
        dur = durations.get(ref.clip_id)
        if dur is not None and ref.start > dur + 1.0:
            return False, f"video ref start {ref.start:.1f}s > duration {dur:.1f}s"
        return True, ""
    if ref.type == "frame":
        if not ref.frame:
            return False, "frame ref missing path"
        if not os.path.isfile(os.path.join(runs_root, ref.frame)):
            return False, f"frame file not found: {ref.frame}"
        return True, ""
    if ref.type == "manual_page":
        if not ref.page or ref.page < 1:
            return False, "manual_page ref missing/invalid page"
        if ref.char_start is not None or ref.char_end is not None:
            if ref.char_start is None or ref.char_end is None or ref.char_start < 0 or ref.char_end <= ref.char_start:
                return False, "manual_page ref has malformed char_start/char_end"
        return True, ""
    if ref.type == "dtc_seed":
        return (bool(ref.clip_id or ref.quote), "dtc_seed ref empty")
    return False, f"unknown ref type {ref.type}"


@dataclass
class ItemReport:
    id: str
    errors: list[str] = field(default_factory=list)
    warnings: list[str] = field(default_factory=list)
    n_valid_refs: int = 0

    @property
    def ok(self) -> bool:
        return not self.errors


def validate_item(item: BenchmarkItem, *, runs_root: str, durations: dict[str, float]) -> ItemReport:
    rep = ItemReport(id=item.id)

    if not item.query.strip():
        rep.errors.append("empty query")
    if not item.gold_steps:
        rep.errors.append("no gold_steps")
    if item.query_type in ("symptom", "diagnostic") and not (item.gold_cause or "").strip():
        rep.errors.append(f"{item.query_type} item has no gold_cause")

    # Every step needs >=1 valid ref (the triple-citation guarantee).
    for st in item.gold_steps:
        if not st.source_refs:
            rep.errors.append(f"step {st.step} has no source_refs")
            continue
        step_valid = 0
        for ref in st.source_refs:
            ok, why = ref_is_valid(ref, runs_root=runs_root, durations=durations)
            if ok:
                step_valid += 1
                rep.n_valid_refs += 1
            else:
                rep.warnings.append(f"step {st.step}: {why}")
        if step_valid == 0:
            rep.errors.append(f"step {st.step} has no VALID source_ref")

    # Item-level: success metric requires >=1 valid source ref per item.
    item_valid = any(
        ref_is_valid(r, runs_root=runs_root, durations=durations)[0] for r in item.gold_source_refs
    )
    if not item_valid and rep.n_valid_refs == 0:
        rep.errors.append("item has no valid source ref")

    # A frame and a video ref together = stronger (triple-citation-ready) grounding.
    has_frame = any(r.type == "frame" for st in item.gold_steps for r in st.source_refs)
    has_video = any(r.type == "video" for st in item.gold_steps for r in st.source_refs)
    if not (has_frame and has_video):
        rep.warnings.append("step grounding is not multi-modal (missing frame or video ref)")

    return rep


@dataclass
class BenchmarkReport:
    n_items: int
    n_ok: int
    n_verified: int
    by_type: dict
    by_difficulty: dict
    by_component: dict
    n_clips: int
    total_steps: int
    avg_steps: float
    avg_refs_per_step: float
    items_with_valid_ref: int
    item_reports: list[ItemReport]

    @property
    def all_valid(self) -> bool:
        return self.n_ok == self.n_items

    def summary_dict(self) -> dict:
        return {
            "n_items": self.n_items,
            "n_ok": self.n_ok,
            "n_with_errors": self.n_items - self.n_ok,
            "n_verified": self.n_verified,
            "items_with_valid_ref": self.items_with_valid_ref,
            "pct_items_with_valid_ref": round(100.0 * self.items_with_valid_ref / max(1, self.n_items), 1),
            "by_query_type": self.by_type,
            "by_difficulty": self.by_difficulty,
            "n_components": len(self.by_component),
            "n_source_clips": self.n_clips,
            "total_gold_steps": self.total_steps,
            "avg_steps_per_item": round(self.avg_steps, 2),
            "avg_refs_per_step": round(self.avg_refs_per_step, 2),
        }


def compute_durations(runs) -> dict[str, float]:
    return {rs.run_id: rs.duration for rs in runs}


def validate_benchmark(bench: Benchmark, *, runs_root: str, durations: dict[str, float]) -> BenchmarkReport:
    reports = [validate_item(it, runs_root=runs_root, durations=durations) for it in bench.items]
    n_ok = sum(1 for r in reports if r.ok)
    items_with_valid_ref = sum(1 for r in reports if r.n_valid_refs > 0)
    total_steps = sum(len(it.gold_steps) for it in bench.items)
    total_refs = sum(len(st.source_refs) for it in bench.items for st in it.gold_steps)
    return BenchmarkReport(
        n_items=len(bench.items),
        n_ok=n_ok,
        n_verified=sum(1 for it in bench.items if it.verified),
        by_type=dict(Counter(it.query_type for it in bench.items)),
        by_difficulty=dict(Counter(it.difficulty for it in bench.items)),
        by_component=dict(Counter((it.component or "?") for it in bench.items)),
        n_clips=len({it.clip_id for it in bench.items if it.clip_id}),
        total_steps=total_steps,
        avg_steps=total_steps / max(1, len(bench.items)),
        avg_refs_per_step=total_refs / max(1, total_steps),
        items_with_valid_ref=items_with_valid_ref,
        item_reports=reports,
    )


# --------------------------------------------------------------------------- #
# Scoring stubs for Phase P1/P2 (a system prediction vs. gold). Kept simple and
# dependency-free here; the LLM-as-judge faithfulness scorer lands in P2.
# --------------------------------------------------------------------------- #

def _norm(s: str) -> set[str]:
    return {w for w in "".join(c.lower() if c.isalnum() else " " for c in (s or "")).split() if len(w) > 2}


def cause_hit(gold_cause: str, predicted_causes: list[str], *, thresh: float = 0.34) -> bool:
    """Token-overlap proxy for Top-k cause hit-rate (replace with LLM judge in P1)."""
    g = _norm(gold_cause)
    if not g:
        return False
    for c in predicted_causes:
        p = _norm(c)
        if p and len(g & p) / len(g) >= thresh:
            return True
    return False


def citation_validity(pred_refs: list[dict], *, runs_root: str, durations: dict[str, float]) -> float:
    """Fraction of a system's cited refs that resolve to real evidence."""
    if not pred_refs:
        return 0.0
    ok = 0
    for r in pred_refs:
        try:
            ref = SourceRef.model_validate(r)
        except Exception:
            continue
        if ref_is_valid(ref, runs_root=runs_root, durations=durations)[0]:
            ok += 1
    return ok / len(pred_refs)