"""
Preprocessor Executors

Executors for data preprocessing operations.
"""

from __future__ import annotations

from typing import Any, TYPE_CHECKING

import pandas as pd
from sklearn.preprocessing import (
    StandardScaler,
    MinMaxScaler,
    RobustScaler,
    Normalizer,
    LabelEncoder,
    OneHotEncoder,
)
from sklearn.model_selection import train_test_split

from .base import NodeExecutor, executor_function

if TYPE_CHECKING:
    from ..schemas import Node
    from ..storage import StorageManager


class DataScalerExecutor(NodeExecutor):
    """Executor for scaling/normalizing data."""
    
    SCALER_MAP = {
        "standard": StandardScaler,
        "minmax": MinMaxScaler,
        "robust": RobustScaler,
        "normalizer": Normalizer,
    }
    
    def execute(
        self,
        node: "Node",
        input_data: dict[str, Any],
        storage: "StorageManager",
        run_id: str,
    ) -> dict[str, Any]:
        """
        Scale numerical columns in a DataFrame.
        
        Parameters:
            - scaler_type: Type of scaler (standard, minmax, robust, normalizer)
            - columns: List of columns to scale (empty = all numeric)
            - fit_on_train_only: Whether to fit only on training data
        """
        params = node.data.params
        
        # Get input - could be DataFrame or dict with train/test split
        input_value = self.get_single_input(input_data)
        
        if isinstance(input_value, dict) and "X_train" in input_value:
            # Already split data
            return self._scale_split_data(input_value, params, storage, run_id, node.id)
        elif isinstance(input_value, pd.DataFrame):
            # Single DataFrame
            return self._scale_dataframe(input_value, params, storage, run_id, node.id)
        else:
            raise ValueError(f"Unexpected input type: {type(input_value)}")
    
    def _scale_dataframe(
        self,
        df: pd.DataFrame,
        params: dict,
        storage: "StorageManager",
        run_id: str,
        node_id: str,
    ) -> pd.DataFrame:
        """Scale a single DataFrame."""
        scaler_type = params.get("scaler_type", "standard")
        columns = params.get("columns", [])
        
        # Get scaler class
        scaler_class = self.SCALER_MAP.get(scaler_type, StandardScaler)
        scaler = scaler_class()
        
        # Determine columns to scale
        if not columns:
            columns = df.select_dtypes(include=["number"]).columns.tolist()
        
        # Scale the data
        df_scaled = df.copy()
        if columns:
            df_scaled[columns] = scaler.fit_transform(df[columns])
        
        # Save the scaler
        storage.save_artifact(
            run_id=run_id,
            artifact_name=f"{node_id}_scaler",
            artifact=scaler,
        )
        
        storage.log_event(
            run_id=run_id,
            node_id=node_id,
            event="data_scaled",
            details={
                "scaler_type": scaler_type,
                "columns_scaled": columns,
            },
        )
        
        return df_scaled
    
    def _scale_split_data(
        self,
        data: dict[str, Any],
        params: dict,
        storage: "StorageManager",
        run_id: str,
        node_id: str,
    ) -> dict[str, Any]:
        """Scale already-split train/test data."""
        scaler_type = params.get("scaler_type", "standard")
        columns = params.get("columns", [])
        
        scaler_class = self.SCALER_MAP.get(scaler_type, StandardScaler)
        scaler = scaler_class()
        
        X_train = data["X_train"]
        X_test = data["X_test"]
        
        # Determine columns to scale
        if not columns:
            columns = X_train.select_dtypes(include=["number"]).columns.tolist()
        
        # Fit on train, transform both
        result = data.copy()
        if columns:
            result["X_train"] = X_train.copy()
            result["X_test"] = X_test.copy()
            result["X_train"][columns] = scaler.fit_transform(X_train[columns])
            result["X_test"][columns] = scaler.transform(X_test[columns])
        
        # Save scaler
        storage.save_artifact(
            run_id=run_id,
            artifact_name=f"{node_id}_scaler",
            artifact=scaler,
        )
        
        storage.log_event(
            run_id=run_id,
            node_id=node_id,
            event="split_data_scaled",
            details={
                "scaler_type": scaler_type,
                "columns_scaled": columns,
            },
        )
        
        return result


class TrainTestSplitExecutor(NodeExecutor):
    """Executor for splitting data into train/test sets."""
    
    def execute(
        self,
        node: "Node",
        input_data: dict[str, Any],
        storage: "StorageManager",
        run_id: str,
    ) -> dict[str, Any]:
        """
        Split a DataFrame into train and test sets.
        
        Parameters:
            - test_size: Proportion for test set (default: 0.2)
            - random_state: Random seed (default: 42)
            - stratify_column: Column to stratify by (optional)
            - shuffle: Whether to shuffle before split (default: True)
            - target_column: Column to separate as y (optional)
            - drop_columns: List of columns to exclude from features (optional)
        """
        params = node.data.params
        
        df = self.get_single_input(input_data)
        if not isinstance(df, pd.DataFrame):
            raise ValueError("TrainTestSplit requires a DataFrame input")
        
        test_size = params.get("test_size", 0.2)
        random_state = params.get("random_state", 42)
        stratify_column = params.get("stratify_column")
        shuffle = params.get("shuffle", True)
        target_column = params.get("target_column")
        drop_columns = params.get("drop_columns", [])
        
        # Drop specified columns first
        if drop_columns:
            df = df.drop(columns=[col for col in drop_columns if col in df.columns])
        
        # Prepare stratify array
        stratify = None
        if stratify_column and stratify_column in df.columns:
            stratify = df[stratify_column]
        
        # Split features and target if specified
        if target_column and target_column in df.columns:
            X = df.drop(columns=[target_column])
            y = df[target_column]
            
            X_train, X_test, y_train, y_test = train_test_split(
                X, y,
                test_size=test_size,
                random_state=random_state,
                shuffle=shuffle,
                stratify=stratify if stratify is not None else (y if stratify_column == target_column else None),
            )
            
            result = {
                "X_train": X_train,
                "X_test": X_test,
                "y_train": y_train,
                "y_test": y_test,
            }
        else:
            # Just split the DataFrame
            train_df, test_df = train_test_split(
                df,
                test_size=test_size,
                random_state=random_state,
                shuffle=shuffle,
                stratify=stratify,
            )
            
            result = {
                "train": train_df,
                "test": test_df,
            }
        
        storage.log_event(
            run_id=run_id,
            node_id=node.id,
            event="train_test_split",
            details={
                "test_size": test_size,
                "train_samples": len(result.get("X_train", result.get("train"))),
                "test_samples": len(result.get("X_test", result.get("test"))),
            },
        )
        
        return result


class MissingValueHandlerExecutor(NodeExecutor):
    """Executor for handling missing values."""
    
    def execute(
        self,
        node: "Node",
        input_data: dict[str, Any],
        storage: "StorageManager",
        run_id: str,
    ) -> pd.DataFrame:
        """
        Handle missing values in a DataFrame.
        
        Parameters:
            - strategy: How to handle missing values (drop, mean, median, mode, constant)
            - fill_value: Value to use if strategy is "constant"
            - columns: Columns to process (empty = all)
        """
        params = node.data.params
        
        df = self.get_single_input(input_data)
        if not isinstance(df, pd.DataFrame):
            raise ValueError("MissingValueHandler requires a DataFrame input")
        
        strategy = params.get("strategy", "drop")
        fill_value = params.get("fill_value", 0)
        columns = params.get("columns", [])
        
        df_result = df.copy()
        cols_to_process = columns if columns else df_result.columns.tolist()
        
        missing_before = df_result[cols_to_process].isnull().sum().sum()
        
        if strategy == "drop":
            df_result = df_result.dropna(subset=cols_to_process)
        elif strategy == "mean":
            for col in cols_to_process:
                if df_result[col].dtype in ["float64", "int64"]:
                    df_result[col].fillna(df_result[col].mean(), inplace=True)
        elif strategy == "median":
            for col in cols_to_process:
                if df_result[col].dtype in ["float64", "int64"]:
                    df_result[col].fillna(df_result[col].median(), inplace=True)
        elif strategy == "mode":
            for col in cols_to_process:
                df_result[col].fillna(df_result[col].mode().iloc[0], inplace=True)
        elif strategy == "constant":
            df_result[cols_to_process] = df_result[cols_to_process].fillna(fill_value)
        
        missing_after = df_result[cols_to_process].isnull().sum().sum()
        
        storage.log_event(
            run_id=run_id,
            node_id=node.id,
            event="missing_values_handled",
            details={
                "strategy": strategy,
                "missing_before": int(missing_before),
                "missing_after": int(missing_after),
            },
        )
        
        return df_result


class EncoderExecutor(NodeExecutor):
    """Executor for encoding categorical variables."""
    
    def execute(
        self,
        node: "Node",
        input_data: dict[str, Any],
        storage: "StorageManager",
        run_id: str,
    ) -> pd.DataFrame:
        """
        Encode categorical columns.
        
        Parameters:
            - encoding_type: Type of encoding (label, onehot)
            - columns: Columns to encode (empty = all object columns)
            - drop_first: For one-hot, whether to drop first column (default: False)
        """
        params = node.data.params
        
        df = self.get_single_input(input_data)
        if not isinstance(df, pd.DataFrame):
            raise ValueError("Encoder requires a DataFrame input")
        
        encoding_type = params.get("encoding_type", "label")
        columns = params.get("columns", [])
        drop_first = params.get("drop_first", False)
        
        df_result = df.copy()
        
        # Determine columns to encode
        if not columns:
            columns = df_result.select_dtypes(include=["object", "category"]).columns.tolist()
        
        encoded_columns = []
        
        if encoding_type == "label":
            for col in columns:
                if col in df_result.columns:
                    le = LabelEncoder()
                    df_result[col] = le.fit_transform(df_result[col].astype(str))
                    encoded_columns.append(col)
                    
                    # Save encoder
                    storage.save_artifact(
                        run_id=run_id,
                        artifact_name=f"{node.id}_{col}_encoder",
                        artifact=le,
                    )
        
        elif encoding_type == "onehot":
            df_result = pd.get_dummies(
                df_result,
                columns=columns,
                drop_first=drop_first,
            )
            encoded_columns = columns
        
        storage.log_event(
            run_id=run_id,
            node_id=node.id,
            event="data_encoded",
            details={
                "encoding_type": encoding_type,
                "columns_encoded": encoded_columns,
            },
        )
        
        return df_result


# Functional wrappers
data_scaler_executor = executor_function(DataScalerExecutor)
train_test_split_executor = executor_function(TrainTestSplitExecutor)
missing_value_handler_executor = executor_function(MissingValueHandlerExecutor)
encoder_executor = executor_function(EncoderExecutor)
