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

Generates vectors for CBD sampling with eta=2 and eta=3 using the
Python reference SHA_3.PRF for SHAKE-256 expander output, and a local
CBD implementation that outputs signed 12-bit coefficients.
"""

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


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

    Returns hex string where leftmost char = most significant nibble.
    Each nibble maps bits[3:0] to a hex char (bit 3 = MSB of nibble).
    """
    num_bits = len(bits_str)
    result = []
    for i in range(num_bits - 4, -1, -4):
        nib = 0
        for j in range(4):
            pos = i + j
            if pos < num_bits and bits_str[pos] == '1':
                nib |= (1 << j)
        result.append('0123456789ABCDEF'[nib])
    return ''.join(result)


def _sample_poly_cbd(prf_bits, eta):
    """Centered Binomial Distribution on PRF output bits.

    Reads eta*2 bits per coefficient from the binary string (LSB-first).
    Each coefficient = sum of first eta bits - sum of last eta bits.
    Result is a signed integer in range [-eta, eta].

    Args:
        prf_bits: Binary string (LSB-first) from SHAKE-256 PRF.
        eta: 2 or 3.

    Returns:
        list of 256 signed integers.
    """
    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(256):
        pos_sum = 0
        neg_sum = 0
        for j in range(half):
            pos_sum += (1 if prf_bits[step * i + j] == '1' else 0)
        for j in range(half):
            neg_sum += (1 if prf_bits[step * i + half + j] == '1' else 0)
        coeffs.append(pos_sum - neg_sum)
    return coeffs


def _coeff_to_hex_12signed(val):
    """Convert signed value (range [-eta, eta]) to 12-bit hex string.

    Negative values are represented as 12-bit two's complement.
    e.g., -2 → 0xFFE → "FFE", 3 → 0x003 → "003".
    """
    masked = val & 0xFFF  # 12-bit unsigned representation
    return f'{masked:03X}'


class SampleCbdVectorGenerator(VectorGenerator):
    """Generates test vectors for the sample_cbd_sync module."""

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

        Args:
            params: dict with 'eta' key (2 or 3).

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

        # Generate random seed (256 bits) and nonce (8 bits)
        seed_bits = _random_bits(256)   # sigma, LSB-first
        nonce_bits = _random_bits(8)     # N, LSB-first

        # SHAKE-256 PRF: sigma || N → eta*64 bytes
        prf_bits = SHA_3.PRF(seed_bits, nonce_bits, eta)

        # CBD sampling: output signed 12-bit coefficients
        coeffs = _sample_poly_cbd(prf_bits, eta)

        # Convert seed and nonce to MSB-first hex for RTL input
        seed_hex = _bits_to_hex_msb(seed_bits)
        nonce_hex = _bits_to_hex_msb(nonce_bits)

        return {
            'input': {
                'seed_hex': seed_hex,
                'nonce_hex': nonce_hex,
                'eta': eta
            },
            'expected': {
                'coeffs': coeffs
            }
        }

    def write_hex_file(self, vectors: list[dict], filepath: str) -> None:
        """Write input vectors as "SEED_HEX NONCE_HEX ETA" hex format.

        Each line: "64HEXCHARS 2HEXCHARS ETA_DECIMAL".

        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:
                inp = v['input']
                f.write(f'{inp["seed_hex"]} {inp["nonce_hex"]} {inp["eta"]}\n')

    def write_expected_file(self, vectors: list[dict], filepath: str) -> None:
        """Write expected coefficients as hex strings, one per line.

        Each vector produces 256 lines of 3-char hex (12-bit signed).

        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:
                coeffs = v['expected']['coeffs']
                for c in coeffs:
                    f.write(f'{_coeff_to_hex_12signed(c)}\n')

    def compare_results(self, got: list[str], expected_file: str) -> bool:
        """Compare RTL output against expected values.

        Args:
            got: List of hex result strings from simulation.
            expected_file: Path to expected hex file.

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

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

        for i, (g, e) in enumerate(zip(got, expected)):
            if g.upper() != e.upper():
                if i < 10:  # Only show first 10 mismatches
                    print(f"  MISMATCH[{i}]: got={g.upper()}, expected={e.upper()}")
                return False

        return True