#!/usr/bin/env python3
# One-shot kernel-call sequence diff: canary (oracle) vs ours (xenia-rs).
# Usage: python3 diff.py [--max N] [--window W] [--noisy] [--from K]
# Both canary.log and ours.log are read from this directory.

import argparse
import re
import sys
from pathlib import Path

HERE = Path(__file__).parent
CANARY = HERE / "canary.log"
OURS = HERE / "ours.log"

# canary D-level lines:  d> 1234ABCD ExportName(args...
#   (sometimes the handle is 8 hex digits; the export name is then `[A-Z][A-Za-z0-9_]+`).
# Some lines are F-level (file ops) or i-level — those aren't kernel exports we care about.
CANARY_RE = re.compile(r"^d>\s+[0-9A-Fa-f]+\s+([A-Z][A-Za-z0-9_]+)\(")

# our probe_calls lines:  ... probe_calls: hw=0 call=ExportName r3=...
OURS_RE = re.compile(r"probe_calls.*?call=([A-Za-z_][A-Za-z0-9_]*)")

# Noisy "runtime / CRT" exports we strip when --noisy is OFF (default).
# Heuristic: anything called >100x in canary's intro is noise.
NOISY_PREFIXES = (
    "Rtl", "Mm", "Ke",  # runtime/critical-section/memory/kernel-internal hot paths
)
# But keep these — they mark important structural events:
KEEP_NOISY = {
    "KeSetEvent", "KeResetEvent", "KePulseEvent", "KeReleaseSemaphore",
    "KeWaitForSingleObject", "KeWaitForMultipleObjects",
    "RtlRaiseException",
}


def extract(path: Path, regex: re.Pattern) -> list[str]:
    seq: list[str] = []
    with path.open("r", errors="replace") as f:
        for line in f:
            m = regex.search(line)
            if m:
                seq.append(m.group(1))
    return seq


def filter_noisy(seq: list[str], drop_noisy: bool) -> list[str]:
    if not drop_noisy:
        return seq
    out = []
    for n in seq:
        if n in KEEP_NOISY:
            out.append(n)
        elif any(n.startswith(p) for p in NOISY_PREFIXES):
            continue
        else:
            out.append(n)
    return out


def find_first_divergence(a: list[str], b: list[str]) -> int:
    n = min(len(a), len(b))
    for i in range(n):
        if a[i] != b[i]:
            return i
    return n  # one is a prefix of the other


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--max", type=int, default=2000,
                    help="cap each sequence to N entries before diffing (default 2000)")
    ap.add_argument("--window", type=int, default=20,
                    help="lines of context around the divergence point (default 20)")
    ap.add_argument("--noisy", action="store_true",
                    help="don't drop Rtl/Mm/Ke runtime/CRT calls")
    ap.add_argument("--from", dest="skip", type=int, default=0,
                    help="skip first K matched calls in BOTH sequences before diffing")
    args = ap.parse_args()

    canary = extract(CANARY, CANARY_RE)
    ours = extract(OURS, OURS_RE)
    print(f"raw: canary={len(canary)} ours={len(ours)}")

    canary_f = filter_noisy(canary, drop_noisy=not args.noisy)
    ours_f = filter_noisy(ours, drop_noisy=not args.noisy)
    print(f"filtered ({'noisy kept' if args.noisy else 'CRT/Rtl/Mm dropped'}): "
          f"canary={len(canary_f)} ours={len(ours_f)}")

    if args.skip:
        canary_f = canary_f[args.skip:]
        ours_f = ours_f[args.skip:]
        print(f"after --from={args.skip}: canary={len(canary_f)} ours={len(ours_f)}")

    canary_c = canary_f[:args.max]
    ours_c = ours_f[:args.max]

    div = find_first_divergence(canary_c, ours_c)
    if div == min(len(canary_c), len(ours_c)) and canary_c == ours_c[:len(canary_c)]:
        print(f"\nNo divergence within first {args.max} matched calls "
              f"(canary is a prefix of ours, or sequences are equal up to cap).")
        return 0

    print(f"\nFIRST DIVERGENCE at index {div + args.skip} "
          f"(post-noise-filter, post-skip).")

    lo = max(0, div - args.window)
    hi = min(min(len(canary_c), len(ours_c)) + 1, div + args.window + 1)

    print(f"\n  idx |  canary (oracle)               |  ours (xenia-rs)")
    print(  "  ----+--------------------------------+---------------------------------")
    for i in range(lo, hi):
        c = canary_c[i] if i < len(canary_c) else "<END>"
        o = ours_c[i] if i < len(ours_c) else "<END>"
        marker = "*" if i == div else " "
        print(f"  {marker}{i + args.skip:4d}|  {c:<31}|  {o}")

    # Also show sample of what canary has next that ours doesn't.
    print(f"\nNext 30 canary calls past divergence:")
    for i, n in enumerate(canary_c[div:div + 30]):
        print(f"  +{i:2d}: {n}")
    print(f"\nNext 30 ours calls past divergence:")
    for i, n in enumerate(ours_c[div:div + 30]):
        print(f"  +{i:2d}: {n}")

    return 0


if __name__ == "__main__":
    sys.exit(main())