#!/usr/bin/env python3
"""Summarize e2e runs (scripts/e2e_fixture.sh / run_matrix.sh output).

Usage: scripts/summarize.py <out-dir> [--md summary.md] [--csv results.csv]

Every run directory holds result.json and metrics.jsonl. Per run we derive:
  steps         model calls in the turn
  llm_s         sum of stream durations (request start -> stream end)
  tool_s        sum of tool execution time
  overlap_s     tool time that ran while the model was still streaming
                (eager protocols; 0 for end-of-stream execution)
  confirm_ms    mean call open -> call confirmed (args complete)
  wait_ms       mean call confirmed -> execution start
  ctl_errors    contract errors (model left the CTL/XML grammar)
  tokens        prompt / completion / cached prompt tokens
"""

import csv
import json
import os
import statistics
import sys
from collections import defaultdict

FIELDS = [
    "run", "fixture", "protocol", "provider", "model", "outcome", "pass",
    "wall_s", "turn_s", "steps", "calls", "calls_failed", "llm_s", "tool_s",
    "overlap_s", "confirm_ms", "wait_ms", "ctl_errors", "retries",
    "prompt_tokens", "completion_tokens", "cached_tokens",
]


def load_jsonl(path):
    rows = []
    if not os.path.exists(path):
        return rows
    with open(path, encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            try:
                rows.append(json.loads(line))
            except json.JSONDecodeError:
                pass
    return rows


def analyze(run_dir):
    with open(os.path.join(run_dir, "result.json"), encoding="utf-8") as f:
        res = json.load(f)
    if res.get("outcome") in ("baseline_only", "unexpected_green_fixture"):
        return None
    m = load_jsonl(os.path.join(run_dir, "metrics.jsonl"))
    steps = [r for r in m if r.get("type") == "step"]
    calls = [r for r in m if r.get("type") == "call"]
    turns = [r for r in m if r.get("type") == "turn"]
    end_of = {(s["turn"], s["step"]): s.get("t_stream_end", 0) for s in steps}

    overlap = 0.0
    confirm, wait = [], []
    tool_ms = 0.0
    for c in calls:
        start, end = c.get("t_exec_start", -1), c.get("t_exec_end", -1)
        if start >= 0 and end >= start:
            tool_ms += end - start
            stream_end = end_of.get((c["turn"], c["step"]))
            if stream_end is not None and start < stream_end:
                overlap += min(end, stream_end) - start
        if c.get("t_call_open", -1) >= 0 and c.get("t_call_ready", -1) >= 0:
            confirm.append(c["t_call_ready"] - c["t_call_open"])
        if c.get("t_call_ready", -1) >= 0 and start >= 0:
            wait.append(start - c["t_call_ready"])

    model = res.get("model", "")
    if model in ("", "default") and steps:
        model = steps[0].get("model", model)
    return {
        "run": res.get("run", os.path.basename(run_dir)),
        "fixture": res.get("fixture", ""),
        "protocol": res.get("protocol", ""),
        "provider": res.get("provider", ""),
        "model": model,
        "outcome": res.get("outcome", ""),
        "pass": 1 if res.get("outcome") == "pass" else 0,
        "wall_s": res.get("wall_s", 0),
        "turn_s": round(sum(t.get("t_total", 0) for t in turns) / 1000, 3),
        "steps": len(steps),
        "calls": len(calls),
        "calls_failed": sum(1 for c in calls if c.get("exit", 0) != 0
                            or c.get("status") != "ok"),
        "llm_s": round(sum(s.get("t_stream_end", 0) for s in steps) / 1000, 3),
        "tool_s": round(tool_ms / 1000, 3),
        "overlap_s": round(overlap / 1000, 3),
        "confirm_ms": round(statistics.mean(confirm), 1) if confirm else "",
        "wait_ms": round(statistics.mean(wait), 1) if wait else "",
        "ctl_errors": sum(1 for s in steps if s.get("parse_error")),
        "retries": sum(s.get("retries", 0) for s in steps),
        "prompt_tokens": sum(s.get("prompt_tokens", 0) for s in steps),
        "completion_tokens": sum(s.get("completion_tokens", 0) for s in steps),
        "cached_tokens": sum(s.get("cached_tokens", 0) for s in steps),
    }


def mean(rows, key, nd=2):
    vals = [r[key] for r in rows if r[key] != ""]
    if not vals:
        return "-"
    v = statistics.mean(vals)
    return int(round(v)) if nd == 0 else round(v, nd)


def table(rows, keys, title):
    groups = defaultdict(list)
    for r in rows:
        groups[tuple(r[k] for k in keys)].append(r)
    cols = keys + ["n", "pass%", "steps", "wall_s", "llm_s", "tool_s",
                   "overlap_s", "wait_ms", "ctl_err", "tok_in", "tok_out"]
    out = [f"### {title}", "", "| " + " | ".join(cols) + " |",
           "|" + "---|" * len(cols)]
    for g in sorted(groups):
        rs = groups[g]
        n = len(rs)
        cells = list(g) + [
            n,
            round(100.0 * sum(r["pass"] for r in rs) / n),
            mean(rs, "steps", 1),
            mean(rs, "wall_s"),
            mean(rs, "llm_s"),
            mean(rs, "tool_s"),
            mean(rs, "overlap_s"),
            mean(rs, "wait_ms", 0),
            mean(rs, "ctl_errors", 2),
            mean(rs, "prompt_tokens", 0),
            mean(rs, "completion_tokens", 0),
        ]
        out.append("| " + " | ".join(str(c) for c in cells) + " |")
    out.append("")
    return "\n".join(out)


def main(argv):
    if len(argv) < 2:
        print(__doc__)
        return 64
    root = argv[1]
    md_path = os.path.join(root, "summary.md")
    csv_path = os.path.join(root, "results.csv")
    if "--md" in argv:
        md_path = argv[argv.index("--md") + 1]
    if "--csv" in argv:
        csv_path = argv[argv.index("--csv") + 1]

    rows = []
    for name in sorted(os.listdir(root)):
        d = os.path.join(root, name)
        if os.path.isfile(os.path.join(d, "result.json")):
            r = analyze(d)
            if r:
                rows.append(r)
    if not rows:
        print(f"no runs under {root}")
        return 1

    with open(csv_path, "w", newline="", encoding="utf-8") as f:
        w = csv.DictWriter(f, fieldnames=FIELDS)
        w.writeheader()
        w.writerows(rows)

    md = [f"## Results ({len(rows)} runs)", "",
          "wait_ms = call confirmed -> execution start (~0 for eager "
          "protocols, large when buffered until stream end). overlap_s = "
          "tool time hidden behind generation.", "",
          table(rows, ["model", "protocol"], "By model x protocol"),
          table(rows, ["model", "protocol", "fixture"],
                "By model x protocol x fixture")]
    text = "\n".join(md)
    with open(md_path, "w", encoding="utf-8") as f:
        f.write(text)
    print(text)
    print(f"wrote {csv_path} and {md_path}")
    return 0


if __name__ == "__main__":
    sys.exit(main(sys.argv))