from __future__ import annotations

import argparse
import csv
import html
import json
from dataclasses import dataclass
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Dict, Iterable, List, Optional


def _utc_iso() -> str:
    return datetime.now(timezone.utc).isoformat()


def _read_run_json(path: Path) -> Optional[Dict[str, Any]]:
    try:
        return json.loads(path.read_text(encoding="utf-8"))
    except Exception:
        return None


def _iter_run_files(runs_root: Path) -> Iterable[Path]:
    # Expected: <runs>/<suite>/<provider>/<run_id>_<model>/run.json
    yield from runs_root.glob("**/run.json")


def _safe_get(d: Dict[str, Any], *keys: str) -> Any:
    cur: Any = d
    for k in keys:
        if not isinstance(cur, dict):
            return None
        cur = cur.get(k)
    return cur


@dataclass(frozen=True)
class Row:
    suite: str
    provider: str
    model: str
    entropic_core_version: str
    git_commit: str
    suite_sha256: str
    created_utc: str
    baseline_pass_rate: float
    entropic_pass_rate: float
    delta_pass_rate: float
    baseline_grounded_rate: float
    entropic_grounded_rate: float
    delta_grounded_rate: float
    baseline_contra_rate: float
    entropic_contra_rate: float
    delta_contra_rate: float
    baseline_refusal_rate: float
    entropic_refusal_rate: float
    delta_refusal_rate: float
    baseline_retries_avg: float
    entropic_retries_avg: float
    delta_retries_avg: float
    baseline_claim_support_rate: float
    entropic_claim_support_rate: float
    delta_claim_support_rate: float
    baseline_tool_correct_rate: float
    entropic_tool_correct_rate: float
    delta_tool_correct_rate: float
    baseline_tokens_total: int
    entropic_tokens_total: int
    delta_tokens_total: int
    baseline_latency_ms_avg: float
    entropic_latency_ms_avg: float
    delta_latency_ms_avg: float
    baseline_cost_usd_est: str
    entropic_cost_usd_est: str
    delta_cost_usd_est: str
    path: str

    def to_dict(self) -> Dict[str, Any]:
        return self.__dict__.copy()


def _fmt_cost(v: Any) -> str:
    if v is None:
        return ""
    try:
        return f"{float(v):.6f}"
    except Exception:
        return str(v)


def build_rows(runs_root: Path) -> List[Row]:
    rows: List[Row] = []
    for p in _iter_run_files(runs_root):
        payload = _read_run_json(p)
        if not isinstance(payload, dict):
            continue

        suite = str(_safe_get(payload, "suite", "suite") or payload.get("suite") or "")
        provider = str(payload.get("provider") or "")
        model = str(payload.get("model") or "")
        created = str(payload.get("created_utc") or "")
        suite_sha = str(payload.get("suite_sha256") or "")

        env = payload.get("env") if isinstance(payload.get("env"), dict) else {}
        ent_ver = str(env.get("entropic_core_version") or "")
        git = env.get("git") if isinstance(env.get("git"), dict) else {}
        git_commit = str(git.get("commit") or "")

        base = _safe_get(payload, "summary", "baseline") or {}
        ent = _safe_get(payload, "summary", "entropic") or {}

        def _f(x: Any) -> float:
            try:
                return float(x)
            except Exception:
                return 0.0

        def _i(x: Any) -> int:
            try:
                return int(x)
            except Exception:
                return 0

        base_pass = _f(base.get("pass_rate"))
        ent_pass = _f(ent.get("pass_rate"))
        base_gr = _f(base.get("grounded_rate"))
        ent_gr = _f(ent.get("grounded_rate"))
        base_cr = _f(base.get("contradiction_handling_rate"))
        ent_cr = _f(ent.get("contradiction_handling_rate"))
        base_ref = _f(base.get("refusal_rate"))
        ent_ref = _f(ent.get("refusal_rate"))
        base_ret = _f(base.get("retries_avg"))
        ent_ret = _f(ent.get("retries_avg"))
        base_claim = _f(base.get("groundedness_claim_supported_rate"))
        ent_claim = _f(ent.get("groundedness_claim_supported_rate"))
        base_tool = _f(base.get("tool_correct_rate"))
        ent_tool = _f(ent.get("tool_correct_rate"))

        bt = _safe_get(base, "tokens", "total_tokens")
        et = _safe_get(ent, "tokens", "total_tokens")
        bt_i = _i(bt)
        et_i = _i(et)

        base_lat = _f(base.get("latency_ms_avg"))
        ent_lat = _f(ent.get("latency_ms_avg"))

        base_cost = base.get("cost_usd_est")
        ent_cost = ent.get("cost_usd_est")

        # Deltas.
        row = Row(
            suite=suite,
            provider=provider,
            model=model,
            entropic_core_version=ent_ver,
            git_commit=git_commit[:12] if git_commit else "",
            suite_sha256=suite_sha[:12] if suite_sha else "",
            created_utc=created,
            baseline_pass_rate=base_pass,
            entropic_pass_rate=ent_pass,
            delta_pass_rate=ent_pass - base_pass,
            baseline_grounded_rate=base_gr,
            entropic_grounded_rate=ent_gr,
            delta_grounded_rate=ent_gr - base_gr,
            baseline_contra_rate=base_cr,
            entropic_contra_rate=ent_cr,
            delta_contra_rate=ent_cr - base_cr,
            baseline_refusal_rate=base_ref,
            entropic_refusal_rate=ent_ref,
            delta_refusal_rate=ent_ref - base_ref,
            baseline_retries_avg=base_ret,
            entropic_retries_avg=ent_ret,
            delta_retries_avg=ent_ret - base_ret,
            baseline_claim_support_rate=base_claim,
            entropic_claim_support_rate=ent_claim,
            delta_claim_support_rate=ent_claim - base_claim,
            baseline_tool_correct_rate=base_tool,
            entropic_tool_correct_rate=ent_tool,
            delta_tool_correct_rate=ent_tool - base_tool,
            baseline_tokens_total=bt_i,
            entropic_tokens_total=et_i,
            delta_tokens_total=et_i - bt_i,
            baseline_latency_ms_avg=base_lat,
            entropic_latency_ms_avg=ent_lat,
            delta_latency_ms_avg=ent_lat - base_lat,
            baseline_cost_usd_est=_fmt_cost(base_cost),
            entropic_cost_usd_est=_fmt_cost(ent_cost),
            delta_cost_usd_est=_fmt_cost(
                (_f(ent_cost) - _f(base_cost))
                if (ent_cost is not None and base_cost is not None)
                else None
            ),
            path=str(p.parent),
        )
        rows.append(row)

    # Newest first (best effort).
    rows.sort(key=lambda r: r.created_utc, reverse=True)
    return rows


def _write_json(path: Path, obj: Any) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(json.dumps(obj, indent=2), encoding="utf-8")


def _write_csv(path: Path, rows: List[Row]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("w", encoding="utf-8", newline="") as f:
        w = csv.DictWriter(f, fieldnames=list(rows[0].to_dict().keys()) if rows else [])
        if rows:
            w.writeheader()
            for r in rows:
                w.writerow(r.to_dict())


def _write_html(path: Path, rows: List[Row]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)

    def esc(s: str) -> str:
        return html.escape(str(s), quote=False)

    trs = []
    for r in rows:
        trs.append(
            "<tr>"
            + "".join(
                f"<td>{esc(v)}</td>"
                for v in [
                    r.created_utc,
                    r.suite,
                    r.provider,
                    r.model,
                    f"{r.entropic_pass_rate * 100:.1f}%",
                    f"{r.delta_pass_rate * 100:+.1f}%",
                    f"{r.entropic_grounded_rate * 100:.1f}%",
                    f"{r.delta_grounded_rate * 100:+.1f}%",
                    f"{r.entropic_contra_rate * 100:.1f}%",
                    f"{r.delta_contra_rate * 100:+.1f}%",
                    f"{r.entropic_refusal_rate * 100:.1f}%",
                    f"{r.delta_refusal_rate * 100:+.1f}%",
                    f"{r.entropic_retries_avg:.2f}",
                    f"{r.delta_retries_avg:+.2f}",
                    f"{r.entropic_claim_support_rate * 100:.1f}%",
                    f"{r.delta_claim_support_rate * 100:+.1f}%",
                    f"{r.entropic_tool_correct_rate * 100:.1f}%",
                    f"{r.delta_tool_correct_rate * 100:+.1f}%",
                    r.entropic_tokens_total,
                    r.delta_tokens_total,
                    f"{r.entropic_latency_ms_avg:.1f}",
                    f"{r.delta_latency_ms_avg:+.1f}",
                    r.entropic_cost_usd_est,
                    r.delta_cost_usd_est,
                    r.entropic_core_version,
                    r.git_commit,
                    r.suite_sha256,
                    r.path,
                ]
            )
            + "</tr>"
        )

    doc = f"""<!doctype html>
<html lang="en">
  <head>
    <meta charset="utf-8" />
    <meta name="viewport" content="width=device-width, initial-scale=1" />
    <title>Entropic Leaderboard</title>
    <style>
      :root {{
        --bg: #0b0e12;
        --text: #e8eef8;
        --muted: #a8b3c7;
        --line: rgba(232,238,248,0.12);
      }}
      html, body {{ background: var(--bg); color: var(--text); margin: 0; font-family: ui-sans-serif, system-ui, -apple-system, Segoe UI, sans-serif; }}
      .wrap {{ max-width: 1400px; margin: 0 auto; padding: 20px 16px 60px; }}
      h1 {{ margin: 0 0 8px; font-size: 22px; }}
      .sub {{ color: var(--muted); font-size: 13px; }}
      table {{ width: 100%; border-collapse: collapse; margin-top: 14px; font-size: 12px; }}
      th, td {{ border: 1px solid var(--line); padding: 6px 8px; vertical-align: top; }}
      th {{ color: var(--muted); text-align: left; position: sticky; top: 0; background: #0b0e12; }}
      tr:nth-child(even) {{ background: rgba(255,255,255,0.02); }}
      code {{ font-family: ui-monospace, Menlo, Consolas, monospace; }}
    </style>
  </head>
  <body>
    <div class="wrap">
      <h1>Entropic Leaderboard</h1>
      <div class="sub">generated: <code>{esc(_utc_iso())}</code> | runs: <code>{len(rows)}</code></div>
      <table>
        <thead>
          <tr>
            <th>created_utc</th>
            <th>suite</th>
            <th>provider</th>
            <th>model</th>
            <th>pass</th>
            <th>delta pass</th>
            <th>grounded</th>
            <th>delta grounded</th>
            <th>contra</th>
            <th>delta contra</th>
            <th>refusal</th>
            <th>delta refusal</th>
            <th>retries avg</th>
            <th>delta retries</th>
            <th>claim support</th>
            <th>delta claim</th>
            <th>tool correct</th>
            <th>delta tool</th>
            <th>tokens</th>
            <th>delta tokens</th>
            <th>lat ms avg</th>
            <th>delta lat</th>
            <th>cost usd</th>
            <th>delta cost</th>
            <th>entropic</th>
            <th>git</th>
            <th>suite sha</th>
            <th>path</th>
          </tr>
        </thead>
        <tbody>
          {''.join(trs)}
        </tbody>
      </table>
    </div>
  </body>
</html>
"""
    path.write_text(doc, encoding="utf-8")


def main(argv: Optional[List[str]] = None) -> int:
    p = argparse.ArgumentParser(
        prog="entropic-leaderboard", description="Aggregate entropic-evals runs into a leaderboard."
    )
    p.add_argument(
        "--runs",
        default="entropic_evals_runs",
        help="Root folder where entropic-evals writes runs.",
    )
    p.add_argument(
        "--out", default="evals/leaderboard", help="Output folder for leaderboard artifacts."
    )
    p.add_argument("--limit", type=int, default=0, help="0 means all rows.")
    args = p.parse_args(argv)

    runs_root = Path(args.runs)
    outdir = Path(args.out)
    rows = build_rows(runs_root)
    if args.limit and args.limit > 0:
        rows = rows[: args.limit]

    _write_json(
        outdir / "leaderboard.json",
        {"generated_utc": _utc_iso(), "rows": [r.to_dict() for r in rows]},
    )
    if rows:
        _write_csv(outdir / "leaderboard.csv", rows)
    _write_html(outdir / "leaderboard.html", rows)

    print(str(outdir / "leaderboard.html"))
    return 0


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