"""Extraction behind PLRN-031 §1's measured frontier-call distribution.

Pulls the TraceLab v0.0.2 release DuckDB (SyFI Lab, University of Washington;
CC BY 4.0) and prints, as JSON:

- per-call latency quantiles, overall, by provider, and cut by output tokens
  (>=1,000 and >=2,000; the generation-heavy classes §1 cites);
- per-turn model-time quantiles (sum of call spans within a turn, and the
  first-to-last span);
- turn structure (rounds per turn, tool calls per round, share of rounds with
  a tool call), feeding §4's one-governed-action baseline.

Definitions: call latency is the span between a round's first and last timing
events; turn model time is the sum of call spans within a turn; dates are UTC.

Run with an isolated duckdb dependency:

    uv run --with duckdb python plrn031_tracelab_extract.py [--db PATH]

With no --db, the release asset is downloaded (~160 MB) to a temporary
directory, verified against the pinned SHA-256 below, and deleted on exit.
"""

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

RELEASE_URL = (
    "https://github.com/uw-syfi/TraceLab/releases/download/"
    "v0.0.2/syfi_coding_trace.duckdb"
)
SHA256 = "a7bab286bc640844560850965ccf47975cf66407154132abaab90f27ec9be744"

import duckdb

QS = [0.10, 0.25, 0.50, 0.75, 0.90, 0.95, 0.99]


def quantiles(con, table, col, cond="true"):
    sel = ", ".join(f"quantile_cont({col}, {q})" for q in QS)
    (row,) = con.execute(
        f"select count(*), round(avg({col})/1000, 1), {sel} "
        f"from {table} where {cond}"
    ).fetchall()
    n, mean, vals = row[0], row[1], row[2:]
    return {
        "n": n,
        "mean_s": mean,
        **{f"p{int(q * 100)}_s": round(v / 1000, 1) for q, v in zip(QS, vals)},
    }


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


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--db", help="path to a local copy of the DuckDB release")
    args = ap.parse_args()

    with tempfile.TemporaryDirectory() as tmp:
        db = Path(args.db) if args.db else Path(tmp) / "syfi_coding_trace.duckdb"
        if not db.exists():
            print(f"downloading {RELEASE_URL}", file=sys.stderr)
            fetch(db)
        con = duckdb.connect(str(db), read_only=True)
        run(con)


def run(con):
    con.execute(
        """
        create temp view call_lat as
        select t.round_pk, r.provider, r.session_id, r.turn_id,
               r.output_tokens,
               epoch_ms(max(t.timestamp)) - epoch_ms(min(t.timestamp)) as dur_ms
        from timing_events t join rounds r using(round_pk)
        group by 1,2,3,4,5
        """
    )
    # A turn is Codex's native turn_id (non-null) or, for Claude rounds (which
    # carry no turn_id), the session's user-message-count boundary. Null
    # turn_ids are excluded rather than coalesced.
    con.execute(
        """
        create temp view turn_lat as
        select tkey, count(*) as rounds,
               sum(dur_ms) as model_ms,
               epoch_ms(max(last)) - epoch_ms(min(first)) as span_ms
        from (
            select tkey, round_pk, dur_ms,
                   min(first_ts) over (partition by tkey) as first,
                   max(last_ts)  over (partition by tkey) as last
            from (
                select case when r.provider = 'codex'
                            then r.session_id || '|' || r.turn_id::varchar
                            else r.session_id || '|um'
                                 || r.current_user_message_count::varchar
                       end as tkey,
                       r.round_pk, c.dur_ms,
                       min(t.timestamp) as first_ts, max(t.timestamp) as last_ts
                from timing_events t
                join rounds r using(round_pk)
                join call_lat c using(round_pk)
                where r.provider <> 'codex' or r.turn_id is not null
                group by 1, 2, 3
            )
        )
        group by 1
        """
    )

    out = {
        "dataset": {
            "release": "TraceLab v0.0.2 (2026-07-24)",
            "url": RELEASE_URL,
            "sha256_duckdb": SHA256,
            "license": "CC BY 4.0",
        },
        "window_utc": [str(x) for x in con.execute(
            "select min(timestamp), max(timestamp) from timing_events"
        ).fetchone()],
        "counts": dict(
            zip(
                ["calls", "sessions", "users", "tool_calls_recorded"],
                con.execute(
                    "select (select count(*) from rounds),"
                    " (select count(distinct session_id) from rounds),"
                    " (select count(distinct \"user\") from rounds),"
                    " (select count(*) from tool_calls)"
                ).fetchone(),
            )
        ),
        "per_call_s": {
            "all": quantiles(con, "call_lat", "dur_ms"),
            "claude": quantiles(con, "call_lat", "dur_ms", "provider='claude'"),
            "codex": quantiles(con, "call_lat", "dur_ms", "provider='codex'"),
            "output_ge_1k": quantiles(con, "call_lat", "dur_ms",
                                      "output_tokens>=1000"),
            "output_ge_2k": quantiles(con, "call_lat", "dur_ms",
                                      "output_tokens>=2000"),
        },
        "per_turn_s": {
            "model_time_all": quantiles(con, "turn_lat", "model_ms"),
            "model_time_ge_2_rounds": quantiles(con, "turn_lat", "model_ms",
                                                "rounds>=2"),
            "span_all": quantiles(con, "turn_lat", "span_ms"),
        },
        "turn_structure": {
            "rounds_per_turn": {
                "quantiles": con.execute(
                    "select round(quantile_cont(rounds,0.5),1),"
                    " round(quantile_cont(rounds,0.9),1),"
                    " round(quantile_cont(rounds,0.99),1), round(avg(rounds),1)"
                    " from turn_lat"
                ).fetchone(),
                "max_rounds_in_one_turn": con.execute(
                    "select max(rounds) from turn_lat"
                ).fetchone()[0],
                "mean_excluding_top_turn": con.execute(
                    "select round(avg(rounds),1) from turn_lat"
                    " where rounds < (select max(rounds) from turn_lat)"
                ).fetchone()[0],
            },
            "tool_calls_per_round": con.execute(
                "select round(quantile_cont(n,0.5),1),"
                " round(quantile_cont(n,0.9),1), round(quantile_cont(n,0.99),1)"
                " from (select round_pk, count(*) n from tool_calls group by 1)"
            ).fetchone(),
            "rounds_with_tool_call_pct": con.execute(
                "select round(100.0*count(distinct round_pk)"
                " /(select count(*) from rounds),1) from tool_calls"
            ).fetchone()[0],
            "turns_total": con.execute("select count(*) from turn_lat").fetchone()[0],
        },
        # k = governed actions per note-turn (tool calls emitted by a round).
        # Histogram over all rounds, 0 through 7 and 8+, raw and excluding
        # continuation rows (retries of an already-emitted call).
        "k_distribution": {
            "raw": con.execute(
                "select least(coalesce(n,0), 8) as k, count(*)"
                " from rounds r left join"
                "  (select round_pk, count(*) n from tool_calls group by 1)"
                "  t using(round_pk) group by 1 order by 1"
            ).fetchall(),
            "excluding_continuations": con.execute(
                "select least(coalesce(n,0), 8) as k, count(*)"
                " from rounds r left join"
                "  (select round_pk, count(*) n from tool_calls"
                "   where continuation_of_tool_call_id is null group by 1)"
                "  t using(round_pk) group by 1 order by 1"
            ).fetchall(),
            "mean_k_all_rounds": con.execute(
                "select round(1.0*(select count(*) from tool_calls)"
                " /(select count(*) from rounds), 3)"
            ).fetchone()[0],
        },
    }
    json.dump(out, sys.stdout, indent=2)
    print()


if __name__ == "__main__":
    main()