#!/usr/bin/env python3
"""gen_vectors.py - Parse FIPS 203 KAT test vectors for ML-KEM-512.

Reads kat_MLKEM_512.rsp and generates hex vector files for the
mlkem_top XSIM testbench.

Output files (written to vectors/ subdirectory):
  mlkem_top_input.hex    - d (256b) + msg (256b) + z (256b) per line
  mlkem_top_expected.hex - pk (6400b) + sk (13056b) + ct (6144b) + ss (256b)

KAT file format (per test vector):
  count = N
  z = <64 hex chars>     → 32 bytes, implicit rejection seed
  d = <64 hex chars>     → 32 bytes, KeyGen randomness
  msg = <64 hex chars>   → 32 bytes, Encaps message
  seed = <128 hex chars>  → DRBG seed (ignored)
  pk = <1600 hex chars>  → 800 bytes, public key
  sk = <3264 hex chars>  → 1632 bytes, secret key
  ct_n = <hex>           → invalid ciphertext (ignored)
  ss_n = <hex>            → shared secret from invalid ct (ignored)
  ct = <1536 hex chars>  → 768 bytes, ciphertext
  ss = <64 hex chars>    → 32 bytes, shared secret
"""

import os
import re
import sys

KAT_PATH = "/home/fallensigh/Dev/ml-kem-r/test_data/kat_MLKEM_512.rsp"
OUT_DIR = os.path.join(os.path.dirname(__file__), "vectors")
NUM_VECTORS = 5


def parse_kat(filepath: str, num_vectors: int = NUM_VECTORS) -> list[dict]:
    """Parse KAT .rsp file and return list of test vector dicts."""
    vectors: list[dict] = []
    current: dict = {}

    with open(filepath, "r") as f:
        for line in f:
            line = line.strip()

            # New vector starts with "count = N"
            m = re.match(r"^count\s*=\s*(\d+)$", line)
            if m:
                if current and "count" in current:
                    vectors.append(current)
                    current = {}
                idx = int(m.group(1))
                if idx >= num_vectors:
                    break
                current = {"count": idx}
                continue

            # Parse key = value lines
            m = re.match(r"^(\w+)\s*=\s*([0-9a-fA-F]+)$", line)
            if m and current:
                key = m.group(1)
                value = m.group(2).lower()
                current[key] = value
                continue

        # Add last vector if present
        if current and "count" in current:
            vectors.append(current)

    # Validate each vector has required fields
    required = {"d", "z", "msg", "pk", "sk", "ct", "ss"}
    for v in vectors:
        missing = required - set(v.keys())
        if missing:
            print(f"WARNING: count={v['count']} missing fields: {missing}",
                  file=sys.stderr)

    return vectors


def verify_lengths(vectors: list[dict]) -> bool:
    """Verify hex field lengths match FIPS 203 expected sizes."""
    expected = {
        "d":   64,   # 32 bytes
        "z":   64,   # 32 bytes
        "msg": 64,   # 32 bytes
        "pk":  1600, # 800 bytes
        "sk":  3264, # 1632 bytes
        "ct":  1536, # 768 bytes
        "ss":  64,   # 32 bytes
    }
    ok = True
    for v in vectors:
        for field, elen in expected.items():
            if field not in v:
                continue
            actual = len(v[field])
            if actual != elen:
                print(f"WARNING: count={v['count']} {field}: "
                      f"expected {elen} hex chars, got {actual}",
                      file=sys.stderr)
                ok = False
    return ok


def write_input_hex(vectors: list[dict], out_dir: str) -> str:
    """Write input vectors: d || msg || z (768 bits = 192 hex chars each)."""
    os.makedirs(out_dir, exist_ok=True)
    out_path = os.path.join(out_dir, "mlkem_top_input.hex")
    with open(out_path, "w") as f:
        for v in vectors:
            d = v.get("d", "0" * 64)
            msg = v.get("msg", "0" * 64)
            z = v.get("z", "0" * 64)
            # Concatenate: [d:255:0][msg:255:0][z:255:0]
            f.write(f"{d}{msg}{z}\n")
    print(f"Wrote {len(vectors)} input vectors to {out_path}")
    return out_path


def write_expected_hex(vectors: list[dict], out_dir: str) -> str:
    """Write expected output vectors: pk || sk || ct || ss.

    pk:  800 bytes = 1600 hex chars
    sk:  1632 bytes = 3264 hex chars
    ct:  768 bytes = 1536 hex chars
    ss:  32 bytes = 64 hex chars
    Total: 3232 bytes = 6464 hex chars per line
    """
    os.makedirs(out_dir, exist_ok=True)
    out_path = os.path.join(out_dir, "mlkem_top_expected.hex")
    with open(out_path, "w") as f:
        for v in vectors:
            pk = v.get("pk", "0" * 1600)
            sk = v.get("sk", "0" * 3264)
            ct = v.get("ct", "0" * 1536)
            ss = v.get("ss", "0" * 64)
            f.write(f"{pk}{sk}{ct}{ss}\n")
    print(f"Wrote {len(vectors)} expected vectors to {out_path}")
    return out_path


def main():
    vectors = parse_kat(KAT_PATH, NUM_VECTORS)
    print(f"Parsed {len(vectors)} test vectors from KAT file")

    if len(vectors) == 0:
        print("ERROR: No vectors parsed!", file=sys.stderr)
        sys.exit(1)

    verify_lengths(vectors)

    # Print summary
    for v in vectors:
        print(f"  count={v['count']}: "
              f"d={v.get('d', 'N/A')[:8]}... "
              f"msg={v.get('msg', 'N/A')[:8]}... "
              f"z={v.get('z', 'N/A')[:8]}... "
              f"ss={v.get('ss', 'N/A')[:8]}...")

    write_input_hex(vectors, OUT_DIR)
    write_expected_hex(vectors, OUT_DIR)

    print("Done.")


if __name__ == "__main__":
    main()