#!/usr/bin/env python3
"""Offline analyzer for EG 3.0.1.6+ hs_timing.jsonl traces.

Timestamps:
  - k_* events: kernel URB receipt times (ktime), carried in field "c"
  - us_* events: CLOCK_MONOTONIC in uartd ("mono_ns")
  Cross-domain deltas are informative only (same host, different clocks).

Usage:
  python3 hs_timing_analyze.py path/to/hs_timing.jsonl
  python3 hs_timing_analyze.py --self-test
"""

from __future__ import annotations

import argparse
import json
import math
import sys
from pathlib import Path
from typing import Any, Dict, List, Optional, Sequence, Tuple


def load_events(path: Path) -> List[Dict[str, Any]]:
    events: List[Dict[str, Any]] = []
    with path.open("r", encoding="utf-8") as f:
        for line_no, line in enumerate(f, 1):
            line = line.strip()
            if not line:
                continue
            try:
                ev = json.loads(line)
            except json.JSONDecodeError as exc:
                raise SystemExit(f"{path}:{line_no}: invalid JSON: {exc}") from exc
            events.append(ev)
    return events


def percentile(sorted_vals: Sequence[float], p: float) -> Optional[float]:
    if not sorted_vals:
        return None
    if len(sorted_vals) == 1:
        return float(sorted_vals[0])
    k = (len(sorted_vals) - 1) * (p / 100.0)
    f = math.floor(k)
    c = math.ceil(k)
    if f == c:
        return float(sorted_vals[int(k)])
    d0 = sorted_vals[f] * (c - k)
    d1 = sorted_vals[c] * (k - f)
    return float(d0 + d1)


def summarize_ns(samples: List[float]) -> Dict[str, Any]:
    if not samples:
        return {"n": 0}
    s = sorted(samples)
    return {
        "n": len(s),
        "p50_us": round(percentile(s, 50) / 1000.0, 3),  # type: ignore[operator]
        "p95_us": round(percentile(s, 95) / 1000.0, 3),  # type: ignore[operator]
        "p99_us": round(percentile(s, 99) / 1000.0, 3),  # type: ignore[operator]
        "max_us": round(s[-1] / 1000.0, 3),
        "min_us": round(s[0] / 1000.0, 3),
    }


def analyze(events: List[Dict[str, Any]]) -> Dict[str, Any]:
    by_type: Dict[str, int] = {}
    drops = 0
    rtt: List[float] = []
    poll_gap: List[float] = []
    hs_obs: List[Tuple[int, int, int]] = []  # mono, hs, rtt_ns
    k_gpio: List[Dict[str, Any]] = []
    k_rx: List[Dict[str, Any]] = []
    rounds: List[Dict[str, Any]] = []

    inject_open: Optional[Dict[str, Any]] = None
    x1_at: Optional[int] = None
    de_off_at: Optional[int] = None

    for ev in events:
        t = str(ev.get("type", ""))
        by_type[t] = by_type.get(t, 0) + 1
        if t == "us_drop":
            drops += int(ev.get("a", 0) or 0)
        elif t == "us_gpio_stat_rtt":
            rtt.append(float(ev.get("a", 0) or 0))
        elif t == "us_poll_wake":
            poll_gap.append(float(ev.get("a", 0) or 0))
        elif t == "us_hs_obs":
            hs_obs.append((int(ev.get("mono_ns", 0)), int(ev.get("a", 0)), int(ev.get("b", 0))))
        elif t == "k_gpio_int":
            k_gpio.append(ev)
        elif t == "k_rx_urb":
            k_rx.append(ev)
        elif t == "us_inject_start":
            inject_open = {
                "start_ns": int(ev.get("mono_ns", 0)),
                "opcode": int(ev.get("b", 0)),
                "len": int(ev.get("a", 0)),
            }
            x1_at = None
            de_off_at = None
        elif t == "us_de_off" and inject_open is not None:
            de_off_at = int(ev.get("mono_ns", 0))
            inject_open["de_off_ns"] = de_off_at
            inject_open["de_reason"] = int(ev.get("a", 0))
            inject_open["de_hold_us"] = int(ev.get("b", 0))
        elif t == "us_x1_seen" and inject_open is not None:
            x1_at = int(ev.get("mono_ns", 0))
            inject_open["x1_ns"] = x1_at
            inject_open["x1_msg"] = int(ev.get("a", 0))
            inject_open["harvest"] = int(ev.get("b", 0))
        elif t == "us_auto_close_done" and inject_open is not None:
            done_ns = int(ev.get("mono_ns", 0))
            inject_open["ac_done_ns"] = done_ns
            inject_open["ac_err"] = int(ev.get("a", 0))
            if x1_at is not None:
                inject_open["x1_to_0d_ns"] = done_ns - x1_at
            if de_off_at is not None and x1_at is not None:
                inject_open["de_off_to_x1_ns"] = x1_at - de_off_at
            rounds.append(inject_open)
            inject_open = None
        elif t == "us_inject_done" and inject_open is not None:
            inject_open["inject_done_ns"] = int(ev.get("mono_ns", 0))
            inject_open["inject_err"] = int(ev.get("a", 0))
            # Close non-gameplay injects immediately; leave x9 open for auto-close.
            if (int(inject_open.get("opcode", 0)) & 0x0F) != 0x09:
                rounds.append(inject_open)
                inject_open = None
            elif "x1_ns" not in inject_open:
                # Incomplete x9 (NO_X1 / no auto-close) — still a round sample.
                rounds.append(inject_open)
                inject_open = None

    # Kernel GPIO1 rise→fall pairs (flags bit0x08 = SLV_IN active)
    k_hs_high_ns: Optional[int] = None
    k_hs_widths: List[float] = []
    k_gpio5_falls = 0
    for ev in k_gpio:
        flags = int(ev.get("a", 0) or 0)
        ts = int(ev.get("c", 0) or 0)
        if flags & 0x04:
            k_gpio5_falls += 1
        if flags & 0x01:  # GPIO1 activity in this INT
            active = 1 if (flags & 0x08) else 0
            if active and k_hs_high_ns is None:
                k_hs_high_ns = ts
            elif not active and k_hs_high_ns is not None:
                k_hs_widths.append(float(ts - k_hs_high_ns))
                k_hs_high_ns = None

    # Ordering: for each k_gpio with SLV_IN rise, nearest prior/next k_rx
    order_rx_before_hs = 0
    order_hs_before_rx = 0
    rx_ts = [int(e.get("c", 0) or 0) for e in k_rx]
    for ev in k_gpio:
        flags = int(ev.get("a", 0) or 0)
        if not (flags & 0x01 and flags & 0x08):
            continue
        ts = int(ev.get("c", 0) or 0)
        prior = [t for t in rx_ts if t <= ts]
        after = [t for t in rx_ts if t > ts]
        if prior and (not after or (ts - prior[-1]) <= (after[0] - ts)):
            order_rx_before_hs += 1
        elif after:
            order_hs_before_rx += 1

    # Coalesced HS: consecutive us_hs_obs with same level skipped by design;
    # detect kernel event_count gaps via drop markers and dense GPIO5/GPIO1 mix.
    x1_to_0d = [float(r["x1_to_0d_ns"]) for r in rounds if "x1_to_0d_ns" in r]
    de_to_x1 = [float(r["de_off_to_x1_ns"]) for r in rounds if "de_off_to_x1_ns" in r]

    # Userspace HS obs vs nearest kernel GPIO (same-host clock skew possible)
    edge_to_obs: List[float] = []
    k_gpio_ts = [
        (int(e.get("c", 0) or 0), int(e.get("a", 0) or 0))
        for e in k_gpio
        if int(e.get("a", 0) or 0) & 0x01
    ]
    for mono, hs, _rtt in hs_obs:
        # Best-effort: match by sequence order only when clocks not mixed.
        # Prefer comparing successive us_hs_obs RTT field (b) already captured.
        _ = hs
        _ = mono
    for _mono, _hs, rtt_ns in hs_obs:
        if rtt_ns > 0:
            edge_to_obs.append(float(rtt_ns))

    return {
        "counts": by_type,
        "ring_drops": drops,
        "gpio_stat_rtt": summarize_ns(rtt),
        "poll_wake_late": summarize_ns(poll_gap),
        "kernel_hs_width": summarize_ns(k_hs_widths),
        "kernel_gpio5_falls": k_gpio5_falls,
        "kernel_gpio_events": len(k_gpio),
        "kernel_rx_urbs": len(k_rx),
        "order_rx_before_hs_rise": order_rx_before_hs,
        "order_hs_rise_before_rx": order_hs_before_rx,
        "hs_obs_gpio_stat_rtt_on_edge": summarize_ns(edge_to_obs),
        "rounds": len(rounds),
        "de_off_to_x1": summarize_ns(de_to_x1),
        "x1_to_0d": summarize_ns(x1_to_0d),
        "round_samples": rounds[:20],
        "notes": [
            "k_* timestamps are host USB URB receipt times, not copper edges.",
            "us_* use CLOCK_MONOTONIC in uartd; do not subtract from k_* as absolute truth.",
            "us_gpio_stat_rtt.a is GET_REG round-trip nanoseconds.",
        ],
    }


def format_report(report: Dict[str, Any]) -> str:
    lines = [
        "HS timing analysis",
        f"  events: {sum(report['counts'].values())}",
        f"  by_type: {report['counts']}",
        f"  ring_drops: {report['ring_drops']}",
        f"  gpio_stat_rtt: {report['gpio_stat_rtt']}",
        f"  poll_wake_late(>=1.5ms): {report['poll_wake_late']}",
        f"  kernel_gpio_events: {report['kernel_gpio_events']} "
        f"(gpio5_falls={report['kernel_gpio5_falls']})",
        f"  kernel_rx_urbs: {report['kernel_rx_urbs']}",
        f"  kernel_hs_width: {report['kernel_hs_width']}",
        f"  order rx_before_hs_rise={report['order_rx_before_hs_rise']} "
        f"hs_before_rx={report['order_hs_rise_before_rx']}",
        f"  rounds: {report['rounds']}",
        f"  de_off_to_x1: {report['de_off_to_x1']}",
        f"  x1_to_0d: {report['x1_to_0d']}",
    ]
    for n in report["notes"]:
        lines.append(f"  note: {n}")
    return "\n".join(lines) + "\n"


def _self_test() -> None:
    fixtures = Path(__file__).resolve().parent / "tests" / "fixtures" / "hs_timing"
    cases = {
        "ordered_edges.jsonl": lambda r: r["kernel_gpio_events"] >= 2
        and r["order_hs_rise_before_rx"] >= 0,
        "int_after_bulk.jsonl": lambda r: r["order_rx_before_hs_rise"] >= 1,
        "coalesced_edges.jsonl": lambda r: r["counts"].get("k_gpio_int", 0) >= 1,
        "dropped_ring.jsonl": lambda r: r["ring_drops"] >= 3,
        "x9_x1_0d_complete.jsonl": lambda r: r["rounds"] >= 1
        and r["x1_to_0d"]["n"] >= 1,
        "x9_x1_incomplete.jsonl": lambda r: r["x1_to_0d"]["n"] == 0,
    }
    for name, pred in cases.items():
        path = fixtures / name
        report = analyze(load_events(path))
        if not pred(report):
            raise SystemExit(f"self-test failed: {name}: {report}")
    print("hs_timing_analyze self-test OK")


def main(argv: Optional[Sequence[str]] = None) -> int:
    p = argparse.ArgumentParser(description=__doc__)
    p.add_argument("trace", nargs="?", type=Path, help="hs_timing.jsonl path")
    p.add_argument("--json", action="store_true", help="emit machine JSON")
    p.add_argument("--self-test", action="store_true")
    args = p.parse_args(argv)
    if args.self_test:
        _self_test()
        return 0
    if not args.trace:
        p.error("trace path required (or --self-test)")
    report = analyze(load_events(args.trace))
    if args.json:
        json.dump(report, sys.stdout, indent=2)
        sys.stdout.write("\n")
    else:
        sys.stdout.write(format_report(report))
    return 0


if __name__ == "__main__":
    raise SystemExit(main())