"""
Evaluation Runner for pipeline execution with timing.

Handles:
- Pre-build hooks (Rust compilation, bash scripts, etc.)
- Dynamic module import (baseline vs iteration)
- Per-detector timing instrumentation
- Pattern output collection
"""

import importlib
import importlib.util
import subprocess
import sys
import time
from dataclasses import dataclass, field
from datetime import datetime, timedelta
from pathlib import Path
from typing import Dict, List, Any, Optional, Tuple

import pandas as pd
from loguru import logger


@dataclass
class BuildResult:
    """Result from running a build command."""
    success: bool
    command: str
    stdout: str
    stderr: str
    duration_seconds: float
    return_code: int


@dataclass
class RunResult:
    """Result from running an iteration."""
    iteration: str
    success: bool
    patterns_df: Optional[pd.DataFrame]
    features_df: Optional[pd.DataFrame]
    total_seconds: float
    detector_timing: Dict[str, float]
    build_result: Optional[BuildResult] = None
    error: Optional[str] = None


class EvalRunner:
    """
    Runs evaluation iterations with timing instrumentation.

    Supports:
    - Pre-build hooks for custom builds (Rust, bash scripts, etc.)
    - Dynamic module loading for different iterations
    - Per-detector timing collection
    """

    def __init__(
        self,
        input_path: Path,
        output_base_path: Path,
        network: str = "torus",
    ):
        """
        Initialize the evaluation runner.

        Args:
            input_path: Path to input data (e.g., data/input/torus/2025-12-10/256)
            output_base_path: Base path for outputs (e.g., data/output/torus/2025-12-10/256)
            network: Network name (torus, bittensor, etc.)
        """
        self.input_path = Path(input_path)
        self.output_base_path = Path(output_base_path)
        self.network = network

        # Extract metadata from input path
        self.window_days = int(self.input_path.name)
        self.processing_date = self.input_path.parent.name

    def run_build_hook(
        self,
        build_command: str,
        iteration: str,
        working_dir: Optional[Path] = None,
        timeout: int = 300,
    ) -> BuildResult:
        """
        Run a pre-iteration build command.

        This allows LLM iterations to:
        - Compile Rust code with `cargo build --release`
        - Run custom bash scripts
        - Install dependencies
        - Generate code

        Args:
            build_command: Shell command to execute
            iteration: Iteration name (for logging)
            working_dir: Working directory for command (defaults to project root)
            timeout: Command timeout in seconds

        Returns:
            BuildResult with success status and output
        """
        if not build_command:
            return BuildResult(
                success=True,
                command="",
                stdout="",
                stderr="",
                duration_seconds=0.0,
                return_code=0,
            )

        logger.info(f"Running build hook for {iteration}: {build_command}")

        start_time = time.time()
        cwd = working_dir or Path.cwd()

        try:
            result = subprocess.run(
                build_command,
                shell=True,
                cwd=str(cwd),
                capture_output=True,
                text=True,
                timeout=timeout,
            )

            duration = time.time() - start_time

            build_result = BuildResult(
                success=result.returncode == 0,
                command=build_command,
                stdout=result.stdout,
                stderr=result.stderr,
                duration_seconds=duration,
                return_code=result.returncode,
            )

            if build_result.success:
                logger.info(f"Build hook completed in {duration:.2f}s")
            else:
                logger.error(f"Build hook failed (code {result.returncode}): {result.stderr[:500]}")

            return build_result

        except subprocess.TimeoutExpired:
            duration = time.time() - start_time
            logger.error(f"Build hook timed out after {timeout}s")
            return BuildResult(
                success=False,
                command=build_command,
                stdout="",
                stderr=f"Command timed out after {timeout} seconds",
                duration_seconds=duration,
                return_code=-1,
            )
        except Exception as e:
            duration = time.time() - start_time
            logger.error(f"Build hook exception: {e}")
            return BuildResult(
                success=False,
                command=build_command,
                stdout="",
                stderr=str(e),
                duration_seconds=duration,
                return_code=-1,
            )

    def run_iteration(
        self,
        iteration: str,
        build_command: Optional[str] = None,
        build_timeout: int = 300,
    ) -> RunResult:
        """
        Run a single iteration with timing.

        Args:
            iteration: Iteration name ("baseline" or "v_XXXX")
            build_command: Optional pre-build command to run
            build_timeout: Timeout for build command

        Returns:
            RunResult with patterns, timing, and status
        """
        logger.info(f"Starting iteration: {iteration}")

        # Run build hook if specified
        build_result = None
        if build_command:
            build_result = self.run_build_hook(
                build_command=build_command,
                iteration=iteration,
                timeout=build_timeout,
            )
            if not build_result.success:
                return RunResult(
                    iteration=iteration,
                    success=False,
                    patterns_df=None,
                    features_df=None,
                    total_seconds=build_result.duration_seconds,
                    detector_timing={},
                    build_result=build_result,
                    error=f"Build failed: {build_result.stderr[:200]}",
                )

        # Determine module to import
        if iteration == "baseline":
            module_name = "chainswarm_analyzers_baseline"
        else:
            module_name = f"chainswarm_analyzers_{iteration}"

        try:
            # Run the pipeline with timing
            start_time = time.time()

            patterns_df, features_df, detector_timing = self._run_pipeline(
                module_name=module_name,
                iteration=iteration,
            )

            total_seconds = time.time() - start_time

            return RunResult(
                iteration=iteration,
                success=True,
                patterns_df=patterns_df,
                features_df=features_df,
                total_seconds=total_seconds,
                detector_timing=detector_timing,
                build_result=build_result,
            )

        except Exception as e:
            total_seconds = time.time() - start_time if 'start_time' in locals() else 0.0
            logger.error(f"Iteration {iteration} failed: {e}")
            return RunResult(
                iteration=iteration,
                success=False,
                patterns_df=None,
                features_df=None,
                total_seconds=total_seconds,
                detector_timing={},
                build_result=build_result,
                error=str(e),
            )

    def _run_pipeline(
        self,
        module_name: str,
        iteration: str,
    ) -> Tuple[pd.DataFrame, pd.DataFrame, Dict[str, float]]:
        """
        Run the analysis pipeline for a specific module.

        Args:
            module_name: Python module name to import
            iteration: Iteration name for output path

        Returns:
            Tuple of (patterns_df, features_df, detector_timing)
        """
        # Import the module dynamically
        try:
            module = importlib.import_module(module_name)
            importlib.reload(module)  # Ensure fresh import
        except ImportError as e:
            raise ImportError(f"Cannot import module {module_name}: {e}")

        # Import required components
        from_module = lambda name: getattr(module, name, None)

        # Try to get components from the module
        ParquetAdapter = from_module('ParquetAdapter')
        SettingsLoader = from_module('SettingsLoader')
        create_pipeline = from_module('create_pipeline')

        # Fallback to baseline if module doesn't have all components
        if ParquetAdapter is None:
            from chainswarm_analyzers_baseline.adapters.parquet import ParquetAdapter
        if SettingsLoader is None:
            from chainswarm_analyzers_baseline.config import SettingsLoader
        if create_pipeline is None:
            from chainswarm_analyzers_baseline.pipeline import create_pipeline

        # Set up output path for this iteration
        output_path = self.output_base_path / iteration
        output_path.mkdir(parents=True, exist_ok=True)

        # Calculate time range
        processing_dt = datetime.strptime(self.processing_date, "%Y-%m-%d")
        end_timestamp_ms = int(processing_dt.timestamp() * 1000)
        start_timestamp_ms = int(
            (processing_dt - timedelta(days=self.window_days)).timestamp() * 1000
        )

        # Create adapter and pipeline
        adapter = ParquetAdapter(
            input_path=self.input_path,
            output_path=output_path,
        )

        settings_loader = SettingsLoader()

        pipeline = create_pipeline(
            adapter=adapter,
            network=self.network,
            settings_loader=settings_loader,
        )

        # Run with timing instrumentation
        detector_timing = {}

        # Check if pattern_analyzer has timing support
        pattern_analyzer = pipeline.pattern_analyzer
        if hasattr(pattern_analyzer, 'analyze_with_timing'):
            # Use timing-aware method
            money_flows = adapter.read_money_flows(start_timestamp_ms, end_timestamp_ms)
            addresses = self._extract_addresses_from_flows(money_flows)
            address_labels = adapter.read_address_labels(addresses)
            timestamp_data = adapter.read_transfer_timestamps(
                start_timestamp_ms, end_timestamp_ms, addresses
            )

            patterns, detector_timing = pattern_analyzer.analyze_with_timing(
                money_flows=money_flows,
                address_labels=address_labels,
                timestamp_data=timestamp_data,
                window_days=self.window_days,
                processing_date=self.processing_date,
            )

            adapter.write_patterns(patterns, self.window_days, self.processing_date)
        else:
            # Fallback to standard pipeline run
            result = pipeline.run(
                start_timestamp_ms=start_timestamp_ms,
                end_timestamp_ms=end_timestamp_ms,
                window_days=self.window_days,
                processing_date=self.processing_date,
                run_features=True,
                run_patterns=True,
            )

        # Load output files
        patterns_df = self._load_patterns(output_path)
        features_df = self._load_features(output_path)

        return patterns_df, features_df, detector_timing

    def _extract_addresses_from_flows(self, flows: List[Dict[str, Any]]) -> List[str]:
        """Extract unique addresses from money flows."""
        addresses = set()
        for flow in flows:
            addresses.add(flow['from_address'])
            addresses.add(flow['to_address'])
        return list(addresses)

    def _load_patterns(self, output_path: Path) -> pd.DataFrame:
        """Load all pattern files from output directory."""
        pattern_files = list(output_path.glob("patterns_*.parquet"))

        if not pattern_files:
            return pd.DataFrame()

        dfs = []
        for f in pattern_files:
            try:
                df = pd.read_parquet(f, engine='pyarrow')
                dfs.append(df)
            except Exception as e:
                logger.warning(f"Failed to load {f}: {e}")

        if not dfs:
            return pd.DataFrame()

        return pd.concat(dfs, ignore_index=True)

    def _load_features(self, output_path: Path) -> pd.DataFrame:
        """Load features file from output directory."""
        features_file = output_path / "features.parquet"

        if not features_file.exists():
            return pd.DataFrame()

        try:
            return pd.read_parquet(features_file, engine='pyarrow')
        except Exception as e:
            logger.warning(f"Failed to load features: {e}")
            return pd.DataFrame()

    def load_ground_truth(self) -> pd.DataFrame:
        """Load ground truth from input directory."""
        gt_file = self.input_path / "ground_truth.parquet"

        if not gt_file.exists():
            raise FileNotFoundError(f"Ground truth not found: {gt_file}")

        return pd.read_parquet(gt_file, engine='pyarrow')

    def load_transfers(self) -> pd.DataFrame:
        """Load transfers from input directory."""
        transfers_file = self.input_path / "transfers.parquet"

        if not transfers_file.exists():
            raise FileNotFoundError(f"Transfers not found: {transfers_file}")

        return pd.read_parquet(transfers_file, engine='pyarrow')
