#!/usr/bin/env python3
# ---------------------------------------------------------------------------
# mlkem768_reference.py
#
# Independent, dependency-free ML-KEM-768 (FIPS 203) ENCAPSULATION oracle,
# written directly from FIPS 203 and using only Python's standard-library
# hashlib (SHA3-256, SHA3-512, SHAKE128, SHAKE256). It shares no code with
# Bouncy Castle, which is what the 0x0AE6 precompile wraps, so agreement
# between the two is real cross-implementation evidence rather than a tautology.
#
# Purpose: the ML-KEM-768 ACVP vectors in ../vectors/mlkem768_acvp.txt were
# previously only checkable by running MlKemAcvpKat.java against a patched Besu
# build. This script makes them checkable from a clean checkout with no build
# step, so an outside researcher can reproduce the claim.
#
#   python mlkem768_reference.py ../vectors/mlkem768_acvp.txt [out.json]
#
# Vector line format: ek_hex|m_hex|c_hex|k_hex
# ML-KEM.Encaps_internal(ek, m) must reproduce BOTH the ciphertext c and the
# shared secret k. Checking c alone would not catch a wrong shared secret.
#
# NOTE ON SCOPE: 0x0AE6 is TESTNET-ONLY. It is NOT live on Aere mainnet 2800.
# ---------------------------------------------------------------------------
import hashlib
import json
import sys

N = 256
Q = 3329
K = 3       # ML-KEM-768
ETA1 = 2
ETA2 = 2
DU = 10
DV = 4

# zeta^BitRev7(i) mod q, the standard FIPS 203 NTT constant table.
ZETAS = [
    1, 1729, 2580, 3289, 2642, 630, 1897, 848, 1062, 1919, 193, 797, 2786,
    3260, 569, 1746, 296, 2447, 1339, 1476, 3046, 56, 2240, 1333, 1426, 2094,
    535, 2882, 2393, 2879, 1974, 821, 289, 331, 3253, 1756, 1197, 2304, 2277,
    2055, 650, 1977, 2513, 632, 2865, 33, 1320, 1915, 2319, 1435, 807, 452,
    1438, 2868, 1534, 2402, 2647, 2617, 1481, 648, 2474, 3110, 1227, 910, 17,
    2761, 583, 2649, 1637, 723, 2288, 1100, 1409, 2662, 3281, 233, 756, 2156,
    3015, 3050, 1703, 1651, 2789, 1789, 1847, 952, 1461, 2687, 939, 2308, 2437,
    2388, 733, 2337, 268, 641, 1584, 2298, 2037, 3220, 375, 2549, 2090, 1645,
    1063, 319, 2773, 757, 2099, 561, 2466, 2594, 2804, 1092, 403, 1026, 1143,
    2150, 2775, 886, 1722, 1212, 1874, 1029, 2110, 2935, 885, 2154,
]

# zeta^(2*BitRev7(i)+1) for the degree-2 base multiplications.
GAMMAS = [
    17, 3312, 2761, 568, 583, 2746, 2649, 680, 1637, 1692, 723, 2606, 2288,
    1041, 1100, 2229, 1409, 1920, 2662, 667, 3281, 48, 233, 3096, 756, 2573,
    2156, 1173, 3015, 314, 3050, 279, 1703, 1626, 1651, 1678, 2789, 540, 1789,
    1540, 1847, 1482, 952, 2377, 1461, 1868, 2687, 642, 939, 2390, 2308, 1021,
    2437, 892, 2388, 941, 733, 2596, 2337, 992, 268, 3061, 641, 2688, 1584,
    1745, 2298, 1031, 2037, 1292, 3220, 109, 375, 2954, 2549, 780, 2090, 1239,
    1645, 1684, 1063, 2266, 319, 3010, 2773, 556, 757, 2572, 2099, 1230, 561,
    2768, 2466, 863, 2594, 735, 2804, 525, 1092, 2237, 403, 2926, 1026, 2303,
    1143, 2186, 2150, 1179, 2775, 554, 886, 2443, 1722, 1607, 1212, 2117, 1874,
    1455, 1029, 2300, 2110, 1219, 2935, 394, 885, 2444, 2154, 1175,
]


# --- hashes (FIPS 203 section 4.1) -----------------------------------------
def _H(b):
    return hashlib.sha3_256(b).digest()


def _G(b):
    d = hashlib.sha3_512(b).digest()
    return d[:32], d[32:]


class _Squeezer:
    """Incremental reader over a SHAKE128 stream.

    hashlib's SHAKE objects have no incremental read, so squeeze a block and
    grow it if the rejection sampler exhausts it. digest(n) is a prefix of
    digest(n+m) for a fixed absorbed input, so growing is safe.
    """

    def __init__(self, xof):
        self._xof = xof
        self._buf = xof.digest(1024)
        self._pos = 0

    def read(self, n):
        if self._pos + n > len(self._buf):
            self._buf = self._xof.digest(len(self._buf) * 2)
        out = self._buf[self._pos:self._pos + n]
        self._pos += n
        return out


def _XOF(rho, i, j):
    return _Squeezer(hashlib.shake_128(rho + bytes([i, j])))


def _PRF(eta, s, b):
    return hashlib.shake_256(s + bytes([b])).digest(64 * eta)


# --- bit / byte conversion (Algorithms 3-6) --------------------------------
def _bytes_to_bits(bs):
    out = []
    for byte in bs:
        for i in range(8):
            out.append((byte >> i) & 1)
    return out


def _byte_encode(f, d):
    bits = []
    for x in f:
        for i in range(d):
            bits.append((x >> i) & 1)
    out = bytearray(len(bits) // 8)
    for i, bit in enumerate(bits):
        out[i // 8] |= bit << (i % 8)
    return bytes(out)


def _byte_decode(bs, d):
    bits = _bytes_to_bits(bs)
    m = Q if d == 12 else (1 << d)
    return [
        sum(bits[i * d + j] << j for j in range(d)) % m
        for i in range(len(bits) // d)
    ]


# --- sampling (Algorithms 7-8) ---------------------------------------------
def _sample_ntt(xof):
    a = []
    while len(a) < N:
        c = xof.read(3)
        d1 = c[0] + 256 * (c[1] % 16)
        d2 = (c[1] // 16) + 16 * c[2]
        if d1 < Q:
            a.append(d1)
        if d2 < Q and len(a) < N:
            a.append(d2)
    return a


def _sample_poly_cbd(bs, eta):
    bits = _bytes_to_bits(bs)
    f = []
    for i in range(N):
        x = sum(bits[2 * i * eta + j] for j in range(eta))
        y = sum(bits[2 * i * eta + eta + j] for j in range(eta))
        f.append((x - y) % Q)
    return f


# --- NTT (Algorithms 9-11) --------------------------------------------------
def _ntt(f):
    f = list(f)
    i = 1
    length = 128
    while length >= 2:
        start = 0
        while start < N:
            z = ZETAS[i]
            i += 1
            for j in range(start, start + length):
                t = (z * f[j + length]) % Q
                f[j + length] = (f[j] - t) % Q
                f[j] = (f[j] + t) % Q
            start += 2 * length
        length //= 2
    return f


def _ntt_inv(f):
    f = list(f)
    i = 127
    length = 2
    while length <= 128:
        start = 0
        while start < N:
            z = ZETAS[i]
            i -= 1
            for j in range(start, start + length):
                t = f[j]
                f[j] = (t + f[j + length]) % Q
                f[j + length] = (z * (f[j + length] - t)) % Q
            start += 2 * length
        length *= 2
    return [(x * 3303) % Q for x in f]


def _multiply_ntts(f, g):
    h = [0] * N
    for i in range(N // 2):
        a0, a1 = f[2 * i], f[2 * i + 1]
        b0, b1 = g[2 * i], g[2 * i + 1]
        gm = GAMMAS[i]
        h[2 * i] = (a0 * b0 + a1 * b1 % Q * gm) % Q
        h[2 * i + 1] = (a0 * b1 + a1 * b0) % Q
    return h


def _add(f, g):
    return [(a + b) % Q for a, b in zip(f, g)]


# --- compression (section 4.2.1) -------------------------------------------
def _compress(f, d):
    return [(((x << d) + Q // 2) // Q) % (1 << d) for x in f]


def _decompress(f, d):
    return [((x * Q) + (1 << (d - 1))) >> d for x in f]


# --- K-PKE.Encrypt (Algorithm 14) ------------------------------------------
def _kpke_encrypt(ek, m, r):
    t_hat = [_byte_decode(ek[384 * i:384 * (i + 1)], 12) for i in range(K)]
    rho = ek[384 * K:384 * K + 32]

    a_hat = [[_sample_ntt(_XOF(rho, i, j)) for j in range(K)] for i in range(K)]

    y = [_ntt(_sample_poly_cbd(_PRF(ETA1, r, i), ETA1)) for i in range(K)]
    e1 = [_sample_poly_cbd(_PRF(ETA2, r, K + i), ETA2) for i in range(K)]
    e2 = _sample_poly_cbd(_PRF(ETA2, r, 2 * K), ETA2)

    # u = NTT^-1(A^T . y) + e1
    u = []
    for i in range(K):
        acc = [0] * N
        for j in range(K):
            acc = _add(acc, _multiply_ntts(a_hat[j][i], y[j]))
        u.append(_add(_ntt_inv(acc), e1[i]))

    mu = _decompress(_byte_decode(m, 1), 1)
    acc = [0] * N
    for i in range(K):
        acc = _add(acc, _multiply_ntts(t_hat[i], y[i]))
    v = _add(_add(_ntt_inv(acc), e2), mu)

    c1 = b"".join(_byte_encode(_compress(u[i], DU), DU) for i in range(K))
    c2 = _byte_encode(_compress(v, DV), DV)
    return c1 + c2


def encaps_internal(ek, m):
    """ML-KEM.Encaps_internal (Algorithm 17). Returns (shared_secret_K, ciphertext_c)."""
    k_bar, r = _G(m + _H(ek))
    c = _kpke_encrypt(ek, m, r)
    return k_bar, c


def main():
    vectors = sys.argv[1] if len(sys.argv) > 1 else "../vectors/mlkem768_acvp.txt"
    outpath = sys.argv[2] if len(sys.argv) > 2 else "kat-results-mlkem-reference.json"

    passed = failed = total = 0
    results = []
    with open(vectors) as f:
        for raw in f:
            line = raw.strip()
            if not line or line.startswith("#"):
                continue
            p = line.split("|")
            if len(p) != 5:
                continue
            total += 1
            tc, ek_hex, m_hex, c_hex, k_hex = p
            ek = bytes.fromhex(ek_hex)
            m = bytes.fromhex(m_hex)
            c_want = bytes.fromhex(c_hex)
            k_want = bytes.fromhex(k_hex)
            k_got, c_got = encaps_internal(ek, m)
            ok = (k_got == k_want) and (c_got == c_want)
            if ok:
                passed += 1
            else:
                failed += 1
            results.append({
                "tcId": tc,
                "ekLen": len(ek),
                "ctLen": len(c_got),
                "ciphertextMatch": c_got == c_want,
                "sharedSecretMatch": k_got == k_want,
                "status": "PASS" if ok else "FAIL",
            })
            print(f"ML-KEM-768 tc{tc} ekLen={len(ek)} ctLen={len(c_got)} "
                  f"c={'OK' if c_got == c_want else 'MISMATCH'} "
                  f"k={'OK' if k_got == k_want else 'MISMATCH'} -> "
                  f"{'PASS' if ok else 'FAIL'}")

    report = {
        "precompile": "ML-KEM-768 (FIPS 203) deterministic encapsulate",
        "address": "0x0000000000000000000000000000000000000ae6",
        "liveOnMainnet2800": False,
        "note": "TESTNET-ONLY precompile. Verified here against an independent "
                "pure-Python FIPS 203 oracle (stdlib hashlib only), not Bouncy Castle.",
        "oracle": "mlkem768_reference.py (this file)",
        "passed": passed, "failed": failed, "total": total, "results": results,
    }
    with open(outpath, "w") as f:
        json.dump(report, f, indent=2)

    print(f"ML-KEM-768 KAT: passed={passed} failed={failed} total={total}")
    if failed != 0 or total == 0:
        print("MLKEM_KAT_FAILED")
        sys.exit(3)
    print("MLKEM_KAT_OK")


if __name__ == "__main__":
    main()