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

Generates (a, b) input pairs and computes expected (a + b) mod Q.
"""

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


class ModAddVectorGenerator(VectorGenerator):
    """Generates test vectors for the mod_add_sync module."""

    def generate_one(self, params: dict) -> dict:
        """Generate a single (a, b) pair with expected modular sum.

        Uses random values covering:
        - No overflow: a + b < Q
        - With overflow: a + b >= Q
        - Edge cases: a=0, b=0, a=Q-1, b=Q-1

        Args:
            params: Unused for mod_add basic case.

        Returns:
            dict with 'input' and 'expected' keys.
        """
        # Mix of random and edge cases for coverage
        roll = random.random()

        if roll < 0.05:
            # Both zero
            a, b = 0, 0
        elif roll < 0.10:
            # a=0, b random
            a = 0
            b = random.randint(0, Q - 1)
        elif roll < 0.15:
            # b=0, a random
            a = random.randint(0, Q - 1)
            b = 0
        elif roll < 0.20:
            # Both Q-1 (max overflow)
            a, b = Q - 1, Q - 1
        elif roll < 0.25:
            # a random, b = Q-1
            a = random.randint(0, Q - 1)
            b = Q - 1
        elif roll < 0.65:
            # No overflow: both small
            a = random.randint(0, Q // 2 - 1)
            b = random.randint(0, Q // 2 - 1)
        else:
            # Potential overflow
            a = random.randint(0, Q - 1)
            b = random.randint(0, Q - 1)

        # Compute expected: (a + b) mod Q
        expected_sum = (a + b) % Q

        return {
            'input': {'a': a, 'b': b},
            'expected': {'sum': expected_sum}
        }

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

        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:
                a = v['input']['a']
                b = v['input']['b']
                f.write(f'{a:03X} {b:03X}\n')

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

        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:
                s = v['expected']['sum']
                f.write(f'{s:03X}\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):
            return False

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

        return True