#!/usr/bin/env python3
"""gen_vectors.py — Generate SHA3-512 test vectors for tb_sha3_chain_xsim.v

Computes expected rho and sigma values for sha3_chain_top (ML-KEM G function).

The RTL computes: SHA3-512(d_in || 0x02)
  - d_in: 256-bit input (32 bytes, big-endian)
  - 0x02: single byte appended (k=2 parameter for ML-KEM)
  - G(d || 0x02) uses SHA3-512 mode (rate=576, suffix=01)
  - rho  = hash[255:0]   (first 256 bits of hash)
  - sigma = hash[511:256] (next 256 bits of hash)

Output:
  vectors/sha3_chain_input.hex — 256-bit d_in values (64 hex chars per line)

Usage:
  python3 gen_vectors.py
"""

import hashlib
import os

VECTORS_DIR = os.path.join(os.path.dirname(__file__), "vectors")
OUTPUT_FILE = os.path.join(VECTORS_DIR, "sha3_chain_input.hex")

# Test vectors: (label, d_in value)
TEST_VECTORS = [
    ("all-zeros",    0x0000000000000000000000000000000000000000000000000000000000000000),
    ("all-ones",     0xFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFF),
    ("lsb-one",      0x0000000000000000000000000000000000000000000000000000000000000001),
    ("msb-one",      0x8000000000000000000000000000000000000000000000000000000000000000),
    ("pattern-aa",   0xAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA),
    ("pattern-55",   0x5555555555555555555555555555555555555555555555555555555555555555),
    ("random-1",     0x1A2B3C4D5E6F708192A3B4C5D6E7F8091A2B3C4D5E6F708192A3B4C5D6E7F80),
    ("random-2",     0xF0E1D2C3B4A5968778695A4B3C2D1E0FF0E1D2C3B4A5968778695A4B3C2D1E0F),
]


def sha3_512_g(d_in: int) -> tuple:
    """Compute rho, sigma = G(d_in || 0x02) using SHA3-512.

    Returns (rho, sigma) as 256-bit integers.
    """
    # d_in as 32 bytes, big-endian (MSB first)
    d_bytes = d_in.to_bytes(32, "big")
    # Append k=2 as single byte 0x02
    message = d_bytes + b"\x02"
    # SHA3-512 hash → 64 bytes
    h = hashlib.sha3_512(message).digest()
    # rho = first 256 bits, sigma = next 256 bits
    rho = int.from_bytes(h[0:32], "big")
    sigma = int.from_bytes(h[32:64], "big")
    return rho, sigma


def main():
    os.makedirs(VECTORS_DIR, exist_ok=True)

    d_in_values = []
    print("// Expected values computed by gen_vectors.py")
    print("// SHA3-512(d_in || 0x02)")
    print("//" + "=" * 72)

    with open(OUTPUT_FILE, "w") as f:
        for label, d_in in TEST_VECTORS:
            # Write d_in to hex file (256 bits = 64 hex chars)
            f.write(f"{d_in:064X}\n")
            d_in_values.append(d_in)

            # Compute expected values
            rho, sigma = sha3_512_g(d_in)
            print(f"// {label}")
            print(f"//   d_in  = 0x{d_in:064x}")
            print(f"//   rho   = 0x{rho:064x}")
            print(f"//   sigma = 0x{sigma:064x}")
            print()

    print(f"Generated {len(TEST_VECTORS)} vectors → {OUTPUT_FILE}")
    print()
    print("// Copy the expected values above into tb_sha3_chain_xsim.v")
    print("// as hardcoded parameters/initial blocks.")


if __name__ == "__main__":
    main()