"""gen_vectors.py - Test vector generator for sha3_chain module.

Generates random 256-bit d, computes G(d||k=2) using Python reference,
and outputs rho (first 256 bits) and sigma (next 256 bits).
"""

import os
import random
import sys

# Add test_framework/lib to path for VectorGenerator base class
sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', '..', 'lib'))

# Add Python reference path
_REF_PATH = os.path.expanduser(
    "~/Dev/server_code/python_project/PQC_2025/A_ML_KEM_v0")
sys.path.insert(0, _REF_PATH)

import SHA_3

from vector_gen import VectorGenerator


def _bits_to_hex(bits_str, num_bits):
    """Convert binary string (LSB-first: bits_str[0] = bit 0) to hex.

    Returns hex with MSB-nibble first (matching Verilog %X format).
    Each group of 4 bits forms a nibble: bit[i+3:i] → two hex chars.
    """
    result = []
    for i in range(0, num_bits, 4):
        nib = 0
        for j in range(4):
            pos = i + j
            if pos < len(bits_str) and bits_str[pos] == '1':
                nib |= (1 << j)
        result.insert(0, '0123456789ABCDEF'[nib])
    return ''.join(result)


def _random_bits(length):
    """Generate a random binary string of given length (LSB-first)."""
    val = random.getrandbits(length)
    bits = ''
    for i in range(length):
        bits += '1' if (val & (1 << i)) else '0'
    return bits


def _k_bits(k_val):
    """Convert integer k to 8-bit LSB-first binary string.

    k=2 → "01000000" (bit 0=0, bit 1=1)
    """
    bits = ''
    for i in range(8):
        bits += '1' if (k_val & (1 << i)) else '0'
    return bits


class Sha3ChainVectorGenerator(VectorGenerator):
    """Generates test vectors for sha3_chain_top module."""

    def generate_one(self, params: dict) -> dict:
        """Generate a single test vector.

        Args:
            params: dict with 'k' key (default 2 for ML-KEM-512).

        Returns:
            dict with 'input' and 'expected' keys.
        """
        k_val = params.get('k', 2)

        # Generate random 256-bit d (LSB-first binary string)
        d = _random_bits(256)
        k_str = _k_bits(k_val)

        # Compute G(d, k): SHA3-512 of d||k (both LSB-first binary)
        expected_bits = SHA_3.G(d, k_str)

        # rho = first 256 bits, sigma = next 256 bits
        rho_bits   = expected_bits[0:256]
        sigma_bits = expected_bits[256:512]

        # Convert to hex (MSB-first)
        d_hex       = _bits_to_hex(d, 256)
        rho_hex     = _bits_to_hex(rho_bits, 256)
        sigma_hex   = _bits_to_hex(sigma_bits, 256)

        return {
            'input': {
                'd': d_hex
            },
            'expected': {
                'rho': rho_hex,
                'sigma': sigma_hex
            }
        }

    def write_hex_file(self, vectors: list[dict], filepath: str) -> None:
        """Write input vectors: one 64-char d-hex per line.

        Args:
            vectors: List of vector dicts from generate_one().
            filepath: Path to write the hex file.
        """
        os.makedirs(os.path.dirname(filepath), exist_ok=True)
        with open(filepath, 'w') as f:
            for v in vectors:
                d_hex = v['input']['d']
                f.write(f'{d_hex}\n')

    def write_expected_file(self, vectors: list[dict], filepath: str) -> None:
        """Write expected output: "RHO_HEX SIGMA_HEX" per line (128 chars).

        Args:
            vectors: List of vector dicts from generate_one().
            filepath: Path to write the expected hex file.
        """
        os.makedirs(os.path.dirname(filepath), exist_ok=True)
        with open(filepath, 'w') as f:
            for v in vectors:
                rho_hex   = v['expected']['rho']
                sigma_hex = v['expected']['sigma']
                f.write(f'{rho_hex} {sigma_hex}\n')