"""test_secp256k1_known_answers.py — bend form A validator.

Pipes BSCP-format scalar*G requests at our secp256k1-batch-mul worker
via --binary mode, parses the BSCR response, and compares each output
point against a host reference (coincurve, fallback ecdsa).

Wire format matches secp256k1-batch-mul.cu header doc:
  BSCP + u32 op_id + u32 n + base_x(BE,32) + base_y(BE,32) + scalars(BE,32*n)
  BSCR + u32 status + u32 n + points(BE,64*n  x||y)

Usage:
  python3 -u test_secp256k1_known_answers.py [worker-path]
  python3 -u test_secp256k1_known_answers.py ./secp256k1-batch-mul --n 1000
"""
import argparse
import os
import random
import struct
import subprocess
import sys
import tempfile
import time

# secp256k1 curve order and base point.
SECP256K1_N = 0xFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFEBAAEDCE6AF48A03BBFD25E8CD0364141
SECP256K1_GX = 0x79BE667EF9DCBBAC55A06295CE870B07029BFCDB2DCE28D959F2815B16F81798
SECP256K1_GY = 0x483ADA7726A3C4655DA4FBFC0E1108A8FD17B448A68554199C47D08FFB10D4B8

HAS_COINCURVE = False
HAS_ECDSA = False
try:
    from coincurve import PrivateKey, PublicKey
    HAS_COINCURVE = True
except ImportError:
    try:
        from ecdsa import SigningKey, SECP256k1, ellipticcurve
        from ecdsa.ecdsa import generator_secp256k1
        HAS_ECDSA = True
    except ImportError:
        pass

if not HAS_COINCURVE and not HAS_ECDSA:
    print("ERROR: install coincurve (pip install coincurve) or python-ecdsa")
    sys.exit(1)


def host_scalar_mul_G(k):
    """Return (x, y) of k*G on secp256k1, or (0, 0) for infinity."""
    k = k % SECP256K1_N
    if k == 0:
        return (0, 0)
    if HAS_COINCURVE:
        # coincurve's PrivateKey requires 1 <= k < n; encode public_key.format(False)
        # gives 0x04 + x32 + y32.
        priv = PrivateKey.from_int(k)
        pub = priv.public_key.format(compressed=False)
        assert pub[0] == 0x04
        x = int.from_bytes(pub[1:33], "big")
        y = int.from_bytes(pub[33:65], "big")
        return (x, y)
    else:
        pt = k * generator_secp256k1
        if pt == ellipticcurve.INFINITY:
            return (0, 0)
        return (pt.x(), pt.y())


def host_scalar_mul_point(k, px, py):
    """k * (px, py) on secp256k1. Used only when base != G."""
    if HAS_ECDSA:
        from ecdsa.ecdsa import generator_secp256k1
        # Build the affine point on the same curve.
        curve = generator_secp256k1.curve()
        pt = ellipticcurve.Point(curve, px, py)
        result = k * pt
        if result == ellipticcurve.INFINITY:
            return (0, 0)
        return (result.x(), result.y())
    # coincurve only exposes G-based ops; for non-G base fall through
    # to a pure-Python double-and-add.
    return _python_scalar_mul(k, px, py)


_SECP_P = 0xFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFEFFFFFC2F

def _python_scalar_mul(k, px, py):
    def inv(a, m): return pow(a, -1, m)
    def add(P, Q):
        if P is None: return Q
        if Q is None: return P
        if P[0] == Q[0]:
            if (P[1] + Q[1]) % _SECP_P == 0:
                return None
            lam = (3 * P[0] * P[0]) * inv(2 * P[1], _SECP_P) % _SECP_P
        else:
            lam = (Q[1] - P[1]) * inv(Q[0] - P[0], _SECP_P) % _SECP_P
        x3 = (lam * lam - P[0] - Q[0]) % _SECP_P
        y3 = (lam * (P[0] - x3) - P[1]) % _SECP_P
        return (x3, y3)
    R = None
    P = (px, py)
    while k:
        if k & 1: R = add(R, P)
        P = add(P, P)
        k >>= 1
    return (0, 0) if R is None else R


def int_to_be32(v):
    return int(v % (1 << 256)).to_bytes(32, "big")


def build_bscp(n, base_x, base_y, scalars):
    parts = [b"BSCP",
             (0x01).to_bytes(4, "little"),
             n.to_bytes(4, "little"),
             int_to_be32(base_x),
             int_to_be32(base_y)]
    for s in scalars:
        parts.append(int_to_be32(s))
    return b"".join(parts)


def parse_bscr(blob, n):
    if len(blob) < 12:
        raise RuntimeError(f"short BSCR: {len(blob)} bytes")
    magic = blob[:4]
    status = int.from_bytes(blob[4:8], "little")
    n_echo = int.from_bytes(blob[8:12], "little")
    if magic != b"BSCR":
        raise RuntimeError(f"bad magic {magic!r}, status={status}")
    if status != 0:
        raise RuntimeError(f"worker status={status} (nonzero)")
    if n_echo != n:
        raise RuntimeError(f"n mismatch: req={n} resp={n_echo}")
    expected = 12 + n * 64
    if len(blob) != expected:
        raise RuntimeError(f"BSCR size {len(blob)} != expected {expected}")
    out = []
    for i in range(n):
        off = 12 + i * 64
        x = int.from_bytes(blob[off:off+32], "big")
        y = int.from_bytes(blob[off+32:off+64], "big")
        out.append((x, y))
    return out


def run_worker(worker, blob):
    with tempfile.NamedTemporaryFile(delete=False, suffix=".bscp") as fin:
        fin.write(blob)
        in_path = fin.name
    out_path = in_path.replace(".bscp", ".bscr")
    try:
        r = subprocess.run([worker, "--binary", in_path, out_path],
                           capture_output=True, timeout=600)
        if r.returncode != 0:
            raise RuntimeError(f"worker exit {r.returncode}: {r.stderr[:400].decode(errors='replace')}")
        with open(out_path, "rb") as f:
            return f.read(), r.stderr.decode(errors="replace")
    finally:
        for p in (in_path, out_path):
            try: os.unlink(p)
            except OSError: pass


def test_known_small(worker):
    """Tiny fixed-scalar sanity. Includes 1, 2, 3, n-1, n (->infinity)."""
    scalars = [1, 2, 3, 7, 0xdeadbeef,
               SECP256K1_N - 1,
               SECP256K1_N,           # infinity
               (1 << 128) - 1]
    n = len(scalars)
    blob = build_bscp(n, SECP256K1_GX, SECP256K1_GY, scalars)
    out, _ = run_worker(worker, blob)
    pts = parse_bscr(out, n)
    fails = 0
    for i, k in enumerate(scalars):
        exp = host_scalar_mul_G(k)
        got = pts[i]
        if exp != got:
            fails += 1
            if fails <= 3:
                print(f"  FAIL  k=0x{k:x}")
                print(f"    exp x=0x{exp[0]:064x} y=0x{exp[1]:064x}")
                print(f"    got x=0x{got[0]:064x} y=0x{got[1]:064x}")
    label = "known-small"
    if fails == 0:
        print(f"  PASS  {label:25s}  n={n}")
        return True
    print(f"  FAIL  {label:25s}  n={n}  mismatches={fails}/{n}")
    return False


def test_random_n(worker, n, seed=0xC0FFEE):
    rng = random.Random(seed ^ n)
    scalars = [rng.randrange(1, SECP256K1_N) for _ in range(n)]
    blob = build_bscp(n, SECP256K1_GX, SECP256K1_GY, scalars)
    t0 = time.monotonic()
    out, stderr_text = run_worker(worker, blob)
    wall = time.monotonic() - t0
    pts = parse_bscr(out, n)
    fails = 0
    first = None
    t_h0 = time.monotonic()
    for i, k in enumerate(scalars):
        exp = host_scalar_mul_G(k)
        if pts[i] != exp:
            fails += 1
            if first is None:
                first = (i, k, exp, pts[i])
    host_ms = (time.monotonic() - t_h0) * 1000
    label = f"random n={n}"
    if fails == 0:
        thr = n / wall / 1e3
        print(f"  PASS  {label:25s}  n={n}  worker={wall*1000:.1f}ms host={host_ms:.0f}ms  kkeys/s={thr:.1f}")
        print(f"        {stderr_text.strip().splitlines()[-1] if stderr_text else ''}")
        return True
    i, k, exp, got = first
    print(f"  FAIL  {label:25s}  n={n}  mismatches={fails}/{n}")
    print(f"    first @ i={i}")
    print(f"      k   = 0x{k:064x}")
    print(f"      exp = ({exp[0]:#x}, {exp[1]:#x})")
    print(f"      got = ({got[0]:#x}, {got[1]:#x})")
    return False


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("worker", nargs="?", default="./secp256k1-batch-mul")
    ap.add_argument("--n", type=int, default=None,
                    help="single N override (default: sweep)")
    args = ap.parse_args()
    if not os.path.exists(args.worker):
        print(f"worker not found: {args.worker}")
        sys.exit(1)
    print(f"=== secp256k1-batch-mul validation ({args.worker}) ===")
    print(f"    reference: {'coincurve' if HAS_COINCURVE else 'python-ecdsa'}")
    ok = test_known_small(args.worker)
    ns = [args.n] if args.n else [32, 1000, 10000]
    for n in ns:
        ok = test_random_n(args.worker, n) and ok
    if ok:
        print("\nALL PASS")
        sys.exit(0)
    print("\nFAILED")
    sys.exit(1)


if __name__ == "__main__":
    main()
