"""test_cgbn_known_answers.py — validation for cgbn-batch-worker.

Pipes BCGB-format ops at our CGBN worker via --binary mode, parses the
BCGR response, and compares result bytes against a host bignum
reference. Falls back to pure-Python if gmpy2 is unavailable.

Op coverage (all 9 ops in plan §2 dispatcher table):
  0x01 mod-add        a + b mod m
  0x02 mod-sub        a - b mod m
  0x03 mod-mul        a * b mod m
  0x04 mod-sqr        a^2 mod m
  0x05 mod-inv        a^-1 mod m
  0x06 mod-exp        a^b mod m
  0x07 mod-reduce     a mod m
  0x08 add-no-mod     a + b (truncated to 256 bits)
  0x09 mul-no-mod     a * b (full 512-bit product)
  0x0A mod-inv-batch  Montgomery's batched-inversion trick (N inverses)

Usage:
  python3 test_cgbn_known_answers.py [worker-path]   ; default: ./cgbn-batch-worker
  python3 test_cgbn_known_answers.py ./cgbn-batch-worker --n 1000
  python3 test_cgbn_known_answers.py ./cgbn-batch-worker --op mod-mul
"""
import argparse
import os
import random
import subprocess
import sys
import tempfile
import time

SECP256K1_P = 0xFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFEFFFFFC2F
MASK_256 = (1 << 256) - 1

# Op codes — match cgbn-batch-worker.cu OP_* enum.
OP_MOD_ADD    = 0x01
OP_MOD_SUB    = 0x02
OP_MOD_MUL    = 0x03
OP_MOD_SQR    = 0x04
OP_MOD_INV    = 0x05
OP_MOD_EXP    = 0x06
OP_MOD_REDUCE = 0x07
OP_ADD_NO_MOD = 0x08
OP_MUL_NO_MOD = 0x09
OP_MOD_INV_BATCH = 0x0A

try:
    from gmpy2 import mpz, f_mod, invert, powmod
    def h_mul(a, b, m): return int(f_mod(mpz(a) * mpz(b), mpz(m)))
    def h_add(a, b, m): return int(f_mod(mpz(a) + mpz(b), mpz(m)))
    def h_sub(a, b, m): return int(f_mod(mpz(a) - mpz(b), mpz(m)))
    def h_sqr(a, m):    return int(f_mod(mpz(a) * mpz(a), mpz(m)))
    def h_inv(a, m):    return int(invert(mpz(a), mpz(m)))
    def h_exp(a, b, m): return int(powmod(mpz(a), mpz(b), mpz(m)))
    def h_rem(a, m):    return int(f_mod(mpz(a), mpz(m)))
except ImportError:
    def h_mul(a, b, m): return (a * b) % m
    def h_add(a, b, m): return (a + b) % m
    def h_sub(a, b, m): return (a - b) % m
    def h_sqr(a, m):    return (a * a) % m
    def h_inv(a, m):    return pow(a, -1, m)
    def h_exp(a, b, m): return pow(a, b, m)
    def h_rem(a, m):    return a % m


def int_to_le_bytes(x, nbytes=32):
    return x.to_bytes(nbytes, "little")


def le_bytes_to_int(b):
    return int.from_bytes(b, "little")


def build_bcgb(op_id, modulus, a_list, b_list, bitwidth=256):
    """Build BCGB payload. modulus=None for no-mod ops; b_list=None for unary."""
    bpi = bitwidth // 8
    n = len(a_list)
    if b_list is not None:
        assert len(b_list) == n
    parts = [
        b"BCGB",
        op_id.to_bytes(4, "little"),
        bitwidth.to_bytes(4, "little"),
        n.to_bytes(4, "little"),
    ]
    if modulus is not None:
        parts.append(int_to_le_bytes(modulus, bpi))
    for a in a_list:
        parts.append(int_to_le_bytes(a, bpi))
    if b_list is not None:
        for b in b_list:
            parts.append(int_to_le_bytes(b, bpi))
    return b"".join(parts)


def parse_bcgr(blob, n_expected, out_bytes_per_inst=32):
    if len(blob) < 12 or blob[:4] != b"BCGR":
        raise RuntimeError(f"bad BCGR 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"BCGR n {n} != expected {n_expected}")
    expected_len = 12 + n * out_bytes_per_inst
    if len(blob) != expected_len:
        raise RuntimeError(
            f"BCGR size {len(blob)} != expected {expected_len} for n={n} bpi={out_bytes_per_inst}"
        )
    out = []
    for i in range(n):
        out.append(le_bytes_to_int(
            blob[12 + i * out_bytes_per_inst : 12 + (i + 1) * out_bytes_per_inst]
        ))
    return out


def run_worker_binary(worker, in_blob):
    """Spawn worker --binary in_path out_path; return (out_blob, stderr_str)."""
    with tempfile.NamedTemporaryFile(delete=False, suffix=".bcgb") as fin:
        fin.write(in_blob)
        in_path = fin.name
    out_path = in_path.replace(".bcgb", ".bcgr")
    try:
        result = subprocess.run(
            [worker, "--binary", in_path, out_path],
            capture_output=True, timeout=120,
        )
        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 report_pass_fail(label, n, fail, first_fail, modulus):
    if fail == 0:
        print(f"  PASS  {label:30s}  n={n}")
        return True
    i, a, b, exp, got = first_fail
    print(f"  FAIL  {label:30s}  n={n}  mismatches={fail}/{n}")
    print(f"    first @ i={i}")
    print(f"      a   = 0x{a:064x}")
    if b is not None:
        print(f"      b   = 0x{b:064x}")
    if modulus is not None:
        print(f"      m   = 0x{modulus:064x}")
    print(f"      exp = 0x{exp:x}")
    print(f"      got = 0x{got:x}")
    return False


def test_binary_mod(worker, op_id, op_name, host_fn, n, modulus, label):
    rng = random.Random(0xC0FFEE57C0DEC0DE ^ n ^ op_id)
    a_list = [rng.randrange(0, modulus) for _ in range(n)]
    b_list = [rng.randrange(0, modulus) for _ in range(n)]
    in_blob = build_bcgb(op_id, modulus, a_list, b_list)
    out_blob, stderr = run_worker_binary(worker, in_blob)
    gpu = parse_bcgr(out_blob, n, out_bytes_per_inst=32)
    fail = 0; first_fail = None
    for i in range(n):
        expected = host_fn(a_list[i], b_list[i], modulus)
        if gpu[i] != expected:
            fail += 1
            if first_fail is None:
                first_fail = (i, a_list[i], b_list[i], expected, gpu[i])
    return report_pass_fail(f"{op_name} {label}", n, fail, first_fail, modulus)


def test_unary_mod(worker, op_id, op_name, host_fn, n, modulus, label,
                   ensure_invertible=False):
    rng = random.Random(0xC0FFEE57C0DEC0DE ^ n ^ op_id)
    if ensure_invertible:
        # secp256k1_p and our other moduli are prime, so any nonzero a < m
        # is invertible. Still pick from [1, m) to be safe.
        a_list = [rng.randrange(1, modulus) for _ in range(n)]
    else:
        # For mod-reduce we also feed values that exceed m so the reduce
        # is non-trivial. We cap at 2^256-1 because the wire is 256-bit.
        if op_id == OP_MOD_REDUCE:
            a_list = [rng.randrange(0, MASK_256 + 1) for _ in range(n)]
        else:
            a_list = [rng.randrange(0, modulus) for _ in range(n)]
    in_blob = build_bcgb(op_id, modulus, a_list, None)
    out_blob, stderr = run_worker_binary(worker, in_blob)
    gpu = parse_bcgr(out_blob, n, out_bytes_per_inst=32)
    fail = 0; first_fail = None
    for i in range(n):
        expected = host_fn(a_list[i], modulus)
        if gpu[i] != expected:
            fail += 1
            if first_fail is None:
                first_fail = (i, a_list[i], None, expected, gpu[i])
    return report_pass_fail(f"{op_name} {label}", n, fail, first_fail, modulus)


def test_mod_exp_small(worker, n, modulus, label):
    """mod-exp: keep b in [0, 2^16) so the binary ladder finishes fast
    enough at n=10k without blocking the whole suite for minutes."""
    rng = random.Random(0xC0FFEE57C0DEC0DE ^ n ^ OP_MOD_EXP)
    a_list = [rng.randrange(0, modulus) for _ in range(n)]
    b_list = [rng.randrange(0, 1 << 16) for _ in range(n)]
    in_blob = build_bcgb(OP_MOD_EXP, modulus, a_list, b_list)
    out_blob, stderr = run_worker_binary(worker, in_blob)
    gpu = parse_bcgr(out_blob, n, out_bytes_per_inst=32)
    fail = 0; first_fail = None
    for i in range(n):
        expected = h_exp(a_list[i], b_list[i], modulus)
        if gpu[i] != expected:
            fail += 1
            if first_fail is None:
                first_fail = (i, a_list[i], b_list[i], expected, gpu[i])
    return report_pass_fail(f"mod-exp {label}", n, fail, first_fail, modulus)


def test_add_no_mod(worker, n, label):
    rng = random.Random(0xC0FFEE57C0DEC0DE ^ n ^ OP_ADD_NO_MOD)
    a_list = [rng.randrange(0, MASK_256 + 1) for _ in range(n)]
    b_list = [rng.randrange(0, MASK_256 + 1) for _ in range(n)]
    in_blob = build_bcgb(OP_ADD_NO_MOD, None, a_list, b_list)
    out_blob, stderr = run_worker_binary(worker, in_blob)
    gpu = parse_bcgr(out_blob, n, out_bytes_per_inst=32)
    fail = 0; first_fail = None
    for i in range(n):
        expected = (a_list[i] + b_list[i]) & MASK_256
        if gpu[i] != expected:
            fail += 1
            if first_fail is None:
                first_fail = (i, a_list[i], b_list[i], expected, gpu[i])
    return report_pass_fail(f"add-no-mod {label}", n, fail, first_fail, None)


def test_mul_no_mod(worker, n, label):
    rng = random.Random(0xC0FFEE57C0DEC0DE ^ n ^ OP_MUL_NO_MOD)
    a_list = [rng.randrange(0, MASK_256 + 1) for _ in range(n)]
    b_list = [rng.randrange(0, MASK_256 + 1) for _ in range(n)]
    in_blob = build_bcgb(OP_MUL_NO_MOD, None, a_list, b_list)
    out_blob, stderr = run_worker_binary(worker, in_blob)
    # 64 bytes per instance (512-bit product, full).
    gpu = parse_bcgr(out_blob, n, out_bytes_per_inst=64)
    fail = 0; first_fail = None
    for i in range(n):
        expected = a_list[i] * b_list[i]
        if gpu[i] != expected:
            fail += 1
            if first_fail is None:
                first_fail = (i, a_list[i], b_list[i], expected, gpu[i])
    return report_pass_fail(f"mul-no-mod {label}", n, fail, first_fail, None)


# Op driver registry: name -> callable(worker, n, modulus, label).
def driver_mod_add(worker, n, modulus, label):
    return test_binary_mod(worker, OP_MOD_ADD, "mod-add", h_add, n, modulus, label)

def driver_mod_sub(worker, n, modulus, label):
    return test_binary_mod(worker, OP_MOD_SUB, "mod-sub", h_sub, n, modulus, label)

def driver_mod_mul(worker, n, modulus, label):
    return test_binary_mod(worker, OP_MOD_MUL, "mod-mul", h_mul, n, modulus, label)

def driver_mod_sqr(worker, n, modulus, label):
    return test_unary_mod(worker, OP_MOD_SQR, "mod-sqr", h_sqr, n, modulus, label)

def driver_mod_inv(worker, n, modulus, label):
    return test_unary_mod(worker, OP_MOD_INV, "mod-inv", h_inv, n, modulus, label,
                          ensure_invertible=True)

def driver_mod_inv_batch(worker, n, modulus, label):
    """Montgomery's batched inversion (op 0x0A). Same wire shape as 0x05
    but result[i] must still equal x[i]^-1 mod m byte-for-byte."""
    return test_unary_mod(worker, OP_MOD_INV_BATCH, "mod-inv-batch", h_inv,
                          n, modulus, label, ensure_invertible=True)

def driver_mod_exp(worker, n, modulus, label):
    return test_mod_exp_small(worker, n, modulus, label)

def driver_mod_reduce(worker, n, modulus, label):
    return test_unary_mod(worker, OP_MOD_REDUCE, "mod-reduce", h_rem, n, modulus, label)

def driver_add_no_mod(worker, n, modulus, label):
    return test_add_no_mod(worker, n, label)

def driver_mul_no_mod(worker, n, modulus, label):
    return test_mul_no_mod(worker, n, label)


OPS = [
    ("mod-add",       driver_mod_add,       True),
    ("mod-sub",       driver_mod_sub,       True),
    ("mod-mul",       driver_mod_mul,       True),
    ("mod-sqr",       driver_mod_sqr,       True),
    ("mod-inv",       driver_mod_inv,       True),
    ("mod-exp",       driver_mod_exp,       True),
    ("mod-reduce",    driver_mod_reduce,    True),
    ("add-no-mod",    driver_add_no_mod,    False),
    ("mul-no-mod",    driver_mul_no_mod,    False),
    ("mod-inv-batch", driver_mod_inv_batch, True),
]


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("worker", nargs="?", default="./cgbn-batch-worker")
    ap.add_argument("--n", type=int, default=None,
                    help="single N override (default: sweep 32, 1k, 10k, 100k)")
    ap.add_argument("--op", default=None,
                    help="restrict to one op name (e.g. mod-mul)")
    args = ap.parse_args()
    if not os.path.exists(args.worker):
        print(f"worker not found: {args.worker}")
        sys.exit(1)

    ns = [args.n] if args.n else [32, 1_000, 10_000, 100_000]
    moduli = [
        (SECP256K1_P, "secp256k1 prime"),
    ]
    selected = [(name, fn, uses_mod) for name, fn, uses_mod in OPS
                if args.op is None or args.op == name]
    if not selected:
        print(f"no op matched --op {args.op}; valid: {[n for n,_,_ in OPS]}")
        sys.exit(1)

    all_pass = True
    for n in ns:
        for op_name, fn, uses_mod in selected:
            if uses_mod:
                for m, m_name in moduli:
                    ok = fn(args.worker, n, m, m_name)
                    all_pass = all_pass and ok
            else:
                ok = fn(args.worker, n, None, "")
                all_pass = all_pass and ok

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


if __name__ == "__main__":
    main()
