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

Generates random rho+k+i+j test vectors, computes expected coefficients using
Keccak-p[1600, 24] matching the RTL's per-permutation squeeze pattern.

The RTL sample_ntt_sync absorbs rho || j || i, then repeatedly permutes
the keccak state, extracting d1,d2 from bits [23:0] of each permuted state
with rejection sampling (d < Q=3329).

Algorithm (matching RTL exactly):
  1. Build absorb state: pad10*1 padding for SHAKE-128 (rate=1344) on msg(272b)
  2. state = keccak_p(absorb_state)
  3. Extract d1 = (state[11:8] << 8) | state[7:0]
            d2 = (state[23:16] << 4) | state[15:12]
  4. If d1 < Q: output d1 (12-bit unsigned)
  5. If d2 < Q: output d2
  6. state = keccak_p(state)
  7. Repeat 3-6 until 256 coefficients collected

Bit ordering (FIPS 202 / RTL match):
  The RTL feeds rho_i[0] as the first message bit into SHA3.
  $readmemh stores hex MSB-first → rho_i[255:0] in MSB-first order.
  Python: input bytes = reversed(bytes.fromhex(rho_hex)) + bytes([j, i])

Usage:
  python3 gen_vectors.py              # Generate vectors
  python3 gen_vectors.py --verify     # Verify results against expected
  python3 gen_vectors.py --selftest   # Run Keccak-p self-test
"""

import os
import random
import sys

Q = 3329
N_COEFFS = 256

# ======================================================================
# Keccak-p[1600, 24] implementation (stdlib only, matches RTL exactly)
# ======================================================================

# Rotation offsets for rho step: RHO[x][y]
RHO = [
    [ 0, 36,  3, 41, 18],
    [ 1, 44, 10, 45,  2],
    [62,  6, 43, 15, 61],
    [28, 55, 25, 21, 56],
    [27, 20, 39,  8, 14],
]

# Round constants for iota step
RC = [
    0x0000000000000001, 0x0000000000008082, 0x800000000000808A, 0x8000000080008000,
    0x000000000000808B, 0x0000000080000001, 0x8000000080008081, 0x8000000000008009,
    0x000000000000008A, 0x0000000000000088, 0x0000000080008009, 0x000000008000000A,
    0x000000008000808B, 0x800000000000008B, 0x8000000000008089, 0x8000000000008003,
    0x8000000000008002, 0x8000000000000080, 0x000000000000800A, 0x800000008000000A,
    0x8000000080008081, 0x8000000000008080, 0x0000000080000001, 0x8000000080008008,
]


def ROTL64(val, n):
    """Rotate 64-bit value left by n bits."""
    n = n & 63
    return ((val << n) | (val >> (64 - n))) & 0xFFFFFFFFFFFFFFFF


def keccak_round(lanes, rnd_idx):
    """Single Keccak-f round: theta, rho, pi, chi, iota.

    lanes: list of 25 64-bit integers, indexed as lane_idx = 5*y + x.
    Returns new list of 25 lanes.
    """
    # === Theta ===
    C = [0] * 5  # C[x] = A[x,0] ^ A[x,1] ^ A[x,2] ^ A[x,3] ^ A[x,4]
    for x in range(5):
        for y in range(5):
            C[x] ^= lanes[5 * y + x]

    # D[x] = C[x-1] ^ ROTL(C[x+1], 1)
    D = [0] * 5
    for x in range(5):
        D[x] = C[(x - 1) % 5] ^ ROTL64(C[(x + 1) % 5], 1)

    new_lanes = [0] * 25
    for y in range(5):
        for x in range(5):
            idx = 5 * y + x
            new_lanes[idx] = lanes[idx] ^ D[x]

    # === Rho + Pi ===
    B = [0] * 25
    for x in range(5):
        for y in range(5):
            src_idx = 5 * y + x
            dst_row = (2 * x + 3 * y) % 5
            dst_col = y
            dst_idx = 5 * dst_row + dst_col
            B[dst_idx] = ROTL64(new_lanes[src_idx], RHO[x][y])

    # === Chi ===
    result = [0] * 25
    for y in range(5):
        for x in range(5):
            idx = 5 * y + x
            idx1 = 5 * y + ((x + 1) % 5)
            idx2 = 5 * y + ((x + 2) % 5)
            result[idx] = B[idx] ^ ((~B[idx1] & 0xFFFFFFFFFFFFFFFF) & B[idx2])

    # === Iota ===
    result[0] ^= RC[rnd_idx]

    return result


def keccak_p(state_int):
    """Keccak-p[1600, 24] permutation of a 1600-bit integer.

    Returns 1600-bit integer.
    """
    # Convert 1600-bit integer to 25 lanes of 64 bits
    lanes = [0] * 25
    for y in range(5):
        for x in range(5):
            idx = 5 * y + x
            lanes[idx] = (state_int >> (64 * idx)) & 0xFFFFFFFFFFFFFFFF

    for rnd in range(24):
        lanes = keccak_round(lanes, rnd)

    # Convert back to 1600-bit integer
    result = 0
    for y in range(5):
        for x in range(5):
            idx = 5 * y + x
            result |= lanes[idx] << (64 * idx)
    return result & ((1 << 1600) - 1)


def keccak_p_selftest():
    """Verify Keccak-p against known test vectors.

    Validates by computing SHA3-256('') and SHA3-256('abc')
    against Python's hashlib reference.
    """
    import hashlib

    def sha3_256_via_keccak(data):
        """Compute SHA3-256 using our keccak_p, compare with hashlib.

        SHA3-256(M) = Keccak[c=512](M || 01, 256)
        Rate = 1088 bits, capacity = 512 bits.
        """
        # Convert message to bits, LSB-first per byte
        msg_bits = ''.join(format(b, '08b')[::-1] for b in data)

        # Append SHA-3 suffix '01' (bit 0, then bit 1)
        msg_bits += '01'

        # Apply pad10*1: '1' + required zeros + '1'
        pad_len = (1088 - (len(msg_bits) % 1088)) % 1088
        if pad_len == 0:
            pad_len = 1088
        # pad_len >= 2 always for pad10*1
        msg_bits += '1' + '0' * (pad_len - 2) + '1'

        # Build state by XOR-ing padded message into rate portion
        state = 0
        for i, bit in enumerate(msg_bits):
            if bit == '1':
                state ^= (1 << i)

        # Absorb (single block for short messages)
        state = keccak_p(state)

        # Squeeze 256 bits from rate portion [255:0]
        output_bytes = bytes(
            (state >> (8 * i)) & 0xFF for i in range(32)
        )

        return output_bytes.hex()

    # Test SHA3-256("") against hashlib
    expected_empty = hashlib.sha3_256(b"").hexdigest()
    got_empty = sha3_256_via_keccak(b"")
    assert got_empty == expected_empty, \
        f"SHA3-256('') mismatch:\n  got:      {got_empty}\n  expected: {expected_empty}"

    # Test SHA3-256("abc") against hashlib
    expected_abc = hashlib.sha3_256(b"abc").hexdigest()
    got_abc = sha3_256_via_keccak(b"abc")
    assert got_abc == expected_abc, \
        f"SHA3-256('abc') mismatch:\n  got:      {got_abc}\n  expected: {expected_abc}"

    # Also test that our keccak_p is not identity
    assert keccak_p(0) != 0, "keccak_p(0) must not equal 0"

    print("Keccak-p self-test PASSED (SHA3-256 verified against hashlib)")
    return True


# ======================================================================
# Message construction (matching RTL pad10*1 for SHAKE-128)
# ======================================================================

def build_absorb_state(rho_hex, j, i):
    """Build the 1600-bit absorb state matching RTL sample_ntt_sync.

    RTL absorb_state = {capacity(256b), pad10*1, suffix(1111), msg(272b)}

    Message order (bit 0 = state[0]):
      rho_i[0], ..., rho_i[255], j[0], ..., j[7], i[0], ..., i[7]

    Args:
        rho_hex: 64-char hex string (MSB-first).
        j: 2-bit index (j_idx).
        i: 2-bit index (i_idx).

    Returns:
        1600-bit integer representing the padded absorb state.
    """
    # Build message bytes
    rho_bytes = bytes.fromhex(rho_hex)  # MSB-first → byte 0 = rho_i[255:248]
    rho_rev = rho_bytes[::-1]           # Reversed: byte 0 = rho_i[7:0]
    msg_bytes = rho_rev + bytes([j & 0xFF, i & 0xFF])
    # msg_bytes[0] = rho_i[7:0], msg_bytes[1] = rho_i[15:8], ...
    # msg_bytes[32] = j, msg_bytes[33] = i

    # Build 1600-bit state
    # state[7:0] = msg_bytes[0]
    # state[15:8] = msg_bytes[1]
    # ...
    # state[275:272] = 4'b1111  (SHAKE suffix)
    # state[276] = 1'b1       (pad10*1 first 1)
    # state[1342:277] = 0
    # state[1343] = 1'b1      (pad10*1 final 1)
    # state[1599:1344] = 0    (capacity)

    state = 0
    # Message bytes at bits [0:271]
    for idx, b in enumerate(msg_bytes):
        state |= (b & 0xFF) << (8 * idx)

    # SHAKE suffix (1111) at bits [272:275]
    state |= 0xF << 272

    # pad10*1: 1 at bit 276
    state |= 1 << 276

    # pad10*1: 1 at bit 1343
    state |= 1 << 1343

    return state & ((1 << 1600) - 1)


def extract_d1_d2(state):
    """Extract d1, d2 from state[23:0] matching RTL.

    Returns (d1, d2) as 12-bit unsigned integers [0, 4095].
    """
    c0 = (state >> 0) & 0xFF
    c1 = (state >> 8) & 0xFF
    c2 = (state >> 16) & 0xFF

    d1 = ((c1 & 0xF) << 8) | c0
    d2 = (c2 << 4) | (c1 >> 4)
    return d1, d2


def sample_ntt_compute(rho_hex, k, i, j):
    """Compute expected coefficients matching RTL sample_ntt_sync.

    The RTL absorbs rho || j || i, then repeatedly permutes,
    extracting d1,d2 from each permuted state with rejection sampling.

    Args:
        rho_hex: 64-char hex string.
        k: 3-bit k value (unused in computation, passed through).
        i: 2-bit i_idx.
        j: 2-bit j_idx.

    Returns:
        list of 256 integers in [0, Q-1].
    """
    state = build_absorb_state(rho_hex, j, i)
    state = keccak_p(state)

    coeffs = []
    while len(coeffs) < N_COEFFS:
        d1, d2 = extract_d1_d2(state)
        if d1 < Q:
            coeffs.append(d1)
            if len(coeffs) == N_COEFFS:
                break
        if d2 < Q:
            coeffs.append(d2)
            if len(coeffs) == N_COEFFS:
                break
        # Re-permute for next squeeze
        state = keccak_p(state)

    return coeffs


# ======================================================================
# Helper functions
# ======================================================================

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 coeff_to_hex(val):
    """Convert 12-bit unsigned coefficient to 3-char hex string."""
    return f"{val & 0xFFF:03X}"


# ======================================================================
# Vector generation and file I/O
# ======================================================================

def generate_one(k, i, j):
    """Generate a single test vector.

    Returns dict with 'rho_hex', 'k', 'i', 'j', 'coeffs'.
    """
    rho_hex = random_hex(256)
    coeffs = sample_ntt_compute(rho_hex, k, i, j)
    return {
        "rho_hex": rho_hex,
        "k": k,
        "i": i,
        "j": j,
        "coeffs": coeffs,
    }


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

    Each line: {1'b0, j_idx[1:0], i_idx[1:0], k_i[2:0], rho_i[255:0]}
    = 264 bits = 66 hex chars.
    """
    os.makedirs(os.path.dirname(filepath), exist_ok=True)
    with open(filepath, "w") as f:
        for v in vectors:
            rho_hex = v["rho_hex"]
            k = v["k"] & 0x7
            i = v["i"] & 0x3
            j = v["j"] & 0x3
            # Pack: {1'b0, j[1:0], i[1:0], k[2:0], rho[255:0]}
            header = (j << 5) | (i << 3) | k
            # header occupies bits [262:256] (when zero-padded properly)
            # rho occupies bits [255:0]
            # Total: 256 + 7 = 263 bits → 66 hex chars (264 bits, top bit = 0)
            # header_hex: 2 chars
            header_hex = f"{header:02X}"
            packed = header_hex + rho_hex  # 2 + 64 = 66 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 idx, v in enumerate(vectors):
            f.write(f"# VECTOR_{idx} k={v['k']} i={v['i']} j={v['j']} "
                    f"rho={v['rho_hex']}\n")
            for c in v["coeffs"]:
                f.write(coeff_to_hex(c) + "\n")


def verify_results(result_file, vectors):
    """Verify RTL output against expected values."""
    with open(result_file, "r") as f:
        lines = [line.strip() for line in f
                 if line.strip() and not line.strip().startswith("#")]

    # Remove trailing comments
    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(c))

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

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

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

    return True


# ======================================================================
# Main
# ======================================================================

def main():
    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_ntt_input.hex")
    expected_file = os.path.join(vectors_dir, "sample_ntt_expected.hex")
    result_file = os.path.join(vectors_dir, "sample_ntt_result.hex")

    if "--selftest" in sys.argv:
        keccak_p_selftest()
        print("Running quick smoke test on sample_ntt computation...")
        rho = "000102030405060708090A0B0C0D0E0F101112131415161718191A1B1C1D1E1F"
        coeffs = sample_ntt_compute(rho, k=4, i=0, j=0)
        print(f"  rho={rho[:16]}...: {len(coeffs)} coefficients")
        print(f"  first 4: {coeffs[:4]}")
        assert len(coeffs) == 256, f"Expected 256 coefficients, got {len(coeffs)}"
        assert all(0 <= c < Q for c in coeffs), "Coefficient out of range"
        print("Smoke test PASSED")
        return

    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}")
            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) != 66:
                print(f"WARNING: Skipping line with unexpected length {len(line)}: {line[:20]}...")
                continue
            header_hex = line[:2]
            rho_hex = line[2:]
            header = int(header_hex, 16)
            j = (header >> 5) & 0x3
            i = (header >> 3) & 0x3
            k = header & 0x7
            coeffs = sample_ntt_compute(rho_hex, k, i, j)
            vectors.append({"rho_hex": rho_hex, "k": k, "i": i, "j": j, "coeffs": coeffs})

        ok = verify_results(result_file, vectors)
        if ok:
            print("ALL VECTORS PASSED")
        else:
            print("VERIFICATION FAILED")
            sys.exit(1)
    else:
        # Generate mode
        vector_count = 4
        print(f"Generating {vector_count} test vectors...")
        print(f"Running Keccak-p self-test first...")
        keccak_p_selftest()

        vectors = []
        for idx in range(vector_count):
            k = random.choice([2, 3, 4])
            i = random.randint(0, 3)
            j = random.randint(0, 3)
            v = generate_one(k, i, j)
            vectors.append(v)
            print(f"  Vector {idx}: k={k}, i={i}, j={j}, "
                  f"rho={v['rho_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 0 <= c < Q, f"Coefficient {c} out of range [0, {Q-1}]"
            assert len(v["coeffs"]) == N_COEFFS

        print("All sanity checks passed.")


if __name__ == "__main__":
    main()