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

For ML-KEM-512/768/1024, parse the NIST .rsp and write, per case:
  enc_k{K}_c{N}_ek.hex  : ek (=pk) bytes, one hex byte per line, byte 0 first
  enc_k{K}_c{N}_m.hex   : msg bytes, one per line, byte 0 first
  enc_k{K}_c{N}_ct.hex  : ct bytes, one per line, byte 0 first (golden)
  enc_k{K}_c{N}_ss.hex  : ss bytes, one per line, byte 0 first (golden)

Per-byte layout removes all endianness ambiguity: the TB preloads ek_bram with
addr i = line i, builds m_i[8*i +: 8] = m byte i, and compares ct/ss byte-wise.

Run: python3 gen_encaps_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, msg, pk, ct, ss} (hex strings) 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 so "ct =" doesn't catch "ct_n ="
            m = re.match(r"^(msg|pk|ct|ss)\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"enc_k{k}_c{n}")
            write_bytes(f"{base}_ek.hex", v["pk"])
            write_bytes(f"{base}_m.hex", v["msg"])
            write_bytes(f"{base}_ct.hex", v["ct"])
            write_bytes(f"{base}_ss.hex", v["ss"])
        print(f"K={k}: wrote {len(vecs)} cases "
              f"(ek={len(vecs[0]['pk'])//2}B ct={len(vecs[0]['ct'])//2}B)")


if __name__ == "__main__":
    main()