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

Generates vectors for G (SHA3-512), H (SHA3-256), J (SHAKE-256) modes
and computes expected outputs using the Python reference implementation.
"""

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")
_REF_PATH = os.path.expanduser(_REF_PATH)
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 _hex_to_bits(hex_str, num_bits):
    """Convert hex string (MSB-first) to binary string (LSB-first).

    Returns a binary string where result[0] = bit 0 (LSB).
    """
    hex_val = int(hex_str, 16)
    bits = ''
    for i in range(num_bits):
        bits += '1' if (hex_val & (1 << i)) else '0'
    return bits


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


class Sha3VectorGenerator(VectorGenerator):
    """Generates test vectors for the sha3_top module."""

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

        For G mode:
          - d: 256-bit random binary string
          - k: 8-bit random binary string
          - Input to RTL: d||k packed into 512-bit value (data_i[263:0])
          - Expected: G(d, k) = 512-bit binary string

        For H mode:
          - ek: 256-bit random binary string (single-block)
          - Input to RTL: ek in data_i[255:0]
          - Expected: H(ek) = 256-bit binary string

        For J mode:
          - z: 256-bit random binary string
          - c: 256-bit random binary string
          - Input to RTL: z||c in data_i[511:0]
          - Expected: J(z, c) = 256-bit binary string

        Args:
            params: dict with 'mode' key ('G', 'H', 'J').

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

        if mode == 'G':
            # d: 256 bits, k: 8 bits
            d = _random_bits(256)
            k = _random_bits(8)
            msg = d + k  # 264 bits, LSB first

            # Pack into 512-bit hex value: data_i[511:0]
            # pad msg to 512 bits with zeros (upper bits = 0)
            msg_512 = msg + '0' * (512 - len(msg))
            input_hex = _bits_to_hex(msg_512, 512)

            # Expected: G(d, k) returns binary string
            expected_bits = SHA_3.G(d, k)
            expected_hex = _bits_to_hex(expected_bits, 512)

            return {
                'input': {
                    'mode': 0,
                    'data': input_hex
                },
                'expected': {
                    'hash': expected_hex
                }
            }

        elif mode == 'H':
            # ek: 256-bit random message (single-block test)
            ek = _random_bits(256)

            # Pack into 512-bit hex: data_i[511:0]
            msg_512 = ek + '0' * (512 - len(ek))
            input_hex = _bits_to_hex(msg_512, 512)

            # Expected: H(ek) returns 256-bit binary string
            expected_bits = SHA_3.H(ek)
            expected_hex = _bits_to_hex(expected_bits, 256)

            return {
                'input': {
                    'mode': 1,
                    'data': input_hex
                },
                'expected': {
                    'hash': expected_hex
                }
            }

        elif mode == 'J':
            # z: 256 bits, c: 256 bits
            z = _random_bits(256)
            c = _random_bits(256)
            msg = z + c  # 512 bits

            # Pack into 512-bit hex
            input_hex = _bits_to_hex(msg, 512)

            # Expected: J(z, c) returns 256-bit binary string
            expected_bits = SHA_3.J(z, c)
            expected_hex = _bits_to_hex(expected_bits, 256)

            return {
                'input': {
                    'mode': 2,
                    'data': input_hex
                },
                'expected': {
                    'hash': expected_hex
                }
            }

        else:
            raise ValueError(f"Unknown mode: {mode}")

    def write_hex_file(self, vectors: list[dict], filepath: str) -> None:
        """Write input vectors as "MM DDDD..." hex format (mode, data).

        Each line: "MM DDDD..." where MM is hex mode (00/01/10) and DDDD...
        is the hex-encoded 512-bit message.

        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:
                mode_val = v['input']['mode']
                data_hex = v['input']['data']
                f.write(f'{mode_val:02X} {data_hex}\n')

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

        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:
                expected_hex = v['expected']['hash']
                f.write(f'{expected_hex}\n')