"""
Integration Test Harness base class that encapsulates common test logic.
Reduces boilerplate code and provides a consistent testing interface.
"""

import os
import shutil
import subprocess
import sys
from pathlib import Path
from typing import Any, Dict, Optional, Tuple

from .utils import set_test_seeds

# Add paths to sys.path for proper imports
sys.path.append("./")
sys.path.append("./data_engine")


# Import data_engine modules only when needed to avoid import errors during framework testing
def _import_data_engine_modules():
    """Import data_engine modules when they are actually needed."""
    try:
        from nimbus import run_data_engine
        from nimbus.utils.config_processor import ConfigProcessor
        from nimbus.utils.utils import init_env

        return init_env, run_data_engine, ConfigProcessor
    except ImportError as e:
        raise ImportError(f"Failed to import data_engine modules. Ensure dependencies are installed: {e}")


class IntegrationTestHarness:
    """
    Base class for integration tests that provides common functionality.

    This class encapsulates:
    - Configuration loading and processing
    - Output directory cleanup
    - Test pipeline execution (direct or subprocess)
    - Output validation
    - Data comparison with reference
    """

    def __init__(self, config_path: str, seed: int = 42, load_num: int = 0, random_num: int = 0, episodes: int = 0):
        """
        Initialize the test harness with configuration and seed.

        Args:
            config_path: Path to the test configuration YAML file
            seed: Random seed for reproducible results (default: 42)
        """
        self.config_path = config_path
        self.seed = seed
        self.load_num = load_num
        self.random_num = random_num
        self.episodes = episodes
        self.output_dir = None

        # Import data_engine modules
        self._init_env, self._run_data_engine, self._ConfigProcessor = _import_data_engine_modules()
        self.processor = self._ConfigProcessor()

        # Initialize environment (same as launcher.py)
        self._init_env()

        self.config = self.load_and_process_config() if config_path else None
        # Set random seeds for reproducibility
        self.modify_config()

        set_test_seeds(seed)

        from nimbus.utils.flags import set_debug_mode, set_random_seed

        set_debug_mode(True)  # Enable debug mode for better error visibility during tests
        set_random_seed(seed)

    def modify_config(self):
        """Modify configuration parameters as needed before running the pipeline."""
        if self.config and "load_stage" in self.config:
            if "layout_random_generator" in self.config.load_stage:
                if "args" in self.config.load_stage.layout_random_generator:
                    if self.random_num > 0:
                        self.config.load_stage.layout_random_generator.args.random_num = self.random_num
                    if "input_dir" in self.config.load_stage.layout_random_generator.args:
                        if "simbox" in self.config.name:
                            input_path = (
                                "/shared/smartbot_new/zhangyuchang/CI/manip/"
                                "simbox/simbox_plan_ci/seq_path/BananaBaseTask/plan"
                            )
                            self.config.load_stage.layout_random_generator.args.input_dir = input_path
        if self.config and "plan_stage" in self.config:
            if "seq_planner" in self.config.plan_stage:
                if "args" in self.config.plan_stage.seq_planner:
                    if self.episodes > 0:
                        self.config.plan_stage.seq_planner.args.episodes = self.episodes
                        if self.load_num > 0:
                            self.config.load_stage.scene_loader.args.load_num = self.load_num
        if self.config and "name" in self.config:
            self.config.name = self.config.name + "_ci"
        if self.config and "store_stage" in self.config:
            if hasattr(self.config.store_stage.writer.args, "obs_output_dir"):
                self.config.store_stage.writer.args.obs_output_dir = f"output/{self.config.name}/obs_path/"
            if hasattr(self.config.store_stage.writer.args, "seq_output_dir"):
                self.config.store_stage.writer.args.seq_output_dir = f"output/{self.config.name}/seq_path/"
            if hasattr(self.config.store_stage.writer.args, "output_dir"):
                self.config.store_stage.writer.args.output_dir = f"output/{self.config.name}/"
            self._extract_output_dir()

    def load_and_process_config(self) -> Dict[str, Any]:
        """
        Load and process the test configuration file.

        Returns:
            Processed configuration dictionary

        Raises:
            AssertionError: If config file does not exist
        """
        assert os.path.exists(self.config_path), f"Config file not found: {self.config_path}"

        self.config = self.processor.process_config(self.config_path)
        self.processor.print_final_config(self.config)

        # Extract output directory from config
        self._extract_output_dir()

        return self.config

    def _extract_output_dir(self):
        """Extract and expand the output directory path from config."""
        # Try navigation test output path
        output_dir = self.config.get("store_stage", {}).get("writer", {}).get("args", {}).get("seq_output_dir")

        # If not found, try common output_dir used in render configs
        if not output_dir:
            output_dir = self.config.get("store_stage", {}).get("writer", {}).get("args", {}).get("output_dir")

        # Process the output directory if found
        if output_dir and isinstance(output_dir, str):
            name = self.config.get("name", "test_output")
            output_dir = output_dir.replace("${name}", name)
            self.output_dir = os.path.abspath(output_dir)

    def cleanup_output_directory(self):
        """Clean up existing output directory if it exists."""
        if self.output_dir and os.path.exists(self.output_dir):
            # Use ignore_errors=True to handle NFS caching issues
            shutil.rmtree(self.output_dir, ignore_errors=True)
            # If directory still exists (NFS delay), try removing with onerror handler
            if os.path.exists(self.output_dir):

                def handle_remove_error(func, path, exc_info):  # pylint: disable=W0613
                    """Handle errors during removal, with retry for NFS issues."""
                    import time

                    time.sleep(0.1)  # Brief delay for NFS sync
                    try:
                        if os.path.isdir(path):
                            os.rmdir(path)
                        else:
                            os.remove(path)
                    except OSError:
                        pass  # Ignore if still fails

                shutil.rmtree(self.output_dir, onerror=handle_remove_error)
            print(f"Cleaned up existing output directory: {self.output_dir}")

    def run_data_engine_direct(self, **kwargs) -> None:
        """
        Run the test pipeline directly in the current process.

        Args:
            **kwargs: Additional arguments to pass to run_data_engine

        Raises:
            Exception: If pipeline execution fails
        """
        if not self.config:
            self.load_and_process_config()
        self.modify_config()
        self.cleanup_output_directory()

        try:
            self._run_data_engine(self.config, **kwargs)
        except Exception as e:
            raise e

    def run_data_engine_subprocess(
        self,
        runner_script: str,
        interpreter: str = "python",
        timeout: int = 1800,
        compare_output: bool = False,
        reference_dir: str = "",
        comparator: str = "simbox",
        comparator_args: Optional[Dict[str, Any]] = None,
    ) -> subprocess.CompletedProcess:
        """
        Run the test pipeline in a subprocess using a runner script.

        Args:
            runner_script: Path to the runner script to execute
            interpreter: Command to use for running the script (default: "python")
            timeout: Timeout in seconds for subprocess execution (default: 1800)
            compare_output: Whether to compare generated output with a reference directory
            reference_dir: Path to reference directory containing meta_info.pkl and lmdb
            comparator: Which comparator to use (default: "simbox")
            comparator_args: Optional extra arguments for comparator (e.g. {"tolerance": 1e-6, "image_psnr": 37.0})

        Returns:
            subprocess.CompletedProcess object with execution results

        Raises:
            subprocess.TimeoutExpired: If execution exceeds timeout
            AssertionError: If subprocess returns non-zero exit code
        """
        self.cleanup_output_directory()

        # Build command based on interpreter type
        if interpreter == "blenderproc":
            cmd = ["blenderproc", "run", runner_script]
            if os.environ.get("BLENDER_CUSTOM_PATH"):
                cmd.extend(["--custom-blender-path", os.environ["BLENDER_CUSTOM_PATH"]])
        elif interpreter.endswith(".sh"):
            # For scripts like /isaac-sim/python.sh
            cmd = [interpreter, runner_script]
        else:
            cmd = [interpreter, runner_script]

        if not self.output_dir:
            self._extract_output_dir()
        output_dir = self.output_dir

        print(f"Running command: {' '.join(cmd)}")
        print(f"Expected output directory: {output_dir}")

        try:
            result = subprocess.run(cmd, capture_output=False, text=True, timeout=timeout, check=False)

            # Print subprocess output for debugging
            if result.stdout:
                print("STDOUT:", result.stdout[-2000:])  # Last 2000 chars
            if result.stderr:
                print("STDERR:", result.stderr[-1000:])  # Last 1000 chars
            print("Return code:", result.returncode)

            if compare_output and result.returncode == 0:
                if not output_dir:
                    raise RuntimeError(
                        "Output directory not configured. Expected one of "
                        "store_stage.writer.args.(seq_output_dir|obs_output_dir|output_dir)."
                    )

                if output_dir:
                    for root, dirs, files in os.walk(output_dir):
                        if "lmdb" in dirs and "meta_info.pkl" in files:
                            output_dir = root
                            break

                if os.path.exists(reference_dir):
                    # Find the reference render directory
                    for root, dirs, files in os.walk(reference_dir):
                        if "lmdb" in dirs and "meta_info.pkl" in files:
                            reference_dir = root
                            break

                # Build comparator command according to requested comparator
                comp = (comparator or "simbox").lower()
                comparator_args = comparator_args or {}

                if comp == "simbox":
                    compare_cmd = [
                        "/isaac-sim/python.sh",
                        "tests/integration/data_comparators/simbox_comparator.py",
                        "--dir1",
                        output_dir,
                        "--dir2",
                        reference_dir,
                    ]
                    # Optional numeric/image thresholds
                    if "tolerance" in comparator_args and comparator_args["tolerance"] is not None:
                        compare_cmd += ["--tolerance", str(comparator_args["tolerance"])]
                    if "image_psnr" in comparator_args and comparator_args["image_psnr"] is not None:
                        compare_cmd += ["--image-psnr", str(comparator_args["image_psnr"])]

                    print(f"Running comparison: {' '.join(compare_cmd)}")
                    compare_result = subprocess.run(
                        compare_cmd, capture_output=True, text=True, timeout=600, check=False
                    )

                    print("Comparison STDOUT:")
                    print(compare_result.stdout)
                    print("Comparison STDERR:")
                    print(compare_result.stderr)

                    if compare_result.returncode != 0:
                        raise RuntimeError("Simbox comparison failed: outputs differ")

                    if "Successfully loaded data from both directories" in compare_result.stdout:
                        print("✓ Both output directories have valid structure (meta_info.pkl + lmdb)")

                    print("✓ Simbox render test completed with numeric-aligned comparison")

                else:
                    raise ValueError(f"Unknown comparator: {comp}. Use 'simbox'.")

            return result

        except subprocess.TimeoutExpired:
            raise RuntimeError(f"Test timed out after {timeout} seconds")

    def validate_output_generated(self, min_files: int = 1) -> bool:
        """
        Validate that output was generated in the expected directory.

        Args:
            min_files: Minimum number of files expected in output (default: 1)

        Returns:
            True if validation passes

        Raises:
            AssertionError: If output directory doesn't exist or has too few files
        """
        assert self.output_dir is not None, "Output directory not configured"
        assert os.path.exists(self.output_dir), f"Expected output directory was not created: {self.output_dir}"

        output_files = list(Path(self.output_dir).rglob("*"))
        assert (
            len(output_files) >= min_files
        ), f"Expected at least {min_files} files but found {len(output_files)} in: {self.output_dir}"

        print(f"✓ Pipeline completed successfully. Output generated in: {self.output_dir}")
        print(f"✓ Generated {len(output_files)} files/directories")
        return True

    def compare_with_reference(self, reference_dir: str, comparator_func, **comparator_kwargs) -> Tuple[bool, str]:
        """
        Compare generated output with reference data using provided comparator function.

        Args:
            reference_dir: Path to reference data directory
            comparator_func: Function to use for comparison
            **comparator_kwargs: Additional arguments for comparator function

        Returns:
            Tuple of (success: bool, message: str)
        """
        if not os.path.exists(reference_dir):
            print(f"Reference directory not found, skipping comparison: {reference_dir}")
            return True, "Reference data not available"

        success, message = comparator_func(self.output_dir, reference_dir, **comparator_kwargs)

        if success:
            print(f"✓ Results match reference data: {message}")
        else:
            print(f"✗ Result comparison failed: {message}")

        return success, message

    def run_test_end_to_end(
        self,
        reference_dir: Optional[str] = None,
        comparator_func=None,
        comparator_kwargs: Optional[Dict] = None,
        min_output_files: int = 1,
        **pipeline_kwargs,
    ) -> bool:
        """
        Run a complete end-to-end test including comparison with reference.

        Args:
            reference_dir: Optional path to reference data for comparison
            comparator_func: Optional function for comparing results
            comparator_kwargs: Optional kwargs for comparator function
            min_output_files: Minimum expected output files
            **pipeline_kwargs: Additional arguments for pipeline execution

        Returns:
            True if test passes, False otherwise

        Raises:
            Exception: If any test step fails
        """
        # Load configuration
        self.load_and_process_config()
        # Run pipeline
        self.run_data_engine_direct(**pipeline_kwargs, master_seed=self.seed)

        # Validate output
        self.validate_output_generated(min_files=min_output_files)

        # Compare with reference if provided
        if reference_dir and comparator_func:
            comparator_kwargs = comparator_kwargs or {}
            success, message = self.compare_with_reference(reference_dir, comparator_func, **comparator_kwargs)
            assert success, f"Comparison with reference failed: {message}"

        return True