#!/usr/bin/env python3
"""gen_decaps_vectors.py - Emit per-byte Decaps KAT vectors for the XSIM TB.

For ML-KEM-512/768/1024, parse the NIST .rsp and write, per case:
  dec_k{K}_c{N}_dk.hex  : dk (=sk) bytes, one hex byte per line, byte 0 first
  dec_k{K}_c{N}_ct.hex  : ct bytes (valid ciphertext), one per line, byte 0 first
  dec_k{K}_c{N}_ss.hex  : ss bytes (shared secret for valid ct), byte 0 first
  dec_k{K}_c{N}_ctn.hex : ct_n bytes (corrupted ct -> implicit reject)
  dec_k{K}_c{N}_ssn.hex : ss_n bytes (shared secret K-bar for the reject path)

The TB streams dk into the design via dk_in_* (routed to dk_pke/ek_pke/h/z by
region) and ct into c_in_*, runs Decaps, and checks ss:
  valid ct  -> ss == ss   (c'==c, K' kept)
  ct_n      -> ss == ss_n (c'!=c, K-bar via J(z||c_n))

Run: python3 gen_decaps_vectors.py [num_cases]
"""

import os
import re
import sys

ML_KEM_R = os.environ.get("ML_KEM_R", os.path.expanduser("~/Dev/ml-kem-r"))
OUT_DIR = os.path.join(os.path.dirname(__file__), "vectors")
KATS = {2: "kat_MLKEM_512.rsp", 3: "kat_MLKEM_768.rsp", 4: "kat_MLKEM_1024.rsp"}


def parse_kat(path, n):
    """Return list of dicts {count, sk, ct, ss, ct_n, ss_n} (hex) for first n."""
    vecs, cur = [], {}
    with open(path) as f:
        for line in f:
            line = line.strip()
            m = re.match(r"^count\s*=\s*(\d+)$", line)
            if m:
                if cur:
                    vecs.append(cur)
                if len(vecs) >= n:
                    break
                cur = {"count": int(m.group(1))}
                continue
            # exact-key match: 'ct =' vs 'ct_n =', 'ss =' vs 'ss_n ='
            m = re.match(r"^(sk|ct|ss|ct_n|ss_n)\s*=\s*([0-9a-fA-F]+)$", line)
            if m and cur:
                cur[m.group(1)] = m.group(2).lower()
    if cur and len(vecs) < n:
        vecs.append(cur)
    return vecs


def write_bytes(path, hexstr):
    """Write hex string as one byte per line (byte 0 = first 2 hex chars)."""
    with open(path, "w") as f:
        for i in range(0, len(hexstr), 2):
            f.write(hexstr[i:i + 2] + "\n")


def main():
    ncases = int(sys.argv[1]) if len(sys.argv) > 1 else 3
    os.makedirs(OUT_DIR, exist_ok=True)
    for k, fname in KATS.items():
        path = os.path.join(ML_KEM_R, "test_data", fname)
        if not os.path.exists(path):
            print(f"skip K={k}: {path} not found", file=sys.stderr)
            continue
        vecs = parse_kat(path, ncases)
        for v in vecs:
            n = v["count"]
            base = os.path.join(OUT_DIR, f"dec_k{k}_c{n}")
            write_bytes(f"{base}_dk.hex", v["sk"])
            write_bytes(f"{base}_ct.hex", v["ct"])
            write_bytes(f"{base}_ss.hex", v["ss"])
            write_bytes(f"{base}_ctn.hex", v["ct_n"])
            write_bytes(f"{base}_ssn.hex", v["ss_n"])
        print(f"K={k}: wrote {len(vecs)} cases "
              f"(dk={len(vecs[0]['sk'])//2}B ct={len(vecs[0]['ct'])//2}B)")


if __name__ == "__main__":
    main()