#!/usr/bin/env python3
"""Phase C+23 jitter profile probe.

Reads canary jsonls (jitter-1/2/3 + cold-c21 + any fresh runs) and extracts:
- total tid=6 events seen within the first ~120k indices
- the exact event sequence on tid=6 around idx [104,595..104,620]
- count of wait.begin events on tid=6 by SID
- count of contention-prone events (wait.begin, kernel.call RtlEnter / RtlLeave)

Designed to stream line-by-line and not load multi-GB jsonls into RAM.
"""

import json
import os
import sys
from collections import Counter, defaultdict

WINDOW_LO = 104_595
WINDOW_HI = 104_620
TID = 6
TID_EVENT_LIMIT = 120_000


def profile(path: str):
    if not os.path.exists(path):
        return None
    tid6_events = 0
    waitbegins = Counter()
    importcalls = Counter()
    kernelcalls = Counter()
    window_events = []
    tid_idx = -1

    with open(path, "rb") as fh:
        for raw in fh:
            # Cheap reject before json parse: must contain `"tid":6,`
            if b'"tid":6,' not in raw and b'"tid": 6,' not in raw:
                continue
            try:
                ev = json.loads(raw)
            except Exception:
                continue
            if ev.get("tid") != TID:
                continue
            tid_idx = ev.get("tid_event_idx", tid_idx + 1)
            tid6_events += 1
            kind = ev.get("kind", "")
            if kind == "wait.begin":
                sids = ev.get("payload", {}).get("handles_semantic_ids") or []
                for s in sids:
                    waitbegins[s] += 1
            elif kind == "import.call":
                name = ev.get("payload", {}).get("name", "")
                importcalls[name] += 1
            elif kind == "kernel.call":
                name = ev.get("payload", {}).get("name", "")
                kernelcalls[name] += 1

            if WINDOW_LO <= tid_idx <= WINDOW_HI:
                summary = {
                    "idx": tid_idx,
                    "kind": kind,
                    "name": ev.get("payload", {}).get("name", ""),
                }
                if kind == "wait.begin":
                    summary["sid"] = (ev.get("payload", {}).get("handles_semantic_ids") or [None])[0]
                    summary["timeout_ns"] = ev.get("payload", {}).get("timeout_ns")
                window_events.append(summary)

            if tid_idx > TID_EVENT_LIMIT:
                break

    return {
        "path": path,
        "tid6_total_seen": tid6_events,
        "waitbegins_by_sid": dict(waitbegins.most_common(10)),
        "rtlenter_calls": importcalls.get("RtlEnterCriticalSection", 0),
        "rtlleave_calls": importcalls.get("RtlLeaveCriticalSection", 0),
        "window_events": window_events,
    }


def main(paths):
    out = {}
    for p in paths:
        print(f"profiling {p}...", file=sys.stderr)
        r = profile(p)
        if r is None:
            print(f"  (missing)", file=sys.stderr)
            continue
        out[os.path.basename(p)] = r
    json.dump(out, sys.stdout, indent=2)
    print()


if __name__ == "__main__":
    main(sys.argv[1:])