#!/usr/bin/env python3
"""Analyze Send Comm Request (SCR) rounds: 0x03 request → 0x0B reply.

Pairs events across one or more fixture data dirs (typically pos 1 + 4, or all four).

Usage:
  python3 analyze_scr.py "../4/data" "../1/data"
  python3 analyze_scr.py --tid BC37 "sessions/.../4/data" "sessions/.../1/data"
"""

from __future__ import annotations

import argparse
import json
import sys
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, Optional


@dataclass
class ScrRound:
    tid: str
    request_perf_ns: int
    request_utc: str
    request_pos: int
    request_byte_span_us: Optional[float]
    request_checksum_ok: Optional[bool] = None
    request_wire_len_ok: Optional[bool] = None
    request_origin: Optional[str] = None
    reply_perf_ns: Optional[int] = None
    reply_utc: Optional[str] = None
    reply_pos: Optional[int] = None
    reply_checksum_ok: Optional[bool] = None
    reply_wire_len_ok: Optional[bool] = None
    reply_origin: Optional[str] = None
    scr_rtt_ms: Optional[float] = None
    reply_queue_wait_us: Optional[float] = None
    reply_inject_path: Optional[str] = None
    reply_inject_timeout: bool = False
    extra_requests: list[dict[str, Any]] = field(default_factory=list)
    extra_replies: list[dict[str, Any]] = field(default_factory=list)
    echo_reinject_03: bool = False
    eth_relay_skipped: bool = False

    @property
    def rtt_ms(self) -> Optional[float]:
        if self.scr_rtt_ms is not None:
            return self.scr_rtt_ms
        if self.reply_perf_ns is None:
            return None
        return (self.reply_perf_ns - self.request_perf_ns) / 1_000_000.0

    def classify(self) -> str:
        if self.reply_perf_ns is None and self.scr_rtt_ms is None:
            return "no_reply"
        if self.request_checksum_ok is False or self.reply_checksum_ok is False:
            return "bad_frame"
        if self.extra_requests:
            return "retry"
        if self.extra_replies:
            return "duplicate_reply"
        if self.echo_reinject_03 or self.eth_relay_skipped:
            return "echo_polluted"
        rtt = self.rtt_ms or 0.0
        if self.reply_inject_timeout or rtt > 120.0 or (self.reply_queue_wait_us or 0) > 15_000:
            return "delayed"
        if rtt > 80.0 or (self.reply_queue_wait_us or 0) > 5_000:
            return "marginal"
        return "clean"


def load_events(path: Path) -> list[dict[str, Any]]:
    out: list[dict[str, Any]] = []
    with path.open(encoding="utf-8") as fh:
        for line in fh:
            line = line.strip()
            if not line:
                continue
            try:
                out.append(json.loads(line))
            except json.JSONDecodeError:
                continue
    return out


def events_by_pos(dirs: list[Path]) -> dict[int, list[dict[str, Any]]]:
    by_pos: dict[int, list[dict[str, Any]]] = {}
    for d in dirs:
        ev_path = d / "events.jsonl"
        if not ev_path.is_file():
            continue
        for ev in load_events(ev_path):
            pos = int(ev.get("fixture_pos", 0))
            by_pos.setdefault(pos, []).append(ev)
    for pos in by_pos:
        by_pos[pos].sort(key=lambda e: e.get("t_perf_ns", 0))
    return by_pos


def find_scr_rounds(
    events: list[dict[str, Any]],
    *,
    tid_filter: Optional[str],
    anchor_pos: int,
    prefer_scr_round: bool = False,
) -> list[ScrRound]:
    """Build SCR rounds from scr_round events or packet_complete 0x03 on anchor_pos."""
    if prefer_scr_round and any(ev.get("ev") == "scr_round" for ev in events):
        return _find_scr_rounds_from_scr_round(events, tid_filter=tid_filter, anchor_pos=anchor_pos)
    return _find_scr_rounds_from_packet_complete(
        events, tid_filter=tid_filter, anchor_pos=anchor_pos
    )


def _find_scr_rounds_from_scr_round(
    events: list[dict[str, Any]],
    *,
    tid_filter: Optional[str],
    anchor_pos: int,
) -> list[ScrRound]:
    rounds: list[ScrRound] = []
    open_round: Optional[ScrRound] = None

    for ev in events:
        if int(ev.get("fixture_pos", 0)) != anchor_pos:
            continue
        if ev.get("ev") != "scr_round":
            continue
        tid = str(ev.get("tid") or "").upper()
        if tid_filter and tid != tid_filter.upper():
            continue
        phase = ev.get("phase")
        origin = ev.get("origin")
        if phase == "request":
            if open_round and open_round.reply_perf_ns is None and open_round.scr_rtt_ms is None:
                open_round.extra_requests.append({"utc": ev.get("utc"), "origin": origin})
            else:
                open_round = ScrRound(
                    tid=tid or "?",
                    request_perf_ns=int(ev.get("t_perf_ns", 0)),
                    request_utc=str(ev.get("utc", "")),
                    request_pos=anchor_pos,
                    request_byte_span_us=ev.get("byte_span_us"),
                    request_checksum_ok=ev.get("checksum_ok"),
                    request_wire_len_ok=ev.get("wire_len_ok"),
                    request_origin=origin,
                )
                rounds.append(open_round)
        elif phase == "reply" and open_round is not None:
            if open_round.reply_perf_ns is None:
                open_round.reply_perf_ns = int(ev.get("t_perf_ns", 0))
                open_round.reply_utc = str(ev.get("utc", ""))
                open_round.reply_pos = anchor_pos
                open_round.reply_checksum_ok = ev.get("checksum_ok")
                open_round.reply_wire_len_ok = ev.get("wire_len_ok")
                open_round.reply_origin = origin
                if ev.get("rtt_ms") is not None:
                    open_round.scr_rtt_ms = float(ev["rtt_ms"])
            else:
                open_round.extra_replies.append({"utc": ev.get("utc"), "origin": origin})

    return rounds


def _find_scr_rounds_from_packet_complete(
    events: list[dict[str, Any]],
    *,
    tid_filter: Optional[str],
    anchor_pos: int,
) -> list[ScrRound]:
    rounds: list[ScrRound] = []
    open_round: Optional[ScrRound] = None

    for ev in events:
        if int(ev.get("fixture_pos", 0)) != anchor_pos:
            continue
        ev_name = ev.get("ev")
        msg = str(ev.get("msg", "")).upper()
        tid = str(ev.get("tid") or "").upper()
        if tid_filter and tid != tid_filter.upper():
            continue

        if ev_name == "packet_complete" and msg == "03":
            if open_round and open_round.reply_perf_ns is None:
                open_round.extra_requests.append(
                    {"utc": ev.get("utc"), "byte_span_us": ev.get("byte_span_us")}
                )
            else:
                open_round = ScrRound(
                    tid=tid or "?",
                    request_perf_ns=int(ev.get("t_perf_ns", 0)),
                    request_utc=str(ev.get("utc", "")),
                    request_pos=anchor_pos,
                    request_byte_span_us=ev.get("byte_span_us"),
                )
                rounds.append(open_round)

        elif (
            open_round
            and ev_name == "packet_complete"
            and msg == "0B"
            and (not tid or tid == open_round.tid or open_round.tid == "?")
        ):
            if open_round.reply_perf_ns is None:
                open_round.reply_perf_ns = int(ev.get("t_perf_ns", 0))
                open_round.reply_utc = str(ev.get("utc", ""))
                open_round.reply_pos = anchor_pos
            else:
                open_round.extra_replies.append({"utc": ev.get("utc")})

        elif (
            open_round
            and ev_name == "eth_rx_skip"
            and ev.get("reason")
            in (
                "uart_origin_echo",
                "outgoing_pkttype",
                "local_eth_tx",
                "upstream_uart_echo",
                "eth_inject_echo",
                "scr_master_native_only",
            )
            and open_round.reply_perf_ns is None
        ):
            if ev.get("reason") == "scr_master_native_only":
                open_round.eth_relay_skipped = True
            else:
                open_round.echo_reinject_03 = True

        elif (
            open_round
            and ev_name == "eth_to_uart_done"
            and msg == "0B"
            and str(ev.get("tid", "")).upper() in {open_round.tid, ""}
            and open_round.reply_perf_ns is not None
        ):
            if open_round.reply_queue_wait_us is None:
                open_round.reply_queue_wait_us = ev.get("queue_wait_us")

        elif (
            open_round
            and ev_name == "inject_hs_timeout"
            and msg == "0B"
            and str(ev.get("tid", "")).upper() in {open_round.tid, ""}
        ):
            open_round.reply_inject_timeout = True

        elif open_round and ev_name == "slv_out_on" and msg == "0B":
            if open_round.reply_inject_path is None:
                open_round.reply_inject_path = ev.get("inject_path")

    return rounds


def print_scr_diag_summary(events: list[dict[str, Any]], anchor_pos: int) -> None:
    diags = [
        ev
        for ev in events
        if ev.get("ev") == "scr_diag" and int(ev.get("fixture_pos", 0)) == anchor_pos
    ]
    if not diags:
        return
    print(f"\nscr_diag on pos {anchor_pos}: {len(diags)} events")
    by_origin: dict[str, int] = {}
    bad_ck = 0
    bad_len = 0
    for ev in diags:
        origin = str(ev.get("origin", "?"))
        by_origin[origin] = by_origin.get(origin, 0) + 1
        if ev.get("checksum_ok") is False:
            bad_ck += 1
        if ev.get("wire_len_ok") is False:
            bad_len += 1
    for origin, n in sorted(by_origin.items()):
        print(f"  {origin}: {n}")
    if bad_ck or bad_len:
        print(f"  checksum_ok=false: {bad_ck}  wire_len_ok=false: {bad_len}")


def print_session_start(by_pos: dict[int, list[dict[str, Any]]]) -> None:
    print("Session config:")
    for pos in sorted(by_pos):
        for ev in by_pos[pos]:
            if ev.get("ev") == "session_start":
                keys = (
                    "broadcast_engine",
                    "directional_routing",
                    "eth_ports",
                    "route_hold_ms",
                    "drop_outgoing_pkttype",
                    "drop_local_eth_tx",
                    "daisy_chain_router",
                    "inject_on_hs_active",
                    "post_hs_inactive_us",
                    "hs_inactive_wait_ms",
                    "pre_tx_us",
                    "scr_diag",
                    "scr_master_native_only",
                )
                cfg = {k: ev.get(k) for k in keys if k in ev}
                print(f"  pos {pos}: {cfg}")
                break


def main() -> int:
    parser = argparse.ArgumentParser(description="Analyze SCR (0x03→0x0B) timing.")
    parser.add_argument("data_dirs", nargs="+", type=Path, help="fixture data/ directories")
    parser.add_argument("--tid", default="BC37", help="Terminal ID filter (default BC37)")
    parser.add_argument(
        "--anchor-pos",
        type=int,
        default=4,
        help="Position whose UART/scr_round anchors SCR (default 4 machine)",
    )
    parser.add_argument(
        "--master-native",
        action="store_true",
        help="Master-only bench: anchor pos 3, prefer scr_round events, any TID",
    )
    args = parser.parse_args()

    if args.master_native:
        args.anchor_pos = 3
        args.tid = ""

    by_pos = events_by_pos(args.data_dirs)
    if not by_pos:
        print("No events.jsonl found.", file=sys.stderr)
        return 1

    print_session_start(by_pos)

    anchor = args.anchor_pos
    if anchor not in by_pos:
        print(f"No events for anchor pos {anchor}", file=sys.stderr)
        return 1

    tid_filter = args.tid or None
    rounds = find_scr_rounds(
        by_pos[anchor],
        tid_filter=tid_filter,
        anchor_pos=anchor,
        prefer_scr_round=args.master_native,
    )
    print_scr_diag_summary(by_pos[anchor], anchor)

    counts: dict[str, int] = {}
    print(f"\nSCR rounds on pos {anchor} (tid={args.tid or 'any'}): {len(rounds)}\n")
    print(
        f"{'#':>3}  {'class':<16}  {'rtt_ms':>7}  {'q_wait_ms':>9}  "
        f"{'req_span_us':>11}  {'ck_ok':>5}  origin  inject  timeout  notes"
    )
    print("-" * 96)

    for i, r in enumerate(rounds, 1):
        cls = r.classify()
        counts[cls] = counts.get(cls, 0) + 1
        q_ms = (r.reply_queue_wait_us or 0) / 1000.0
        ck = "Y" if r.request_checksum_ok and r.reply_checksum_ok else (
            "N" if r.request_checksum_ok is False or r.reply_checksum_ok is False else "-"
        )
        notes: list[str] = []
        if r.echo_reinject_03:
            notes.append("eth_echo_03")
        if r.eth_relay_skipped:
            notes.append("relay_03_skipped")
        if r.extra_requests:
            notes.append(f"+{len(r.extra_requests)}_retry_03")
        if r.extra_replies:
            notes.append(f"+{len(r.extra_replies)}_dup_0B")
        print(
            f"{i:3d}  {cls:<16}  "
            f"{(r.rtt_ms or 0):7.1f}  {q_ms:9.1f}  "
            f"{(r.request_byte_span_us or 0):11.1f}  "
            f"{ck:>5}  "
            f"{(r.request_origin or '-'):<8}  "
            f"{(r.reply_inject_path or '-'):<8}  "
            f"{'Y' if r.reply_inject_timeout else 'N':7}  "
            f"{' '.join(notes)}"
        )

    print("\nSummary:")
    for cls in (
        "clean",
        "marginal",
        "delayed",
        "retry",
        "duplicate_reply",
        "echo_polluted",
        "bad_frame",
        "no_reply",
    ):
        if counts.get(cls):
            print(f"  {cls}: {counts[cls]}")

    # Trunk witness (pos 1) RTT if available
    if 1 in by_pos and anchor != 1:
        trunk = find_scr_rounds(by_pos[1], tid_filter=args.tid or None, anchor_pos=1)
        trunk_rtts = [t.rtt_ms for t in trunk if t.rtt_ms is not None]
        if trunk_rtts:
            print(f"\nPos 1 trunk SCR RTT ms: min={min(trunk_rtts):.1f} max={max(trunk_rtts):.1f} "
                  f"avg={sum(trunk_rtts)/len(trunk_rtts):.1f} (n={len(trunk_rtts)})")

    return 0


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