"""Cross-check behind PLRN-031 §5/Appendix A: tree structure of agent sessions.

Streams SemiAnalysis's WEKA Claude Code trace corpus
(huggingface.co/datasets/semianalysisai/cc-traces-weka-062126), which counts a
main agent and its subagents as one session tree, and prints, as JSON:

- fan-out share (trees with at least one subagent group);
- subagent groups per fan-out tree and inner calls per tree;
- total model calls per tree, fan-out vs not;
- per-call api_time quantiles for the main stream and for subagent-inner calls
  (an independent corroboration of §1's TraceLab-measured call distribution).

The corpus records each subagent group as an envelope row (type "subagent")
with its inner calls nested under "requests"; inner calls carry their own
api_time. Main-stream calls are the top-level rows of type "s"/"n".

    python3 plrn031_weka_tree_check.py [--input traces.jsonl]

With no --input, the corpus (1.9 GB JSONL) is streamed from Hugging Face to a
temporary file and verified against the pinned SHA-256.
"""

import argparse
import hashlib
import json
import statistics
import sys
import tempfile
import urllib.request
from pathlib import Path

URL = ("https://huggingface.co/datasets/semianalysisai/"
       "cc-traces-weka-062126/resolve/main/traces.jsonl")
SHA256 = ("29b6a19e751ff5230771519aab755f80"
          "a0f43a4ba9cf96b72d3a6a437ec99276")


def quantiles(vals, ps=(10, 25, 50, 75, 90, 95, 99)):
    vals = sorted(vals)
    n = len(vals)
    return {f"p{p}_s": round(vals[min(n - 1, int(n * p / 100))], 2)
            for p in ps}


def fetch(path):
    with urllib.request.urlopen(URL) as r, open(path, "wb") as f:
        while chunk := r.read(1 << 22):
            f.write(chunk)
    got = hashlib.sha256(Path(path).read_bytes()).hexdigest()
    if got != SHA256:
        sys.exit(f"sha256 mismatch: got {got}, want {SHA256}")


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--input", help="local traces.jsonl")
    args = ap.parse_args()

    with tempfile.TemporaryDirectory() as tmp:
        path = Path(args.input) if args.input else Path(tmp) / "traces.jsonl"
        if not path.exists():
            print(f"downloading {URL} (~1.9 GB)", file=sys.stderr)
            fetch(path)
        stats = analyze(path)
    json.dump(stats, sys.stdout, indent=2)
    print()


def analyze(path):
    main_api, sub_api = [], []
    group_counts, group_dur = [], []
    tree_sizes, fan_sizes, flat_sizes = [], [], []
    tot_main = tot_sub = 0

    with open(path) as f:
        for line in f:
            tr = json.loads(line)
            main = [q for q in tr["requests"] if q["type"] in ("s", "n")]
            groups = [q for q in tr["requests"] if q["type"] == "subagent"]
            sub = [r for g in groups for r in g.get("requests", [])]
            tot_main += len(main)
            tot_sub += len(sub)
            size = len(main) + len(sub)
            tree_sizes.append(size)
            if groups:
                group_counts.append(len(groups))
                fan_sizes.append(size)
                group_dur.extend(g["duration_ms"] / 1000 for g in groups)
            else:
                flat_sizes.append(size)
            main_api.extend(q["api_time"] for q in main if "api_time" in q)
            sub_api.extend(q["api_time"] for q in sub if "api_time" in q)

    n = len(tree_sizes)
    return {
        "corpus": {
            "name": "cc-traces-weka-062126", "url": URL, "sha256": SHA256,
            "license": "Apache-2.0",
        },
        "trees": n,
        "fan_out_trees": len(fan_sizes),
        "fan_out_share_pct": round(100 * len(fan_sizes) / n, 1),
        "calls": {"main": tot_main, "subagent_inner": tot_sub,
                  "subagent_share_pct": round(100 * tot_sub / (tot_main + tot_sub), 1)},
        "subagent_groups_per_fanout_tree": {
            "mean": round(statistics.mean(group_counts), 1),
            **quantiles(group_counts)},
        "subagent_group_wall_s": quantiles(group_dur),
        "calls_per_tree": {
            "all": {"mean": round(statistics.mean(tree_sizes), 1),
                    "median": round(statistics.median(tree_sizes), 1)},
            "fan_out": {"mean": round(statistics.mean(fan_sizes), 1),
                        "median": round(statistics.median(fan_sizes), 1)},
            "flat": {"mean": round(statistics.mean(flat_sizes), 1),
                     "median": round(statistics.median(flat_sizes), 1)},
        },
        "api_time_s": {"main": quantiles(main_api),
                       "subagent_inner": quantiles(sub_api)},
    }


if __name__ == "__main__":
    main()