"""test_blake3_known_answers.py — bend form `cuda-blake3-tree` validator.

Pipes BSB3-format requests at our blake3-fanout worker via --binary mode,
parses the BSR3 response, and compares each digest against an in-process
BLAKE3 reference implementation translated directly from the upstream
reference spec (CC0 / Apache-2.0):

    https://github.com/BLAKE3-team/BLAKE3/blob/master/reference_impl/reference_impl.rs

Reference covers chunk-tree walks for inputs of any length, including the
single-chunk fast path (<=1024 B) and multi-chunk Merkle paths. Both code
paths must match byte-for-byte.

Wire (matches header doc in blake3-fanout.cu):
    request:  "BSB3" | u32 out_bytes | u32 n_items | (u32 len + bytes) * n
    response: "BSR3" | u32 n_items   | u32 out_bytes | (out_bytes bytes) * n

Usage:
    python3 -u test_blake3_known_answers.py [worker-path]
    python3 -u test_blake3_known_answers.py ./blake3-fanout --n 1000
    python3 -u test_blake3_known_answers.py ./blake3-fanout --all  # all sizes

Exits 0 on PASS, 1 on FAIL.
"""
import argparse
import os
import random
import struct
import subprocess
import sys
import tempfile
import time


# ── BLAKE3 reference implementation (Python translation) ───────────────
# Translated 1:1 from BLAKE3-team/BLAKE3 reference_impl/reference_impl.rs.

OUT_LEN = 32
KEY_LEN = 32
BLOCK_LEN = 64
CHUNK_LEN = 1024

CHUNK_START = 1 << 0
CHUNK_END   = 1 << 1
PARENT      = 1 << 2
ROOT        = 1 << 3
KEYED_HASH  = 1 << 4

IV = [
    0x6A09E667, 0xBB67AE85, 0x3C6EF372, 0xA54FF53A,
    0x510E527F, 0x9B05688C, 0x1F83D9AB, 0x5BE0CD19,
]

MSG_PERMUTATION = [2, 6, 3, 10, 7, 0, 4, 13, 1, 11, 12, 5, 9, 14, 15, 8]

MASK32 = 0xFFFFFFFF

def _rotr(x, n):
    return ((x >> n) | (x << (32 - n))) & MASK32

def _g(state, a, b, c, d, mx, my):
    state[a] = (state[a] + state[b] + mx) & MASK32
    state[d] = _rotr(state[d] ^ state[a], 16)
    state[c] = (state[c] + state[d]) & MASK32
    state[b] = _rotr(state[b] ^ state[c], 12)
    state[a] = (state[a] + state[b] + my) & MASK32
    state[d] = _rotr(state[d] ^ state[a], 8)
    state[c] = (state[c] + state[d]) & MASK32
    state[b] = _rotr(state[b] ^ state[c], 7)

def _round(state, m):
    _g(state, 0, 4,  8, 12, m[0],  m[1])
    _g(state, 1, 5,  9, 13, m[2],  m[3])
    _g(state, 2, 6, 10, 14, m[4],  m[5])
    _g(state, 3, 7, 11, 15, m[6],  m[7])
    _g(state, 0, 5, 10, 15, m[8],  m[9])
    _g(state, 1, 6, 11, 12, m[10], m[11])
    _g(state, 2, 7,  8, 13, m[12], m[13])
    _g(state, 3, 4,  9, 14, m[14], m[15])

def _permute(m):
    p = [m[MSG_PERMUTATION[i]] for i in range(16)]
    for i in range(16):
        m[i] = p[i]

def _compress(cv, block_words, counter, block_len, flags):
    state = [
        cv[0], cv[1], cv[2], cv[3],
        cv[4], cv[5], cv[6], cv[7],
        IV[0], IV[1], IV[2], IV[3],
        counter & MASK32, (counter >> 32) & MASK32,
        block_len, flags,
    ]
    block = list(block_words)
    for _ in range(6):
        _round(state, block)
        _permute(block)
    _round(state, block)
    for i in range(8):
        state[i]   ^= state[i + 8]
        state[i+8] ^= cv[i]
    return state

def _words_from_le(byts):
    n = len(byts)
    pad = (-n) % 4
    if pad:
        byts = bytes(byts) + b'\x00' * pad
    return [int.from_bytes(byts[4*i:4*i+4], 'little') for i in range((n + pad) // 4)]

def _le_bytes_from_words(words, n_bytes):
    out = bytearray()
    for w in words:
        out.extend(int(w & MASK32).to_bytes(4, 'little'))
    return bytes(out[:n_bytes])

class _Output:
    def __init__(self, input_cv, block_words, counter, block_len, flags):
        self.input_cv = list(input_cv)
        self.block_words = list(block_words)
        self.counter = counter
        self.block_len = block_len
        self.flags = flags
    def chaining_value(self):
        return _compress(self.input_cv, self.block_words, self.counter,
                         self.block_len, self.flags)[:8]
    def root_output_bytes(self, n_bytes):
        out = bytearray()
        block_ctr = 0
        while len(out) < n_bytes:
            words = _compress(self.input_cv, self.block_words, block_ctr,
                              self.block_len, self.flags | ROOT)
            take = min(2 * OUT_LEN, n_bytes - len(out))
            out.extend(_le_bytes_from_words(words, take))
            block_ctr += 1
        return bytes(out[:n_bytes])

class _ChunkState:
    def __init__(self, key, counter, flags):
        self.cv = list(key)
        self.counter = counter
        self.buf = bytearray(BLOCK_LEN)
        self.buf_len = 0
        self.blocks_compressed = 0
        self.flags = flags
    def len(self):
        return BLOCK_LEN * self.blocks_compressed + self.buf_len
    def start_flag(self):
        return CHUNK_START if self.blocks_compressed == 0 else 0
    def update(self, data):
        i = 0
        while i < len(data):
            if self.buf_len == BLOCK_LEN:
                bw = _words_from_le(bytes(self.buf))
                self.cv = _compress(self.cv, bw, self.counter, BLOCK_LEN,
                                    self.flags | self.start_flag())[:8]
                self.blocks_compressed += 1
                self.buf = bytearray(BLOCK_LEN)
                self.buf_len = 0
            take = min(BLOCK_LEN - self.buf_len, len(data) - i)
            self.buf[self.buf_len:self.buf_len+take] = data[i:i+take]
            self.buf_len += take
            i += take
    def output(self):
        bw = _words_from_le(bytes(self.buf))
        return _Output(self.cv, bw, self.counter, self.buf_len,
                       self.flags | self.start_flag() | CHUNK_END)

def _parent_output(left_cv, right_cv, key, flags):
    bw = list(left_cv) + list(right_cv)
    return _Output(list(key), bw, 0, BLOCK_LEN, PARENT | flags)

def _parent_cv(left_cv, right_cv, key, flags):
    return _parent_output(left_cv, right_cv, key, flags).chaining_value()

class _Hasher:
    def __init__(self, key=None, flags=0):
        if key is None: key = IV
        self.chunk_state = _ChunkState(key, 0, flags)
        self.key = list(key)
        self.cv_stack = []
        self.flags = flags
    def _add_chunk_cv(self, new_cv, total_chunks):
        while (total_chunks & 1) == 0:
            new_cv = _parent_cv(self.cv_stack.pop(), new_cv, self.key, self.flags)
            total_chunks >>= 1
        self.cv_stack.append(new_cv)
    def update(self, data):
        i = 0
        while i < len(data):
            if self.chunk_state.len() == CHUNK_LEN:
                cv = self.chunk_state.output().chaining_value()
                total = self.chunk_state.counter + 1
                self._add_chunk_cv(cv, total)
                self.chunk_state = _ChunkState(self.key, total, self.flags)
            want = CHUNK_LEN - self.chunk_state.len()
            take = min(want, len(data) - i)
            self.chunk_state.update(data[i:i+take])
            i += take
    def finalize(self, n_bytes=32):
        out = self.chunk_state.output()
        remaining = list(self.cv_stack)
        while remaining:
            cv = remaining.pop()
            out = _parent_output(cv, out.chaining_value(), self.key, self.flags)
        return out.root_output_bytes(n_bytes)

def blake3_ref(data, out_bytes=32):
    h = _Hasher()
    h.update(data)
    return h.finalize(out_bytes)


# ── wire protocol drivers ─────────────────────────────────────────────

def encode_bsb3(inputs, out_bytes=32):
    buf = bytearray()
    buf.extend(b"BSB3")
    buf.extend(struct.pack("<II", out_bytes, len(inputs)))
    for it in inputs:
        buf.extend(struct.pack("<I", len(it)))
        buf.extend(it)
    return bytes(buf)

def decode_bsr3(blob, expected_n, expected_out_bytes):
    if len(blob) < 12 or blob[:4] != b"BSR3":
        raise ValueError(f"bad magic: {blob[:4]!r}")
    n, ob = struct.unpack("<II", blob[4:12])
    if n != expected_n:
        raise ValueError(f"n mismatch: {n} != {expected_n}")
    if ob != expected_out_bytes:
        raise ValueError(f"out_bytes mismatch: {ob} != {expected_out_bytes}")
    body = blob[12:]
    if len(body) != n * ob:
        raise ValueError(f"body size: {len(body)} != {n * ob}")
    return [bytes(body[i*ob:(i+1)*ob]) for i in range(n)]


def run_worker(worker, inputs, out_bytes=32):
    req = encode_bsb3(inputs, out_bytes)
    with tempfile.NamedTemporaryFile(delete=False, suffix=".bin") as fin:
        fin.write(req)
        in_path = fin.name
    out_path = in_path + ".out"
    t0 = time.time()
    try:
        r = subprocess.run([worker, "--binary", in_path, out_path],
                           check=False, capture_output=True)
        if r.returncode != 0:
            print(f"  worker stderr: {r.stderr.decode(errors='replace')}")
            raise RuntimeError(f"worker exit {r.returncode}")
        with open(out_path, "rb") as f:
            resp = f.read()
        # extract kernel time from stderr ("kernel=X")
        stderr = r.stderr.decode(errors='replace')
        kernel_ms = None
        for tok in stderr.split():
            if tok.startswith("kernel="):
                try: kernel_ms = float(tok.split("=", 1)[1])
                except ValueError: pass
        return decode_bsr3(resp, len(inputs), out_bytes), time.time() - t0, kernel_ms
    finally:
        for p in (in_path, out_path):
            if os.path.exists(p):
                os.unlink(p)


def gen_inputs(n, mode="mixed", seed=1):
    rnd = random.Random(seed)
    ins = []
    if mode == "small":
        for _ in range(n):
            L = rnd.randint(0, 64)
            ins.append(rnd.randbytes(L))
    elif mode == "single-chunk":
        for _ in range(n):
            L = rnd.randint(0, 1024)
            ins.append(rnd.randbytes(L))
    elif mode == "multi-chunk":
        for _ in range(n):
            L = rnd.randint(1025, 8192)
            ins.append(rnd.randbytes(L))
    elif mode == "fixed-64":
        for _ in range(n):
            ins.append(rnd.randbytes(64))
        # cover edge cases at the front
        ins[0:5] = [b"", b"\x00", b"\x00\x01", b"abc", b"a" * 63]
        ins = ins[:n]
    else:  # mixed: cover both single-chunk and multi-chunk
        for i in range(n):
            if i % 3 == 0:
                L = rnd.randint(0, 1024)
            elif i % 3 == 1:
                L = rnd.randint(1025, 4096)
            else:
                L = rnd.randint(65, 512)
            ins.append(rnd.randbytes(L))
    return ins


def run_case(worker, n, mode, max_show_diff=4):
    print(f"  n={n:>7d} mode={mode}")
    ins = gen_inputs(n, mode=mode)
    expected = [blake3_ref(x, 32) for x in ins]
    got, wall, kernel_ms = run_worker(worker, ins, 32)
    total_in = sum(len(x) for x in ins)
    mb_s = (total_in / (kernel_ms / 1e3)) / 1e6 if kernel_ms else float('nan')
    fail = []
    for i, (a, b) in enumerate(zip(expected, got)):
        if a != b:
            fail.append((i, a, b))
    if fail:
        print(f"    FAIL: {len(fail)} mismatches / {n} (kernel {kernel_ms} ms, "
              f"wall {wall*1000:.1f} ms, {total_in/1e6:.2f} MB)")
        for i, a, b in fail[:max_show_diff]:
            print(f"      input[{i}] len={len(ins[i])}")
            print(f"        ref: {a.hex()}")
            print(f"        got: {b.hex()}")
            print(f"        in:  {ins[i][:32].hex()}...")
        return False
    print(f"    OK   kernel={kernel_ms} ms  wall={wall*1000:.1f} ms  "
          f"{total_in/1e6:.2f} MB  {mb_s:.1f} MB/s")
    return True


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("worker", nargs="?", default="./blake3-fanout")
    ap.add_argument("--n", type=int, default=None,
                    help="single n; default = run [32, 1000, 10000, 100000, 1000000]")
    ap.add_argument("--mode", default="mixed",
                    choices=["mixed", "small", "single-chunk", "multi-chunk", "fixed-64"])
    ap.add_argument("--quick", action="store_true",
                    help="just n=32 mixed + n=1000 multi-chunk")
    args = ap.parse_args()

    if not os.path.exists(args.worker):
        print(f"worker not found: {args.worker}", file=sys.stderr); sys.exit(2)

    if args.n is not None:
        ok = run_case(args.worker, args.n, args.mode)
        sys.exit(0 if ok else 1)
    if args.quick:
        ok = True
        ok &= run_case(args.worker, 32,   "mixed")
        ok &= run_case(args.worker, 1000, "multi-chunk")
        sys.exit(0 if ok else 1)

    # full byte-identity sweep across canonical sizes & modes
    cases = [
        (32,      "mixed"),
        (32,      "small"),
        (32,      "single-chunk"),
        (32,      "multi-chunk"),
        (1000,    "mixed"),
        (10000,   "mixed"),
        (100000,  "fixed-64"),
        (1000000, "fixed-64"),
    ]
    ok = True
    for n, mode in cases:
        ok &= run_case(args.worker, n, mode)
    print("PASS" if ok else "FAIL")
    sys.exit(0 if ok else 1)


if __name__ == "__main__":
    main()
