#!/usr/bin/env python3
"""gen_vectors.py - Test vector generator for sample_cbd_sync module.

Generates random seed+nonce+eta test vectors, computes expected CBD coefficients
using Python hashlib (stdlib only), and writes input hex and expected output hex files.

Algorithm (matching RTL sample_cbd_sync.v):
  PRF(sigma, N) = SHAKE-256(sigma || N) → squeeze eta*64 bytes
  For each of 256 coefficients:
    eta=2: read 4 bits, coeff = (b0+b1) - (b2+b3)
    eta=3: read 6 bits, coeff = (b0+b1+b2) - (b3+b4+b5)
  Each coefficient in range [-eta, eta], stored as 12-bit signed.

Bit ordering (FIPS 202 / RTL match):
  The RTL feeds seed_i[0] as the first bit into SHA3.
  Python hashlib expects bytes[0] LSB as the first bit.
  $readmemh stores hex with MSB first → seed_i[255:0] MSB-first.
  So: Python input = reverse(bytes.fromhex(seed_hex)) + bytes([nonce]).

Usage:
  python3 gen_vectors.py              # Generate vectors
  python3 gen_vectors.py --verify     # Verify results against expected
"""

import hashlib
import os
import random
import sys

N_COEFFS = 256
Q = 3329  # not used for CBD, but for reference


def random_hex(bits):
    """Generate a random hex string (MSB-first) of the given bit length."""
    val = random.getrandbits(bits)
    num_nibbles = (bits + 3) // 4
    return f"{val:0{num_nibbles}X}"


def bits_from_bytes(data):
    """Convert bytes to LSB-first bit string.

    data[0] LSB = bits[0], data[0] bit 1 = bits[1], ...
    """
    bits = []
    for b in data:
        for bit in range(8):
            bits.append('1' if (b >> bit) & 1 else '0')
    return ''.join(bits)


def shake256_prf(seed_hex, nonce, eta):
    """Compute SHAKE-256(seed || nonce) matching RTL bit ordering.

    Args:
        seed_hex: 64-char hex string (MSB-first, as stored by $readmemh).
        nonce: 8-bit integer.
        eta: 2 or 3.

    Returns:
        bytes: eta * 64 bytes of SHAKE-256 squeeze output.
    """
    # seed_hex is MSB-first: seed_i[255:0] = 256'h<seed_hex>
    # RTL feeds seed_i[0] first into SHA3.
    # Python: reverse bytes so that byte[0] = seed_i[7:0],
    # then append nonce byte.
    seed_bytes = bytes.fromhex(seed_hex)  # MSB-first bytes
    shake_input = bytes(reversed(seed_bytes)) + bytes([nonce & 0xFF])
    shake = hashlib.shake_256(shake_input)
    return shake.digest(eta * 64)


def cbd_sample(prf_bytes, eta):
    """Apply Centered Binomial Distribution to PRF output bytes.

    Args:
        prf_bytes: bytes from SHAKE-256 squeeze (eta * 64 bytes).
        eta: 2 or 3.

    Returns:
        list of 256 signed integers in range [-eta, eta].
    """
    bits = bits_from_bytes(prf_bytes)
    step = eta * 2  # 4 for eta=2, 6 for eta=3
    half = eta       # 2 for eta=2, 3 for eta=3
    coeffs = []
    for i in range(N_COEFFS):
        pos_sum = sum(1 for j in range(half) if bits[step * i + j] == '1')
        neg_sum = sum(1 for j in range(half) if bits[step * i + half + j] == '1')
        coeffs.append(pos_sum - neg_sum)
    return coeffs


def coeff_to_hex_12signed(val):
    """Convert signed 12-bit value to 3-char hex string (two's complement)."""
    masked = val & 0xFFF  # 12-bit unsigned
    return f"{masked:03X}"


def generate_one(eta):
    """Generate a single test vector.

    Returns:
        dict with 'seed_hex', 'nonce', 'eta', 'coeffs' keys.
    """
    seed_hex = random_hex(256)
    nonce = random.randint(0, 255)
    prf_bytes = shake256_prf(seed_hex, nonce, eta)
    coeffs = cbd_sample(prf_bytes, eta)
    return {
        "seed_hex": seed_hex,
        "nonce": nonce,
        "eta": eta,
        "coeffs": coeffs,
    }


def write_input_hex(vectors, filepath):
    """Write input vectors as packed hex for $readmemh.

    Each line: {eta[1:0], nonce[7:0], seed[255:0]} = 266 bits.
    Written as 67 hex chars (268 bits, top 2 bits zero).
    """
    os.makedirs(os.path.dirname(filepath), exist_ok=True)
    with open(filepath, "w") as f:
        for v in vectors:
            eta = v["eta"]
            nonce = v["nonce"]
            seed_hex = v["seed_hex"]
            # Pack: {2'b00, eta[1:0], nonce[7:0], seed[255:0]} = 268 bits = 67 hex chars
            # The first hex char carries eta in its lower 2 bits:
            #   hex value = 0x0 | eta → nibble = 0b00XX where XX = eta
            # TB reads vec_eta = vector_mem[idx][265:264]
            packed = f"{eta & 0x3:01X}{nonce:02X}{seed_hex}"  # 1+2+64 = 67 chars
            f.write(packed + "\n")


def write_expected_hex(vectors, filepath):
    """Write expected coefficients: one 12-bit hex value per line."""
    os.makedirs(os.path.dirname(filepath), exist_ok=True)
    with open(filepath, "w") as f:
        for i, v in enumerate(vectors):
            f.write(f"# VECTOR_{i} eta={v['eta']} nonce=0x{v['nonce']:02X} seed={v['seed_hex']}\n")
            for c in v["coeffs"]:
                f.write(coeff_to_hex_12signed(c) + "\n")


def verify_results(result_file, vectors):
    """Verify RTL output against expected values.

    Args:
        result_file: Path to RTL result file (coeff hex per line).
        vectors: List of vector dicts with 'coeffs'.

    Returns:
        bool: True if all vectors match.
    """
    with open(result_file, "r") as f:
        lines = [line.strip() for line in f
                 if line.strip() and not line.strip().startswith("#")]

    # lines may contain trailing comments separated by "  #"
    cleaned = []
    for line in lines:
        comment_idx = line.find("  #")
        if comment_idx >= 0:
            line = line[:comment_idx].strip()
        cleaned.append(line)

    expected_all = []
    for v in vectors:
        for c in v["coeffs"]:
            expected_all.append(coeff_to_hex_12signed(c))

    if len(cleaned) != len(expected_all):
        print(f"  COUNT MISMATCH: got {len(cleaned)}, expected {len(expected_all)}")
        return False

    mismatches = 0
    for i, (g, e) in enumerate(zip(cleaned, expected_all)):
        if g.upper() != e.upper():
            if mismatches < 10:
                print(f"  MISMATCH[{i}]: got={g.upper()}, expected={e.upper()}")
            mismatches += 1

    if mismatches > 0:
        print(f"  Total mismatches: {mismatches}")
        return False

    return True


def main():
    vector_count = 4  # 2 vectors with eta=2, 2 with eta=3

    base_dir = os.path.dirname(os.path.abspath(__file__))
    vectors_dir = os.path.join(base_dir, "vectors")
    input_file = os.path.join(vectors_dir, "sample_cbd_input.hex")
    expected_file = os.path.join(vectors_dir, "sample_cbd_expected.hex")
    result_file = os.path.join(vectors_dir, "sample_cbd_result.hex")

    verify_mode = "--verify" in sys.argv

    if verify_mode:
        if not os.path.exists(result_file):
            print(f"ERROR: Result file not found: {result_file}")
            print("  Run simulation first to generate results.")
            sys.exit(1)
        if not os.path.exists(input_file):
            print(f"ERROR: Input file not found: {input_file}")
            sys.exit(1)
        print(f"Verifying results from {result_file}...")
        # Recompute expected from input file
        with open(input_file, "r") as f:
            input_lines = [l.strip() for l in f if l.strip()]
        vectors = []
        for line in input_lines:
            if len(line) != 67:
                print(f"WARNING: Skipping line with unexpected length {len(line)}: {line[:20]}...")
                continue
            # Parse: first char = eta_nibble, next 2 = nonce, next 64 = seed
            eta_nibble = int(line[0], 16)
            eta = eta_nibble >> 2  # bits[3:2]
            nonce = int(line[1:3], 16)
            seed_hex = line[3:]
            prf_bytes = shake256_prf(seed_hex, nonce, eta)
            coeffs = cbd_sample(prf_bytes, eta)
            vectors.append({"seed_hex": seed_hex, "nonce": nonce, "eta": eta, "coeffs": coeffs})

        ok = verify_results(result_file, vectors)
        if ok:
            print("ALL VECTORS PASSED")
        else:
            print("VERIFICATION FAILED")
            sys.exit(1)
    else:
        # Generate mode
        print(f"Generating {vector_count} test vectors...")
        vectors = []
        # eta=2 vectors
        for i in range(vector_count // 2):
            v = generate_one(eta=2)
            vectors.append(v)
            print(f"  Vector {len(vectors)-1}: eta=2, nonce=0x{v['nonce']:02X}, "
                  f"seed={v['seed_hex'][:8]}..., coeffs[0]={v['coeffs'][0]}")
        # eta=3 vectors
        for i in range(vector_count // 2):
            v = generate_one(eta=3)
            vectors.append(v)
            print(f"  Vector {len(vectors)-1}: eta=3, nonce=0x{v['nonce']:02X}, "
                  f"seed={v['seed_hex'][:8]}..., coeffs[0]={v['coeffs'][0]}")

        write_input_hex(vectors, input_file)
        print(f"Wrote {len(vectors)} vectors to {input_file}")

        write_expected_hex(vectors, expected_file)
        print(f"Wrote expected coefficients to {expected_file}")

        # Sanity checks
        for v in vectors:
            for c in v["coeffs"]:
                assert -v["eta"] <= c <= v["eta"], \
                    f"Coefficient {c} out of range [-{v['eta']}, {v['eta']}]"

        print("All sanity checks passed.")


if __name__ == "__main__":
    main()