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

Generates streaming 256-coefficient polynomial add/sub test vectors.
Each vector: "MODE A0..A255 B0..B255" (513 space-separated hex values).
Expected: "C0 C1 ... C255" (256 space-separated hex values).
"""

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'))

from vector_gen import VectorGenerator

Q = 3329  # ML-KEM prime modulus
N = 256   # polynomial degree


class PolyArithVectorGenerator(VectorGenerator):
    """Generates test vectors for the poly_arith_sync module."""

    def generate_one(self, params: dict) -> dict:
        """Generate one polynomial add or sub vector.

        Generates 256 random coefficients for A and B, computes expected
        result as (A[i] +/- B[i]) mod Q.

        Args:
            params: dict with optional 'mode' key ('add' or 'sub').

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

        A = [random.randint(0, Q - 1) for _ in range(N)]
        B = [random.randint(0, Q - 1) for _ in range(N)]

        if mode == 'sub':
            expected = [(a - b) % Q for a, b in zip(A, B)]
        else:
            expected = [(a + b) % Q for a, b in zip(A, B)]

        return {
            'input': {'mode': mode, 'A': A, 'B': B},
            'expected': {'C': expected}
        }

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

        Format: "MODE A0 A1 ... A255 B0 B1 ... B255"
        MODE is 'A' for add, 'S' for sub.

        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']
                mode_char = 'A' if inp['mode'] == 'add' else 'S'
                A_hex = ' '.join(f'{c:03X}' for c in inp['A'])
                B_hex = ' '.join(f'{c:03X}' for c in inp['B'])
                f.write(f'{mode_char} {A_hex} {B_hex}\n')

    def write_expected_file(self, vectors: list[dict], filepath: str) -> None:
        """Write expected results as hex file.

        Format: "C0 C1 ... C255" (one space-separated line per vector).

        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:
                C = v['expected']['C']
                C_hex = ' '.join(f'{c:03X}' for c in C)
                f.write(f'{C_hex}\n')

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

        Args:
            got: List of result strings from simulation.
                 Each is " C0 C1 ... C255" after "RESULT:" prefix.
            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):
            return False

        for i, (g, e) in enumerate(zip(got, expected)):
            if g.strip() != e.strip():
                return False

        return True