#!/usr/bin/env python3
"""Summarize a prime-rl eval traces.jsonl (flat or nested Harbor traces)."""
from __future__ import annotations

import json
import sys
from collections import Counter
from pathlib import Path


def _iter_traces(path: Path):
    with path.open() as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            try:
                row = json.loads(line)
            except json.JSONDecodeError:
                continue
            traces = row.get("traces")
            if isinstance(traces, list) and traces:
                for t in traces:
                    if isinstance(t, dict):
                        yield t
            else:
                yield row


def _name(t: dict) -> str:
    data = (t.get("task") or {}).get("data") or {}
    return str(data.get("name") or t.get("task_id") or t.get("id") or "?")


def _score(t: dict) -> float | None:
    rewards = t.get("rewards")
    if isinstance(rewards, dict):
        solved = rewards.get("solved")
        if isinstance(solved, dict) and "score" in solved:
            try:
                return float(solved["score"])
            except (TypeError, ValueError):
                pass
        if isinstance(solved, (int, float)):
            return float(solved)
    r = t.get("reward")
    if r is None:
        r = (t.get("metrics") or {}).get("reward")
    if r is None:
        return None
    try:
        return float(r)
    except (TypeError, ValueError):
        return None


def _err(t: dict) -> str | None:
    errors = t.get("errors") or t.get("error") or t.get("err")
    if isinstance(errors, list) and errors:
        return str(errors[0])[:120]
    if errors:
        return str(errors)[:120]
    return None


def main() -> None:
    path = Path(sys.argv[1] if len(sys.argv) > 1 else "traces.jsonl")
    if path.is_dir():
        cands = list(path.rglob("traces.jsonl"))
        if not cands:
            print("no traces.jsonl under", path)
            sys.exit(1)
        path = cands[0]
    n = n_ok = n_err = n_scored = 0
    stops: Counter[str] = Counter()
    errs: Counter[str] = Counter()
    solved: list[str] = []
    for t in _iter_traces(path):
        n += 1
        err = _err(t)
        if err:
            n_err += 1
            errs[err[:80]] += 1
        sc = _score(t)
        if sc is not None:
            n_scored += 1
            if sc > 0.5:
                n_ok += 1
                solved.append(_name(t))
        stops[str(t.get("stop_condition") or "?")] += 1
        mark = "OK" if (sc is not None and sc > 0.5) else ("ERR" if err else "no")
        print(f"  {mark:3s} {_name(t):52s} stop={t.get('stop_condition')} score={sc}")
    print(f"file {path}")
    print(f"rows {n} scored {n_scored} errors {n_err} solved {n_ok} {solved}")
    if n_scored:
        print(f"pass_rate {n_ok / n_scored:.4f}")
    print("stops", dict(stops))
    if errs:
        print("errors:")
        for k, v in errs.most_common(8):
            print(f"  {v:4d} {k}")


if __name__ == "__main__":
    main()