"""test_radix_sort_known_answers.py — validation for radix-sort.

Pipes BSRT-format u64 key arrays at the CUB radix-sort worker via
--binary mode, parses the BSRR response, & compares the byte stream
against Python's sorted() reference.

Coverage: op_id 0x01 (sort-u64-asc). Sweeps N across small, medium,
& large workloads so the test catches both edge cases (tiny N) and
sustained-throughput regressions (10 M keys).

Usage:
  python3 -u test_radix_sort_known_answers.py [worker]   # default ./radix-sort
  python3 -u test_radix_sort_known_answers.py ./radix-sort --bench
  python3 -u test_radix_sort_known_answers.py ./radix-sort --n 1000000
"""
import argparse
import os
import random
import re
import subprocess
import sys
import tempfile
import time

OP_SORT_U64_ASC = 0x01


def build_bsrt(op_id, keys):
    """Frame an u64 key list into a BSRT payload."""
    n = len(keys)
    parts = [
        b"BSRT",
        op_id.to_bytes(4, "little"),
        n.to_bytes(4, "little"),
    ]
    for k in keys:
        parts.append(int(k).to_bytes(8, "little"))
    return b"".join(parts)


def parse_bsrr(blob, n_expected):
    if len(blob) < 12 or blob[:4] != b"BSRR":
        raise RuntimeError(f"bad BSRR magic / short: {blob[:40]!r}")
    status = int.from_bytes(blob[4:8], "little")
    n      = int.from_bytes(blob[8:12], "little")
    if status != 0:
        raise RuntimeError(f"worker returned status {status}")
    if n != n_expected:
        raise RuntimeError(f"BSRR n {n} != expected {n_expected}")
    expected_len = 12 + n * 8
    if len(blob) != expected_len:
        raise RuntimeError(
            f"BSRR size {len(blob)} != expected {expected_len} (n={n})"
        )
    body = blob[12:]
    return [int.from_bytes(body[i*8:(i+1)*8], "little") for i in range(n)]


KERNEL_RE = re.compile(
    r"kernel=([\d.]+)\s+write=[\d.]+\s+ms\s+\(([\d.]+)\s+Gkeys/s\)"
)


def run_worker_binary(worker, in_blob):
    with tempfile.NamedTemporaryFile(delete=False, suffix=".bsrt") as fin:
        fin.write(in_blob)
        in_path = fin.name
    out_path = in_path.replace(".bsrt", ".bsrr")
    try:
        result = subprocess.run(
            [worker, "--binary", in_path, out_path],
            capture_output=True, timeout=300,
        )
        if result.returncode != 0:
            raise RuntimeError(
                f"worker exited {result.returncode}: {result.stderr.decode()[:400]}"
            )
        with open(out_path, "rb") as fout:
            return fout.read(), result.stderr.decode()
    finally:
        for p in (in_path, out_path):
            try:
                os.unlink(p)
            except OSError:
                pass


def parse_kernel_telemetry(stderr_text):
    """Extract (kernel_ms, gkeys_per_sec) from the worker's stderr line."""
    m = KERNEL_RE.search(stderr_text)
    if not m:
        return None, None
    return float(m.group(1)), float(m.group(2))


def test_n(worker, n, seed=0xC0FFEE57C0DEC0DE):
    """Sort N random u64s, compare byte-for-byte against Python sorted()."""
    rng = random.Random(seed ^ n)
    keys = [rng.randrange(0, 1 << 64) for _ in range(n)]
    expected = sorted(keys)
    in_blob = build_bsrt(OP_SORT_U64_ASC, keys)

    t_host = time.perf_counter()
    out_blob, stderr = run_worker_binary(worker, in_blob)
    wall_ms = (time.perf_counter() - t_host) * 1000.0
    got = parse_bsrr(out_blob, n)

    kernel_ms, gkeys = parse_kernel_telemetry(stderr)
    # Byte-identity check: rebuild canonical bytes from sorted() and
    # compare to the worker's body (skip the 12-byte BSRR header).
    expected_bytes = b"".join(int(k).to_bytes(8, "little") for k in expected)
    actual_bytes   = out_blob[12:]
    ok = (expected_bytes == actual_bytes)

    if not ok:
        # Locate first mismatch for the report.
        first = None
        for i in range(n):
            if got[i] != expected[i]:
                first = (i, expected[i], got[i])
                break
        kernel_label = f"{kernel_ms:.3f}" if kernel_ms is not None else "n/a"
        print(f"  FAIL  sort-u64-asc  n={n:<10}  wall={wall_ms:.1f}ms  kernel={kernel_label}ms")
        if first is not None:
            i, exp, gv = first
            print(f"    first mismatch @ i={i}: expected=0x{exp:016x} got=0x{gv:016x}")
        return False, kernel_ms, gkeys

    extra = ""
    if kernel_ms is not None and gkeys is not None:
        extra = f"  kernel={kernel_ms:.3f}ms  {gkeys:.3f} Gkeys/s"
    print(f"  PASS  sort-u64-asc  n={n:<10}  wall={wall_ms:.1f}ms{extra}")
    return True, kernel_ms, gkeys


DEFAULT_NS = [32, 1_000, 100_000, 1_000_000, 10_000_000]
BENCH_NS   = [100_000, 1_000_000, 10_000_000]


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("worker", nargs="?", default="./radix-sort")
    ap.add_argument("--n", type=int, default=None,
                    help="single N override (default: sweep)")
    ap.add_argument("--bench", action="store_true",
                    help="report kernel ms / Gkeys/s table at n=100k/1M/10M")
    args = ap.parse_args()

    if not os.path.exists(args.worker):
        print(f"worker not found: {args.worker}")
        sys.exit(1)

    if args.n:
        ns = [args.n]
    elif args.bench:
        ns = BENCH_NS
    else:
        ns = DEFAULT_NS

    all_pass = True
    bench_rows = []
    for n in ns:
        ok, kernel_ms, gkeys = test_n(args.worker, n)
        all_pass = all_pass and ok
        bench_rows.append((n, kernel_ms, gkeys))

    if args.bench:
        print("\nbench summary")
        print(f"  {'n':>10}  {'kernel_ms':>12}  {'Gkeys/s':>10}")
        for n, kernel_ms, gkeys in bench_rows:
            km = f"{kernel_ms:.3f}" if kernel_ms is not None else "n/a"
            gk = f"{gkeys:.3f}" if gkeys is not None else "n/a"
            print(f"  {n:>10}  {km:>12}  {gk:>10}")

    if all_pass:
        print("\nALL PASS")
        sys.exit(0)
    else:
        print("\nFAILED")
        sys.exit(1)


if __name__ == "__main__":
    main()
