from __future__ import annotations

from dataclasses import dataclass
from typing import Any

NORMALIZATION_VERSION = "2026-03-v1"

_TEST_TO_SUITE: dict[str, str] = {
    "TAU-RETAIL": "tau-bench",
    "TAU-AIRLINE": "tau-bench",
    "SWE-LITE": "swe-bench",
    "SWE-VERIFIED": "swe-bench",
    "GAIA2-MINI": "gaia2",
    "GAIA2-EXEC": "gaia2",
    "CLAW-GDPVAL": "clawwork",
    "CLAW-SURVIVAL": "clawwork",
}


@dataclass(frozen=True)
class NormalizedBenchmarkResult:
    test_id: str
    suite: str
    metric_name: str
    raw_score: float
    normalized_score: float
    passed: bool
    mode: str
    source_standard: str
    notes: str = ""

    def to_dict(self) -> dict[str, Any]:
        return {
            "test_id": self.test_id,
            "suite": self.suite,
            "metric_name": self.metric_name,
            "raw_score": round(self.raw_score, 6),
            "normalized_score": round(self.normalized_score, 6),
            "passed": self.passed,
            "mode": self.mode,
            "source_standard": self.source_standard,
            "notes": self.notes,
            "normalization_version": NORMALIZATION_VERSION,
        }


def suite_for_test_id(test_id: str) -> str | None:
    return _TEST_TO_SUITE.get(test_id)


def normalize_external_trace(
    *,
    test_id: str,
    trace: dict[str, Any] | None,
    fallback_passed: bool,
) -> NormalizedBenchmarkResult | None:
    suite = suite_for_test_id(test_id)
    if suite is None:
        return None

    data = trace or {}
    mode = str(data.get("mode", "unknown"))

    if suite == "tau-bench":
        raw = _safe_float(data.get("avg_reward", 0.0))
        return NormalizedBenchmarkResult(
            test_id=test_id,
            suite=suite,
            metric_name="avg_reward",
            raw_score=raw,
            normalized_score=_clamp(raw),
            passed=bool(data.get("passed", fallback_passed)),
            mode=mode,
            source_standard="tau-bench reward",
            notes="Reward already in [0,1].",
        )

    if suite == "swe-bench":
        raw = _safe_float(data.get("resolve_rate", 0.0))
        return NormalizedBenchmarkResult(
            test_id=test_id,
            suite=suite,
            metric_name="resolve_rate",
            raw_score=raw,
            normalized_score=_clamp(raw),
            passed=bool(data.get("passed", fallback_passed)),
            mode=mode,
            source_standard="SWE-bench resolve rate",
            notes="Resolve rate already in [0,1].",
        )

    if suite == "gaia2":
        raw = _safe_float(data.get("completion_rate", 0.0))
        return NormalizedBenchmarkResult(
            test_id=test_id,
            suite=suite,
            metric_name="completion_rate",
            raw_score=raw,
            normalized_score=_clamp(raw),
            passed=bool(data.get("passed", fallback_passed)),
            mode=mode,
            source_standard="GAIA2 completion rate",
            notes="Completion rate already in [0,1].",
        )

    if suite == "clawwork":
        if test_id == "CLAW-SURVIVAL":
            raw = 1.0 if bool(data.get("survived", data.get("passed", fallback_passed))) else 0.0
            metric = "survival"
            notes = "Binary survival metric."
        else:
            raw = _safe_float(data.get("avg_quality", 0.0))
            metric = "avg_quality"
            notes = "GDPVal quality score in [0,1]."
        return NormalizedBenchmarkResult(
            test_id=test_id,
            suite=suite,
            metric_name=metric,
            raw_score=raw,
            normalized_score=_clamp(raw),
            passed=bool(data.get("passed", fallback_passed)),
            mode=mode,
            source_standard="ClawWork GDPVal",
            notes=notes,
        )

    return None


def aggregate_suite_results(
    normalized_results: list[NormalizedBenchmarkResult],
) -> dict[str, dict[str, Any]]:
    by_suite: dict[str, list[NormalizedBenchmarkResult]] = {}
    for row in normalized_results:
        by_suite.setdefault(row.suite, []).append(row)

    scorecard: dict[str, dict[str, Any]] = {}
    for suite, rows in by_suite.items():
        avg_score = sum(r.normalized_score for r in rows) / max(len(rows), 1)
        passed_all = all(r.passed for r in rows)
        mode = _collapse_mode([r.mode for r in rows])
        confidence = _confidence_for_mode(mode)
        scorecard[suite] = {
            "suite": suite,
            "normalized_score": round(avg_score, 4),
            "passed": passed_all,
            "confidence": confidence,
            "mode": mode,
            "normalization_version": NORMALIZATION_VERSION,
            "tests": [r.to_dict() for r in rows],
        }
    return scorecard


def score_from_scorecard(scorecard: dict[str, dict[str, Any]]) -> float:
    if not scorecard:
        return 0.0
    values = [_safe_float(item.get("normalized_score", 0.0)) for item in scorecard.values()]
    return sum(values) / max(len(values), 1)


def _safe_float(value: Any) -> float:
    try:
        return float(value)
    except (TypeError, ValueError):
        return 0.0


def _clamp(v: float, lo: float = 0.0, hi: float = 1.0) -> float:
    return max(lo, min(hi, v))


def _collapse_mode(modes: list[str]) -> str:
    if any(m == "local-cache-full" for m in modes):
        return "local-cache-full"
    if any(m == "offline-miniset" for m in modes):
        return "offline-miniset"
    if any(m == "online-full" for m in modes):
        return "online-full"
    if any(m == "dataset-only" for m in modes):
        return "dataset-only"
    if any(m == "are-cli" for m in modes):
        return "are-cli"
    return "unknown"


def _confidence_for_mode(mode: str) -> float:
    if mode == "online-full":
        return 1.0
    if mode == "local-cache-full":
        return 0.95
    if mode in {"dataset-only", "are-cli"}:
        return 0.85
    if mode == "offline-miniset":
        return 0.7
    return 0.6
