"""test_runner.py - Main test engine.

Discovers modules, loads test plans, generates vectors, runs simulations,
and collects results.
"""

import os
import json
import time
import importlib.util
import sys
from pathlib import Path

from .sim_controller import SimController
from .result_checker import check_results


class TestRunner:
    """Main test engine for the ML-KEM RTL test framework."""

    def __init__(self, config_path: str = 'test_framework/config.json'):
        """Initialize with configuration.

        Args:
            config_path: Path to config.json.
        """
        self.project_root = os.path.dirname(os.path.dirname(
            os.path.dirname(os.path.abspath(__file__))))
        self.config_path = os.path.join(self.project_root, config_path)
        with open(self.config_path) as f:
            self.config = json.load(f)
        self.sim = SimController(self.config)
        self.modules_dir = os.path.join(self.project_root,
                                        'test_framework', 'modules')
        self.rtl_root = os.path.join(self.project_root,
                                     self.config.get('rtl_root', 'sync_rtl'))

    def discover_modules(self) -> dict[str, dict]:
        """Scan test_framework/modules/*/test_plan.json and return plans.

        Returns:
            dict of module_name -> test_plan.
        """
        modules = {}
        if not os.path.isdir(self.modules_dir):
            return modules

        for entry in os.listdir(self.modules_dir):
            mod_dir = os.path.join(self.modules_dir, entry)
            plan_path = os.path.join(mod_dir, 'test_plan.json')
            if os.path.isdir(mod_dir) and os.path.exists(plan_path):
                with open(plan_path) as f:
                    plan = json.load(f)
                modules[entry] = plan
        return modules

    def _get_rtl_files(self, plan: dict) -> list[str]:
        """Resolve RTL source file paths from a test plan.

        Args:
            plan: Test plan dict with 'rtl_top' and optional 'rtl_deps'.

        Returns:
            List of absolute paths: deps first, then top.
        """
        rtl_top = os.path.join(self.project_root, plan['rtl_top'])
        rtl_deps = [os.path.join(self.project_root, d)
                    for d in plan.get('rtl_deps', [])]
        return rtl_deps + [rtl_top]

    def run_module(self, module_name: str, quick: bool = False,
                   single_case: str = None) -> dict:
        """Run all test cases for a module.

        Args:
            module_name: Name of the module to test.
            quick: If True, reduce vector count for faster testing.
            single_case: If set, only run this case ID.

        Returns:
            dict with keys: pass, fail, cases, elapsed, status.
        """
        modules = self.discover_modules()
        if module_name not in modules:
            return {
                'pass': 0, 'fail': 0, 'cases': {},
                'elapsed': 0.0, 'status': 'NOT_FOUND',
                'error': f'Module "{module_name}" not found'
            }

        plan = modules[module_name]
        mod_dir = os.path.join(self.modules_dir, module_name)
        vectors_dir = os.path.join(mod_dir, 'vectors')
        work_dir = os.path.join(mod_dir, 'obj_dir')
        os.makedirs(vectors_dir, exist_ok=True)

        # Import module's gen_vectors.py
        gen_path = os.path.join(mod_dir, 'gen_vectors.py')
        if not os.path.exists(gen_path):
            return {
                'pass': 0, 'fail': 0, 'cases': {},
                'elapsed': 0.0, 'status': 'ERROR',
                'error': f'gen_vectors.py not found for {module_name}'
            }

        spec = importlib.util.spec_from_file_location(
            f'{module_name}_gen', gen_path)
        gen_module = importlib.util.module_from_spec(spec)
        sys.modules[f'{module_name}_gen'] = gen_module
        spec.loader.exec_module(gen_module)

        # Find generator class
        generator = None
        for name in dir(gen_module):
            obj = getattr(gen_module, name)
            if (isinstance(obj, type) and
                    hasattr(obj, 'generate_one') and
                    hasattr(obj, 'write_hex_file') and
                    name != 'VectorGenerator'):
                generator = obj()
                break

        if generator is None:
            return {
                'pass': 0, 'fail': 0, 'cases': {},
                'elapsed': 0.0, 'status': 'ERROR',
                'error': f'No generator class found in {gen_path}'
            }

        # Resolve RTL files and top module name
        rtl_files = self._get_rtl_files(plan)
        tb_cpp = os.path.join(self.project_root, plan['tb_cpp'])

        # Top module name derived from rtl_top filename (strip .v)
        top_module = os.path.basename(plan['rtl_top']).replace('.v', '')

        # Compile once for the module
        print(f"\n  Compiling {module_name} with Verilator...")
        start_time = time.time()

        try:
            executable = self.sim.compile(
                top_module=top_module,
                rtl_files=rtl_files,
                tb_cpp=tb_cpp,
                work_dir=work_dir,
                inc_dirs=[self.project_root]
            )
        except RuntimeError as e:
            return {
                'pass': 0, 'fail': 0, 'cases': {},
                'elapsed': time.time() - start_time,
                'status': 'COMPILE_ERROR',
                'error': str(e)
            }

        case_results = {}
        total_pass = 0
        total_fail = 0

        for case in plan.get('cases', []):
            case_id = case['id']

            # Filter by single case
            if single_case and case_id != single_case:
                continue

            num_vectors = case.get('num_vectors', 10)
            if quick:
                num_vectors = max(1, num_vectors // 5)

            params = case.get('params', {})
            print(f"  Case [{case_id}]: {case.get('description', '')} "
                  f"({num_vectors} vectors)")

            # Generate vectors
            vectors = []
            for i in range(num_vectors):
                vec = generator.generate_one(params)
                vectors.append(vec)

            # Write input vectors hex file
            input_file = os.path.join(vectors_dir, f'{case_id}_input.hex')
            generator.write_hex_file(vectors, input_file)

            # Write expected output hex file
            expected_file = os.path.join(vectors_dir, f'{case_id}_expected.hex')
            generator.write_expected_file(vectors, expected_file)

            # Run simulation
            timeout_s = plan.get('timeout_s', 30)
            run_result = self.sim.run(executable, input_file, timeout_s)

            if not run_result['success']:
                case_results[case_id] = {
                    'pass': 0, 'fail': num_vectors,
                    'elapsed': run_result['elapsed'],
                    'error': run_result.get('stderr', 'Simulation failed')
                }
                total_fail += num_vectors
                continue

            # Compare results
            tolerance = case.get('tolerance', 'bit_exact')
            if tolerance == 'bit_exact':
                num_pass, num_fail = check_results(
                    run_result['results'], expected_file)
            else:
                # Delegate to generator for custom comparison
                if hasattr(generator, 'compare_results'):
                    compare_ok = generator.compare_results(
                        run_result['results'], expected_file)
                    num_pass = num_vectors if compare_ok else 0
                    num_fail = 0 if compare_ok else num_vectors
                else:
                    num_pass, num_fail = check_results(
                        run_result['results'], expected_file)

            case_results[case_id] = {
                'pass': num_pass,
                'fail': num_fail,
                'elapsed': run_result['elapsed']
            }
            total_pass += num_pass
            total_fail += num_fail

            # Clean up vectors unless keep_vectors is true
            if not self.config.get('keep_vectors', False):
                try:
                    os.remove(input_file)
                    os.remove(expected_file)
                except OSError:
                    pass

        elapsed = time.time() - start_time

        return {
            'pass': total_pass,
            'fail': total_fail,
            'cases': case_results,
            'elapsed': elapsed,
            'status': 'PASS' if total_fail == 0 else 'FAIL'
        }
