#!/usr/bin/env python3
"""Match origin pos N packets.raw hex lines against a witness pos packets.raw.

Usage:
  python validate_origin_bytes.py ../sessions/MY_SESSION
  python validate_origin_bytes.py ../sessions/MY_SESSION --origin 4 --witness 1
  python validate_origin_bytes.py ../sessions/MY_SESSION --origin 4 --witness 1 --native-only --passive-only
  python validate_origin_bytes.py ../sessions/MY_SESSION --matrix
  python validate_origin_bytes.py ../sessions/MY_SESSION --matrix --passive-only
"""
from __future__ import annotations

import argparse
import json
from collections import Counter
from pathlib import Path

SOH = 0x05
EOP = 0x0D
WIRE_LEN = {0x03: 32, 0x0A: 82, 0x0B: 12}


def parse_raw_line(line: str) -> bytes | None:
    line = line.strip()
    if not line:
        return None
    try:
        raw = bytes(int(x, 16) for x in line.split(","))
    except ValueError:
        return None
    if len(raw) < 2 or raw[0] != SOH or raw[-1] != EOP:
        return None
    return raw


def validate_frame(raw: bytes) -> tuple[bool, str]:
    expected = WIRE_LEN.get(raw[1])
    if expected is not None and len(raw) != expected:
        return False, f"wire_len expected {expected} got {len(raw)}"
    return True, "ok"


def load_raw(pos: int, root: Path) -> Counter[str]:
    p = root / str(pos) / "data" / "packets.raw"
    if not p.is_file():
        return Counter()
    return Counter(
        ln.strip()
        for ln in p.read_text(encoding="utf-8").splitlines()
        if ln.strip()
    )


def load_events(pos: int, root: Path) -> list[dict]:
    p = root / str(pos) / "data" / "events.jsonl"
    if not p.is_file():
        return []
    rows: list[dict] = []
    with p.open(encoding="utf-8") as fh:
        for line in fh:
            line = line.strip()
            if line:
                rows.append(json.loads(line))
    return rows


def tid_le(raw: bytes) -> str:
    return f"{raw[4]:02X}{raw[3]:02X}" if len(raw) > 4 else "????"


def packet_complete_keys(pos: int, root: Path) -> Counter[tuple[str, str, int]]:
    rows = load_events(pos, root)
    pc_keys: Counter[tuple[str, str, int]] = Counter()
    for r in rows:
        if r.get("ev") != "packet_complete":
            continue
        pc_keys[(r.get("msg", ""), r.get("tid", ""), r.get("len", 0))] += 1
    return pc_keys


def lines_for_packet_complete(pos: int, root: Path) -> Counter[str]:
    """packets.raw lines tied to local passive UART packet_complete events."""
    raw = load_raw(pos, root)
    pc_keys = packet_complete_keys(pos, root)
    matched: Counter[str] = Counter()
    for line, n in raw.items():
        parsed = parse_raw_line(line)
        if not parsed:
            continue
        key = (f"{parsed[1]:02X}", tid_le(parsed), len(parsed))
        if pc_keys.get(key, 0) > 0:
            matched[line] = min(n, pc_keys[key])
    return matched


def native_origin_lines(pos: int, root: Path) -> Counter[str]:
    """Origin pos: packets.raw lines from local UART packet_complete (not eth inject)."""
    return lines_for_packet_complete(pos, root)


def witness_raw_lines(pos: int, root: Path, *, passive_only: bool) -> Counter[str]:
    if passive_only:
        return lines_for_packet_complete(pos, root)
    return load_raw(pos, root)


def compare_origin_to_witness(
    origin: int,
    witness: int,
    root: Path,
    *,
    native_only: bool,
    passive_only: bool,
) -> tuple[int, int, bool]:
    if native_only:
        src_raw = native_origin_lines(origin, root)
    else:
        src_raw = load_raw(origin, root)
    wit_raw = witness_raw_lines(witness, root, passive_only=passive_only)
    sent = sum(src_raw.values())
    matched = sum(min(n, wit_raw.get(line, 0)) for line, n in src_raw.items())
    ok = matched == sent and sent > 0
    return matched, sent, ok


def print_compare(
    origin: int,
    witness: int,
    root: Path,
    *,
    native_only: bool,
    passive_only: bool,
    fail_flag: list[bool],
) -> None:
    if native_only:
        src_raw = native_origin_lines(origin, root)
        label = "native origin"
    else:
        src_raw = load_raw(origin, root)
        label = "all raw"
    wit_raw = witness_raw_lines(witness, root, passive_only=passive_only)
    wit_label = "passive UART witness" if passive_only else "all packets.raw"
    invalid = 0
    for line, n in src_raw.items():
        raw = parse_raw_line(line)
        if raw is None or not validate_frame(raw)[0]:
            invalid += n
    sent = sum(src_raw.values())
    matched = sum(min(n, wit_raw.get(line, 0)) for line, n in src_raw.items())
    print(f"pos{origin} {label} -> pos{witness} {wit_label}")
    print(f"  sent={sent} unique={len(src_raw)} invalid={invalid}")
    print(f"  exact hex match: {matched}/{sent}")
    if passive_only:
        inj = sum(1 for r in load_events(witness, root) if r.get("ev") == "eth_to_uart_done")
        passive_n = sum(wit_raw.values())
        total = sum(load_raw(witness, root).values())
        print(
            f"  witness pos{witness} raw breakdown: passive={passive_n} inject-log={total - passive_n} (eth_to_uart={inj})"
        )
    if matched != sent:
        fail_flag[0] = True
        for line, n in src_raw.items():
            miss = n - wit_raw.get(line, 0)
            if miss > 0:
                print(f"    MISSING x{miss}: {line[:72]}...")
    print()


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("session", type=Path, help="Archive folder under sessions/")
    parser.add_argument("--origin", type=int, default=0, help="Origin pos (0 = all 2-4)")
    parser.add_argument("--witness", type=int, default=1, help="Witness pos (1-4)")
    parser.add_argument(
        "--native-only",
        action="store_true",
        help="Origin side: only lines tied to local packet_complete (machine origin)",
    )
    parser.add_argument(
        "--passive-only",
        action="store_true",
        help="Witness side: only passive UART packet_complete lines (exclude eth inject log)",
    )
    parser.add_argument(
        "--matrix",
        action="store_true",
        help="Print native-origin match grid for all origin/witness pairs",
    )
    args = parser.parse_args()
    root = args.session.resolve()
    fail = [False]

    ss = load_events(1, root)
    for r in ss:
        if r.get("ev") == "session_start":
            print(f"parser_build: {r.get('parser_build')}")
            print(f"log_eth_inject_to_packets_raw: {r.get('log_eth_inject_to_packets_raw')}")
            break
    print(f"Session: {root.name}\n")

    if args.matrix:
        wit_note = "passive UART witness" if args.passive_only else "all packets.raw (incl. inject log)"
        print(f"Native-origin hex match matrix (matched/sent) — witness: {wit_note}:\n")
        print("         witness ->")
        header = "origin   " + "".join(f"  pos{w}" for w in (1, 2, 3, 4))
        print(header)
        for origin in (2, 3, 4):
            cells = [f"pos{origin}   "]
            for witness in (1, 2, 3, 4):
                if witness == origin:
                    cells.append("    — ")
                    continue
                m, s, _ = compare_origin_to_witness(
                    origin,
                    witness,
                    root,
                    native_only=True,
                    passive_only=args.passive_only,
                )
                cells.append(f" {m:3d}/{s:<3d}" if s else "   0/0 ")
            print("".join(cells))
        print()
        for origin in (2, 3, 4):
            for witness in (1, 2, 3, 4):
                if witness == origin:
                    continue
                print_compare(
                    origin,
                    witness,
                    root,
                    native_only=True,
                    passive_only=args.passive_only,
                    fail_flag=fail,
                )
        return 1 if fail[0] else 0

    origins = [args.origin] if args.origin else [2, 3, 4]
    for src in origins:
        print_compare(
            src,
            args.witness,
            root,
            native_only=args.native_only,
            passive_only=args.passive_only,
            fail_flag=fail,
        )
    return 1 if fail[0] else 0


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