#!/usr/bin/env python3
# xcheck_mlkemr.py - Cross-validate the hardware ML-KEM KeyGen KAT vectors
# against the ml-kem-r Rust reference's NIST .rsp files.
#
# WHY THIS SCRIPT EXISTS (byte-order convention):
#   The hardware feeds G(reverse(d) || K), i.e. it byte-reverses the d/z seeds
#   before hashing. Consequently the hardware vector files
#       sync_rtl/top/TB/vectors/kat_k{K}_c{C}_{d,z}.hex
#   store d/z in the REVERSE byte order of the NIST .rsp `d`/`z` fields. Both
#   conventions hash to the SAME pk/sk (reverse(reverse(d)) == d), so both the
#   hardware KAT and ml-kem-r's own KAT pass -- they just disagree on how d/z
#   are written down. A naive byte-for-byte compare of d/z therefore "fails"
#   even though the implementations are equivalent. This script compares with
#   the reversal applied, and compares ek/dk (== pk/sk) directly.
#
#   See also: the TB d/z dump in tb_mlkem_kg_katK_xsim.v and the keygen G-hash
#   reversal in sync_rtl/top/mlkem_top.v.
#
# Usage:
#   python3 sync_rtl/top/TB/xcheck_mlkemr.py
#   ML_KEM_R=/path/to/ml-kem-r python3 sync_rtl/top/TB/xcheck_mlkemr.py
#
# Exit code 0 = all consistent, 1 = discrepancy / missing files.

import os
import sys

# Hardware vector dir (relative to repo root, where this script's ../../.. is).
SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
HW_DIR = SCRIPT_DIR  # vectors live under SCRIPT_DIR/vectors

# ml-kem-r location: env override, else common sibling path.
ML_KEM_R = os.environ.get("ML_KEM_R", os.path.expanduser("~/Dev/ml-kem-r"))
RSP = {
    2: "test_data/kat_MLKEM_512.rsp",
    3: "test_data/kat_MLKEM_768.rsp",
    4: "test_data/kat_MLKEM_1024.rsp",
}
# Cases present in the hardware vector set per K (K=2 has 0..4, K=3/4 have 0..2).
CASES = {2: range(5), 3: range(3), 4: range(3)}


def die(msg):
    print(f"ERROR: {msg}", file=sys.stderr)
    sys.exit(1)


def rsp_cases(path):
    """Parse a NIST .rsp into a list of dicts keyed by count."""
    if not os.path.isfile(path):
        die(f"ml-kem-r .rsp not found: {path}\n"
            f"       set ML_KEM_R=/path/to/ml-kem-r (current: {ML_KEM_R})")
    cur, out = None, []
    with open(path) as f:
        for ln in f:
            ln = ln.strip()
            if ln.startswith("count = "):
                if cur is not None:
                    out.append(cur)
                cur = {"count": int(ln[8:])}
            elif " = " in ln and cur is not None:
                k, v = ln.split(" = ", 1)
                cur[k] = v.strip().lower()
    if cur is not None:
        out.append(cur)
    return out


def hwcat(k, c, which):
    """Concatenate a hardware .hex vector file into one lowercase hex string."""
    path = os.path.join(HW_DIR, "vectors", f"kat_k{k}_c{c}_{which}.hex")
    if not os.path.isfile(path):
        die(f"hardware vector not found: {path}")
    with open(path) as f:
        return "".join(f.read().split()).lower()


def revhex(h):
    """Byte-reverse a hex string."""
    return bytes.fromhex(h)[::-1].hex()


def main():
    print(f"ml-kem-r: {ML_KEM_R}")
    print(f"hardware: {os.path.join(HW_DIR, 'vectors')}")
    print("convention: hardware d/z == reverse(NIST d/z); ek==pk, dk==sk\n")

    all_ok = True
    total = 0
    for k in (2, 3, 4):
        cases = rsp_cases(os.path.join(ML_KEM_R, RSP[k]))
        bycount = {c["count"]: c for c in cases}
        for c in CASES[k]:
            total += 1
            hd, hz = hwcat(k, c, "d"), hwcat(k, c, "z")
            hek, hdk = hwcat(k, c, "ek"), hwcat(k, c, "dk")
            rc = bycount.get(c)
            if rc is None:
                die(f"ml-kem-r .rsp for K={k} has no count={c}")
            d_ok = revhex(hd) == rc["d"]
            z_ok = revhex(hz) == rc["z"]
            ek_ok = hek == rc["pk"]
            dk_ok = hdk == rc["sk"]
            case_ok = d_ok and z_ok and ek_ok and dk_ok
            all_ok = all_ok and case_ok
            print(f"K={k} c{c}: d(rev)={'Y' if d_ok else 'N'} "
                  f"z(rev)={'Y' if z_ok else 'N'} "
                  f"ek={'Y' if ek_ok else 'N'} "
                  f"dk={'Y' if dk_ok else 'N'}  -> "
                  f"{'OK' if case_ok else 'MISMATCH'}")

    print(f"\n{total} cases checked: "
          f"{'ALL CONSISTENT' if all_ok else 'DISCREPANCY FOUND'}")
    sys.exit(0 if all_ok else 1)


if __name__ == "__main__":
    main()