from __future__ import annotations

import argparse
import dataclasses
import json
import os
from dataclasses import dataclass
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple

from .verifier import EvidenceItem, GroundednessConfig, analyze_groundedness


@dataclass(frozen=True)
class ExpectedClaim:
    selector: str  # substring to locate the claim in extracted claims
    supported: bool
    evidence_ids: Tuple[str, ...] = ()
    failure_type: str = ""  # optional, for reporting


@dataclass(frozen=True)
class BenchCase:
    case_id: str
    context: List[EvidenceItem]
    response: str
    expected: Tuple[ExpectedClaim, ...]


def _load_cases(path: Path) -> Tuple[Dict[str, Any], List[BenchCase]]:
    meta: Dict[str, Any] = {}
    cases: List[BenchCase] = []
    for line in path.read_text(encoding="utf-8").splitlines():
        line = line.strip()
        if not line:
            continue
        obj = json.loads(line)
        if obj.get("type") == "meta":
            meta = obj
            continue
        if obj.get("type") != "case":
            continue
        ctx = [EvidenceItem(evidence_id=c["id"], text=c["text"]) for c in obj.get("context", [])]
        expected = []
        for ec in obj.get("expected", {}).get("claims", []):
            expected.append(
                ExpectedClaim(
                    selector=str(ec.get("selector", "")).strip(),
                    supported=bool(ec.get("supported", False)),
                    evidence_ids=tuple(ec.get("evidence_ids", []) or []),
                    failure_type=str(ec.get("failure_type", "") or ""),
                )
            )
        cases.append(
            BenchCase(
                case_id=str(obj["id"]),
                context=ctx,
                response=str(obj.get("response", "")),
                expected=tuple(expected),
            )
        )
    if not meta:
        meta = {"suite": path.stem, "version": "unknown"}
    return meta, cases


def _find_assessment_for_selector(
    report: Dict[str, Any], selector: str
) -> Optional[Dict[str, Any]]:
    sel = (selector or "").strip().lower()
    if not sel:
        return None
    for a in report["assessments"]:
        t = str(a["claim"]["text"]).lower()
        if sel in t:
            return a
    return None


def _metrics_from_cases(
    cases: List[BenchCase], *, cfg: Optional[GroundednessConfig] = None
) -> Dict[str, Any]:
    tp = fp = fn = tn = 0
    extraction_hits = 0
    extraction_total = 0
    evidence_hit = 0
    evidence_total = 0

    per_case: List[Dict[str, Any]] = []
    for c in cases:
        rep = analyze_groundedness(c.response, c.context, config=cfg)
        rep_dict = {
            "case_id": c.case_id,
            "summary": dataclasses.asdict(rep.summary),
            "assessments": [
                {
                    "claim": {
                        "text": a.claim.text,
                        "cited": list(a.claim.cited_evidence_ids),
                        "type": a.claim.claim_type,
                    },
                    "supported": bool(a.supported),
                    "support_score": float(a.support_score),
                    "evidence_ids": list(a.evidence_ids),
                    "failure_types": [str(x.value) for x in a.failure_types],
                }
                for a in rep.assessments
            ],
        }

        mismatches: List[str] = []
        for ec in c.expected:
            extraction_total += 1
            a = _find_assessment_for_selector(rep_dict, ec.selector)
            if a is None:
                mismatches.append(f"missing_claim_selector:{ec.selector}")
                fn += 1 if ec.supported else 0
                # If it was expected unsupported, missing extraction is still a miss for extraction recall,
                # but does not affect supported/unsupported classification.
                continue
            extraction_hits += 1

            pred_supported = bool(a["supported"])
            true_supported = bool(ec.supported)
            if pred_supported and true_supported:
                tp += 1
            elif pred_supported and not true_supported:
                fp += 1
            elif (not pred_supported) and true_supported:
                fn += 1
            else:
                tn += 1

            if true_supported:
                evidence_total += 1
                pred_eids = set(a.get("evidence_ids", []) or [])
                if any(eid in pred_eids for eid in ec.evidence_ids):
                    evidence_hit += 1

        per_case.append({"case_id": c.case_id, "mismatches": mismatches})

    precision = tp / (tp + fp) if (tp + fp) else 0.0
    recall = tp / (tp + fn) if (tp + fn) else 0.0
    f1 = (2 * precision * recall) / (precision + recall) if (precision + recall) else 0.0
    extraction_recall = extraction_hits / extraction_total if extraction_total else 0.0
    evidence_hit_rate = evidence_hit / evidence_total if evidence_total else 0.0

    return {
        "counts": {"tp": tp, "fp": fp, "fn": fn, "tn": tn},
        "precision_supported": precision,
        "recall_supported": recall,
        "f1_supported": f1,
        "extraction_recall": extraction_recall,
        "evidence_hit_rate": evidence_hit_rate,
        "per_case": per_case,
    }


def main(argv: Optional[List[str]] = None) -> int:
    here = Path(__file__).resolve().parent
    ds_default = here.parent / "evals" / "datasets" / "groundedness_v1.jsonl"

    p = argparse.ArgumentParser(
        prog="entropic-groundedness-bench",
        description="Offline groundedness benchmark (experimental).",
    )
    p.add_argument(
        "--dataset", default=str(ds_default), help="Path to groundedness benchmark dataset (jsonl)."
    )
    p.add_argument(
        "--out", default=os.getenv("ENTROPIC_GROUNDEDNESS_OUT", "entropic_groundedness_runs")
    )
    p.add_argument("--min-support-score", type=float, default=None)
    args = p.parse_args(argv)

    ds_path = Path(args.dataset)
    meta, cases = _load_cases(ds_path)

    cfg = GroundednessConfig()
    if args.min_support_score is not None:
        cfg = dataclasses.replace(cfg, min_support_score=float(args.min_support_score))

    metrics = _metrics_from_cases(cases, cfg=cfg)

    run_id = datetime.now(timezone.utc).strftime("%Y%m%d_%H%M%S")
    outdir = Path(args.out) / str(meta.get("suite", "groundedness_v1")) / run_id
    outdir.mkdir(parents=True, exist_ok=True)

    payload = {
        "run_id": run_id,
        "created_utc": datetime.now(timezone.utc).isoformat(),
        "suite": meta,
        "dataset_path": str(ds_path),
        "dataset_sha256": _sha256_file(ds_path),
        "config": dataclasses.asdict(cfg),
        "metrics": metrics,
    }
    (outdir / "report.json").write_text(json.dumps(payload, indent=2), encoding="utf-8")
    print(str(outdir / "report.json"))
    return 0


def _sha256_file(path: Path) -> str:
    import hashlib

    h = hashlib.sha256()
    with path.open("rb") as f:
        for chunk in iter(lambda: f.read(65536), b""):
            h.update(chunk)
    return h.hexdigest()


if __name__ == "__main__":
    raise SystemExit(main())
