"""
Production-Ready Data Cleaning Pipeline
Comprehensive enterprise-grade data cleaning with error handling, validation, and monitoring
"""

import pandas as pd
import numpy as np
from typing import Dict, List, Optional, Callable, Union, Any, Tuple
from dataclasses import dataclass, field
from datetime import datetime, timezone
import logging
import json
import re
from pathlib import Path
from contextlib import contextmanager
import hashlib
from abc import ABC, abstractmethod
import math
import yaml

try:
    from sklearn.ensemble import IsolationForest
    SKLEARN_AVAILABLE = True
except ImportError:
    SKLEARN_AVAILABLE = False

from .config import (
    PipelineConfig, 
    ImputationStrategy, 
    OutlierMethod, 
    ConfigurationError
)


# ============================================================================
# CONFIGURATION
# ============================================================================
# Implementations moved to src.config.py
    

@dataclass
class CleaningMetrics:
    """Metrics collected during cleaning process"""
    rows_processed: int = 0
    rows_removed: int = 0
    columns_modified: int = 0
    errors_caught: int = 0
    warnings_generated: int = 0
    execution_time_seconds: float = 0.0
    memory_usage_mb: float = 0.0
    operations: List[str] = field(default_factory=list)
    

@dataclass
class ValidationResult:
    """Result of data validation"""
    is_valid: bool
    errors: List[str] = field(default_factory=list)
    warnings: List[str] = field(default_factory=list)
    

# ============================================================================
# CUSTOM EXCEPTIONS
# ============================================================================

class PipelineError(Exception):
    """Base exception for pipeline errors"""
    pass


class ValidationError(PipelineError):
    """Raised when data validation fails"""
    pass


class ConfigurationError(PipelineError):
    """Raised when configuration is invalid"""
    pass


class DataQualityError(PipelineError):
    """Raised when data quality issues are severe"""
    pass


# ============================================================================
# LOGGING SETUP
# ============================================================================

class JsonFormatter(logging.Formatter):
    """JSON log formatter"""
    def format(self, record):
        log_record = {
            "timestamp": self.formatTime(record, self.datefmt),
            "level": record.levelname,
            "logger": record.name,
            "message": record.getMessage()
        }
        if record.exc_info:
            log_record['exception'] = self.formatException(record.exc_info)
        return json.dumps(log_record)

def setup_logger(name: str, level: str = "INFO", log_file: Optional[Path] = None) -> logging.Logger:
    """Setup structured logger with JSON formatting"""
    logger = logging.getLogger(name)
    logger.setLevel(getattr(logging, level.upper()))
    
    if not logger.handlers:
        formatter = JsonFormatter()
        
        # Stream Handler
        stream_handler = logging.StreamHandler()
        stream_handler.setFormatter(formatter)
        logger.addHandler(stream_handler)
        
        # File Handler (if configured)
        if log_file:
            try:
                Path(log_file).parent.mkdir(parents=True, exist_ok=True)
                file_handler = logging.FileHandler(log_file)
                file_handler.setFormatter(formatter)
                logger.addHandler(file_handler)
            except Exception as e:
                print(f"Failed to setup file logging: {e}")
    
    return logger


# ============================================================================
# VALIDATORS
# ============================================================================

class DataValidator:
    """Validates data quality and schema"""
    
    @staticmethod
    def validate_dataframe(df: pd.DataFrame) -> ValidationResult:
        """Validate basic DataFrame requirements"""
        errors = []
        warnings = []
        
        if df is None:
            errors.append("DataFrame is None")
            return ValidationResult(False, errors, warnings)
        
        if not isinstance(df, pd.DataFrame):
            errors.append(f"Expected DataFrame, got {type(df)}")
            return ValidationResult(False, errors, warnings)
        
        if df.empty:
            warnings.append("DataFrame is empty")
        
        if len(df.columns) == 0:
            errors.append("DataFrame has no columns")
        
        # Check for unnamed columns
        unnamed_cols = [col for col in df.columns if str(col).startswith('Unnamed:')]
        if unnamed_cols:
            warnings.append(f"Found {len(unnamed_cols)} unnamed columns")
        
        is_valid = len(errors) == 0
        return ValidationResult(is_valid, errors, warnings)
    
    @staticmethod
    def validate_column_exists(df: pd.DataFrame, columns: Union[str, List[str]]) -> ValidationResult:
        """Validate that columns exist in DataFrame"""
        errors = []
        cols = [columns] if isinstance(columns, str) else columns
        
        missing = [col for col in cols if col not in df.columns]
        if missing:
            errors.append(f"Columns not found: {missing}")
        
        return ValidationResult(len(errors) == 0, errors)
    
    @staticmethod
    def validate_numeric_column(df: pd.DataFrame, column: str) -> ValidationResult:
        """Validate that column is numeric"""
        errors = []
        
        if column not in df.columns:
            errors.append(f"Column '{column}' not found")
        elif not pd.api.types.is_numeric_dtype(df[column]):
            errors.append(f"Column '{column}' is not numeric (dtype: {df[column].dtype})")
        
        return ValidationResult(len(errors) == 0, errors)


class SchemaValidator:
    """Validates DataFrame against a YAML schema"""
    
    def __init__(self, schema_path: Union[str, Path]):
        self.schema_path = Path(schema_path)
        self.schema = self._load_schema()
        
    def _load_schema(self) -> Dict[str, Any]:
        if not self.schema_path.exists():
            raise ConfigurationError(f"Schema file not found: {self.schema_path}")
        try:
            with open(self.schema_path, 'r') as f:
                return yaml.safe_load(f)
        except Exception as e:
            raise ConfigurationError(f"Failed to load schema: {e}")

    def validate(self, df: pd.DataFrame, schema_name: str) -> ValidationResult:
        """Validate DataFrame against a specific schema definition"""
        errors = []
        warnings = []
        
        if schema_name not in self.schema:
            raise ConfigurationError(f"Schema '{schema_name}' not found in {self.schema_path}")
            
        rules = self.schema[schema_name]
        columns = rules.get('columns', {})
        
        # Check required columns
        for col_name, col_rules in columns.items():
            if col_rules.get('required', False) and col_name not in df.columns:
                errors.append(f"Missing required column: {col_name}")
                continue
                
            if col_name in df.columns:
                # Type check (basic)
                expected_type = col_rules.get('type')
                if expected_type == 'int64' and not pd.api.types.is_integer_dtype(df[col_name]):
                    errors.append(f"Column {col_name} should be int64")
                elif expected_type == 'float64' and not pd.api.types.is_float_dtype(df[col_name]):
                    errors.append(f"Column {col_name} should be float64")
                elif expected_type == 'datetime' and not pd.api.types.is_datetime64_any_dtype(df[col_name]):
                    errors.append(f"Column {col_name} should be datetime")
                elif expected_type == 'string' and not (pd.api.types.is_string_dtype(df[col_name]) or pd.api.types.is_object_dtype(df[col_name])):
                    errors.append(f"Column {col_name} should be string")

                # Pattern check (Regex)
                if 'pattern' in col_rules:
                    pattern = col_rules['pattern']
                    if pd.api.types.is_string_dtype(df[col_name]) or pd.api.types.is_object_dtype(df[col_name]):
                        # Fill NA with empty string to avoid errors, or handle as invalid? 
                        # Usually keep NA as NA, but regex match on NA yields NaN.
                        # Using na=False means NA doesn't match.
                        matches = df[col_name].astype(str).str.match(pattern)
                        if not matches.all():
                            invalid_count = (~matches).sum()
                            errors.append(f"Column {col_name} contains {invalid_count} values not matching pattern '{pattern}'")
                    else:
                        warnings.append(f"Pattern validation skipped for non-string column {col_name}")
                    
                # Range check
                if 'min' in col_rules:
                    if df[col_name].min() < col_rules['min']:
                        errors.append(f"Column {col_name} contains values < {col_rules['min']}")
                if 'max' in col_rules:
                    if df[col_name].max() > col_rules['max']:
                        errors.append(f"Column {col_name} contains values > {col_rules['max']}")
                        
        return ValidationResult(len(errors) == 0, errors, warnings)


class NumpyEncoder(json.JSONEncoder):
    """Custom encoder for NumPy data types"""
    def default(self, obj):
        if isinstance(obj, (np.int_, np.intc, np.intp, np.int8,
                            np.int16, np.int32, np.int64, np.uint8,
                            np.uint16, np.uint32, np.uint64)):
            return int(obj)
        elif isinstance(obj, (np.float_, np.float16, np.float32, np.float64)):
            return float(obj)
        elif isinstance(obj, (np.ndarray,)):
            return obj.tolist()
        return json.JSONEncoder.default(self, obj)


# ============================================================================
# MAIN PIPELINE CLASS
# ============================================================================

class DataCleaningPipeline:
    """
    Production-ready data cleaning pipeline with comprehensive error handling,
    validation, monitoring, and rollback capabilities.
    """
    
    def __init__(
        self,
        df: pd.DataFrame,
        config: Optional[PipelineConfig] = None,
        logger: Optional[logging.Logger] = None
    ):
        """Initialize the pipeline with a DataFrame and optional configuration"""
        
        # Configuration
        self.config = config or PipelineConfig()
        self.logger = logger or setup_logger(__name__, self.config.log_level, self.config.log_file)
        
        # Validate input
        validation = DataValidator.validate_dataframe(df)
        if not validation.is_valid:
            raise ValidationError(f"Invalid DataFrame: {validation.errors}")
        
        for warning in validation.warnings:
            self.logger.warning(warning)
        
        # State management
        self.df = df.copy()
        self.original_df = df.copy()
        self._history: List[pd.DataFrame] = []
        self._checkpoint_data()
        
        # Metrics and reporting
        self.metrics = CleaningMetrics()
        self.cleaning_report: List[str] = []
        # Metrics and reporting
        self.metrics = CleaningMetrics()
        self.cleaning_report: List[str] = []
        self._start_time = datetime.now(timezone.utc)
        
        # Data fingerprint for integrity
        self._original_hash = self._compute_hash(df)
        
        self.logger.info(
            f"Pipeline initialized: {df.shape[0]} rows, {df.shape[1]} columns"
        )
    
    # ========================================================================
    # HELPER METHODS
    # ========================================================================
    
    def _checkpoint_data(self) -> None:
        """Create a checkpoint of current state"""
        self._history.append(self.df.copy())
        
        # 1. Enforce max_checkpoints count
        while len(self._history) > self.config.max_checkpoints:
            self._history.pop(0)

        # 2. Enforce max_memory_mb limit
        # Calculate total memory of history
        try:
            total_mem_bytes = sum(df.memory_usage(deep=True).sum() for df in self._history)
            max_mem_bytes = self.config.max_memory_mb * 1024 * 1024
            
            if total_mem_bytes > max_mem_bytes:
                self.logger.warning(
                    f"Memory usage ({total_mem_bytes/1024/1024:.2f} MB) exceeds limit "
                    f"({self.config.max_memory_mb} MB). Dropping oldest checkpoints."
                )
                
                # Keep dropping until under limit or only 1 left
                while total_mem_bytes > max_mem_bytes and len(self._history) > 1:
                    removed_df = self._history.pop(0)
                    total_mem_bytes -= removed_df.memory_usage(deep=True).sum()
                    
        except Exception as e:
            self.logger.warning(f"Failed to check memory usage: {e}")
    
    def _compute_hash(self, df: pd.DataFrame) -> str:
        """Compute hash of DataFrame for integrity checks"""
        return hashlib.md5(
            pd.util.hash_pandas_object(df, index=True).values
        ).hexdigest()
    
    def _log_operation(self, operation: str, details: str = "") -> None:
        """Log an operation with details"""
        msg = f"{operation}: {details}" if details else operation
        self.cleaning_report.append(msg)
        self.metrics.operations.append(operation)
        self.logger.info(msg)
    
    def _validate_path(self, path: Union[str, Path], must_exist: bool = False) -> Path:
        """Validate file path to prevent traversal attacks"""
        try:
            # Resolve to absolute path
            path_obj = Path(path).resolve()
            cwd = Path.cwd().resolve()
            
            # Check if path is within current working directory (or specific allowed dirs)
            if not str(path_obj).startswith(str(cwd)):
                self.logger.warning(f"Path traversal detected (outside CWD): {path}")
                # Enforce CWD restriction for security
                if not self.config.allow_file_operations:
                     raise PipelineError(f"Security: File operations restricted to CWD. Access denied: {path}")

            # Explicit check for ".." in the string before resolution to catch obvious attempts
            if '..' in str(path):
                 self.logger.warning(f"Constructed path contains '..': {path}")

            if must_exist and not path_obj.exists():
                raise FileNotFoundError(f"File not found: {path_obj}")
                
            return path_obj
            
        except Exception as e:
            if isinstance(e, (PipelineError, FileNotFoundError)):
                raise
            raise PipelineError(f"Invalid path: {str(e)}")


    def _update_metrics(self, rows_removed: int = 0, cols_modified: int = 0) -> None:
        """Update pipeline metrics"""
        self.metrics.rows_removed += rows_removed
        self.metrics.columns_modified += cols_modified
        self.metrics.rows_processed = len(self.df)
    
    @contextmanager
    def _safe_operation(self, operation_name: str):
        """Context manager for safe operations with rollback"""
        checkpoint = self.df.copy()
        try:
            yield
            self._checkpoint_data()
        except Exception as e:
            self.df = checkpoint
            self.metrics.errors_caught += 1
            self.logger.error(f"{operation_name} failed: {str(e)}")
            if self.config.strict_mode:
                # If it's already a PipelineError, re-raise it to preserve specificity
                if isinstance(e, PipelineError):
                    raise
                raise PipelineError(f"{operation_name} failed: {str(e)}") from e
            else:
                self._log_operation(f"{operation_name} (FAILED)", str(e))
    
    # ========================================================================
    # CORE CLEANING METHODS
    # ========================================================================
    
    def remove_duplicates(
        self,
        subset: Optional[List[str]] = None,
        keep: str = 'first'
    ) -> 'DataCleaningPipeline':
        """
        Remove duplicate rows with validation
        
        Args:
            subset: Columns to consider for identifying duplicates
            keep: Which duplicates to keep ('first', 'last', False)
        """
        with self._safe_operation("remove_duplicates"):
            # Validate subset columns
            if subset:
                validation = DataValidator.validate_column_exists(self.df, subset)
                if not validation.is_valid:
                    raise ValidationError(f"Invalid columns: {validation.errors}")
            
            initial_rows = len(self.df)
            
            try:
                self.df = self.df.drop_duplicates(subset=subset, keep=keep)
            except Exception as e:
                raise PipelineError(f"Failed to remove duplicates: {str(e)}") from e
            
            removed = initial_rows - len(self.df)
            self._update_metrics(rows_removed=removed)
            self._log_operation(
                "remove_duplicates",
                f"Removed {removed} duplicate rows ({removed/initial_rows*100:.1f}%)"
            )
        
        return self
    
    def handle_missing_values(
        self,
        strategy: Dict[str, Union[str, Tuple[str, Any]]]
    ) -> 'DataCleaningPipeline':
        """
        Handle missing values with multiple strategies
        
        Args:
            strategy: Dict mapping column names to imputation strategies
                     Examples: {'age': 'median', 'name': 'drop', 'score': ('constant', 0)}
        """
        with self._safe_operation("handle_missing_values"):
            for col, method in strategy.items():
                # Validate column exists
                if col not in self.df.columns:
                    self.logger.warning(f"Column '{col}' not found, skipping")
                    continue
                
                missing_before = self.df[col].isna().sum()
                if missing_before == 0:
                    continue
                
                try:
                    self._apply_imputation(col, method)
                    missing_after = self.df[col].isna().sum()
                    
                    self._log_operation(
                        f"handle_missing[{col}]",
                        f"Handled {missing_before - missing_after} missing values using {method}"
                    )
                except Exception as e:
                    self.logger.error(f"Failed to handle missing values in '{col}': {str(e)}")
                    if self.config.strict_mode:
                        raise
        
        return self
    
    def _apply_imputation(self, col: str, method: Union[str, Tuple[str, Any]]) -> None:
        """Apply specific imputation method to a column"""
        # Handle constant imputation
        if isinstance(method, tuple) and method[0] == ImputationStrategy.CONSTANT.value:
            self.df[col] = self.df[col].fillna(method[1])
            return
        
        method_str = method if isinstance(method, str) else method[0]
        
        if method_str == ImputationStrategy.DROP.value:
            self.df = self.df.dropna(subset=[col])
        
        elif method_str == ImputationStrategy.MEAN.value:
            if not pd.api.types.is_numeric_dtype(self.df[col]):
                raise ValidationError(f"Cannot compute mean for non-numeric column '{col}'")
            self.df[col] = self.df[col].fillna(self.df[col].mean())
        
        elif method_str == ImputationStrategy.MEDIAN.value:
            if not pd.api.types.is_numeric_dtype(self.df[col]):
                raise ValidationError(f"Cannot compute median for non-numeric column '{col}'")
            self.df[col] = self.df[col].fillna(self.df[col].median())
        
        elif method_str == ImputationStrategy.MODE.value:
            mode_values = self.df[col].mode()
            if len(mode_values) > 0:
                self.df[col] = self.df[col].fillna(mode_values[0])
            else:
                self.logger.warning(f"No mode found for column '{col}'")
        
        elif method_str == ImputationStrategy.FORWARD_FILL.value:
            self.df[col] = self.df[col].ffill()
        
        elif method_str == ImputationStrategy.BACKWARD_FILL.value:
            self.df[col] = self.df[col].bfill()
        
        elif method_str == ImputationStrategy.INTERPOLATE.value:
            if not pd.api.types.is_numeric_dtype(self.df[col]):
                raise ValidationError(f"Cannot interpolate non-numeric column '{col}'")
            self.df[col] = self.df[col].interpolate()
        
        else:
            raise ConfigurationError(f"Unknown imputation strategy: {method_str}")
    
    def remove_outliers(
        self,
        columns: List[str],
        method: str = 'iqr',
        threshold: float = 1.5
    ) -> 'DataCleaningPipeline':
        """
        Remove outliers using various methods
        
        Args:
            columns: List of columns to check for outliers
            method: Detection method ('iqr', 'zscore', 'modified_zscore')
            threshold: Threshold for outlier detection
        """
        with self._safe_operation("remove_outliers"):
            initial_rows = len(self.df)
            mask = pd.Series(True, index=self.df.index)
            
            for col in columns:
                # Validate column
                validation = DataValidator.validate_numeric_column(self.df, col)
                if not validation.is_valid:
                    self.logger.warning(f"Skipping non-numeric column '{col}'")
                    continue
                
                try:
                    col_mask = self._detect_outliers(col, method, threshold)
                    mask &= ~col_mask
                except Exception as e:
                    self.logger.error(f"Failed to detect outliers in '{col}': {str(e)}")
                    if self.config.strict_mode:
                        raise
            
            self.df = self.df[mask]
            removed = initial_rows - len(self.df)
            
            self._update_metrics(rows_removed=removed)
            self._log_operation(
                "remove_outliers",
                f"Removed {removed} outliers using {method} method (threshold={threshold})"
            )
        
        return self
    
    def _detect_outliers(self, col: str, method: str, threshold: float) -> pd.Series:
        """Detect outliers in a column"""
        if method == OutlierMethod.IQR.value:
            Q1 = self.df[col].quantile(0.25)
            Q3 = self.df[col].quantile(0.75)
            IQR = Q3 - Q1
            lower = Q1 - threshold * IQR
            upper = Q3 + threshold * IQR
            self.logger.debug(f"IQR outlier detection for {col}: Q1={Q1}, Q3={Q3}, IQR={IQR}, threshold={threshold}")
            mask = (self.df[col] < lower) | (self.df[col] > upper)
            self.logger.debug(f"Found {mask.sum()} outliers in {col}")
            return mask
        
        elif method == OutlierMethod.ZSCORE.value:
            std = self.df[col].std()
            if std == 0:
                self.logger.warning(f"Standard deviation is 0 for column {col}, cannot compute Z-score")
                return pd.Series([False] * len(self.df), index=self.df.index)
            z_scores = np.abs((self.df[col] - self.df[col].mean()) / std)
            return z_scores > threshold
        
        elif method == OutlierMethod.MODIFIED_ZSCORE.value:
            median = self.df[col].median()
            mad = np.median(np.abs(self.df[col] - median))
            if mad == 0:
                self.logger.warning(f"MAD is 0 for column {col}, cannot compute modified z-score")
                return pd.Series([False] * len(self.df))
                
            modified_z = 0.6745 * (self.df[col] - median) / mad
            return np.abs(modified_z) > threshold
            
        elif method == OutlierMethod.ISOLATION_FOREST.value:
            if not SKLEARN_AVAILABLE:
                raise ConfigurationError("scikit-learn not installed, cannot use IsolationForest")
            
            # Reshape for sklearn
            X = self.df[col].values.reshape(-1, 1)
            clf = IsolationForest(contamination='auto', random_state=42)
            preds = clf.fit_predict(X)
            # IsolationForest returns -1 for outliers
            return pd.Series(preds == -1, index=self.df.index)
        
        else:
            raise ConfigurationError(f"Unknown outlier detection method: {method}")
    
    def standardize_text(
        self,
        columns: List[str],
        lowercase: bool = True,
        strip: bool = True,
        remove_special_chars: bool = False,
        remove_extra_spaces: bool = True
    ) -> 'DataCleaningPipeline':
        """Standardize text columns with multiple options"""
        with self._safe_operation("standardize_text"):
            for col in columns:
                if col not in self.df.columns:
                    self.logger.warning(f"Column '{col}' not found, skipping")
                    continue
                
                try:
                    # Convert to string
                    self.df[col] = self.df[col].astype(str)
                    
                    # Apply transformations
                    # Apply transformations using vectorized string operations
                    if strip:
                        self.df[col] = self.df[col].str.strip()
                    
                    if lowercase:
                        self.df[col] = self.df[col].str.lower()
                    
                    if remove_special_chars:
                        self.df[col] = self.df[col].str.replace(r'[^a-zA-Z0-9\s]', '', regex=True)
                    
                    if remove_extra_spaces:
                        self.df[col] = self.df[col].str.replace(r'\s+', ' ', regex=True)
                    
                    self._update_metrics(cols_modified=1)
                    
                except Exception as e:
                    self.logger.error(f"Failed to standardize text in '{col}': {str(e)}")
                    if self.config.strict_mode:
                        raise
            
            self._log_operation("standardize_text", f"Standardized {len(columns)} text columns")
        
        return self
    
    def validate_email(
        self,
        column: str,
        remove_invalid: bool = True,
        add_validation_column: bool = False
    ) -> 'DataCleaningPipeline':
        """
        Validate email addresses with improved regex
        
        Args:
            column: Column containing email addresses
            remove_invalid: Remove rows with invalid emails
            add_validation_column: Add boolean column for validation results
        """
        with self._safe_operation("validate_email"):
            if column not in self.df.columns:
                raise ValidationError(f"Column '{column}' not found")
            
            # Improved email regex (still basic, consider using email-validator library)
            # CACHED REGEX PATTERN at function scope? No, use re.compile at module level or static
            # Define as class attribute or use cache. Let's use internal cache.
            if not hasattr(self, '_email_pattern'):
                 self._email_pattern = re.compile(r'^[a-zA-Z0-9][a-zA-Z0-9._%+-]*@[a-zA-Z0-9][a-zA-Z0-9.-]*\.[a-zA-Z]{2,}$')
            
            # Handle NaN values
            email_series = self.df[column].fillna('')
            valid_mask = email_series.astype(str).str.match(self._email_pattern, na=False)
            
            invalid_count = (~valid_mask).sum()
            
            if add_validation_column:
                self.df[f'{column}_valid'] = valid_mask
            
            if remove_invalid:
                initial_rows = len(self.df)
                self.df = self.df[valid_mask]
                self._update_metrics(rows_removed=invalid_count)
            
            self._log_operation(
                "validate_email",
                f"Found {invalid_count} invalid emails in '{column}'"
            )
        
        return self
    
    def convert_data_types(
        self,
        type_mapping: Dict[str, str],
        errors: str = 'coerce'
    ) -> 'DataCleaningPipeline':
        """
        Convert column data types with error handling
        
        Args:
            type_mapping: Dict mapping column names to target types
            errors: How to handle conversion errors ('raise', 'coerce', 'ignore')
        """
        with self._safe_operation("convert_data_types"):
            for col, dtype in type_mapping.items():
                if col not in self.df.columns:
                    self.logger.warning(f"Column '{col}' not found, skipping")
                    continue
                
                try:
                    if dtype == 'datetime':
                        self.df[col] = pd.to_datetime(self.df[col], errors=errors)
                    elif dtype in ['int', 'int64', 'int32']:
                        self.df[col] = pd.to_numeric(self.df[col], errors=errors)
                        if errors != 'coerce' or self.df[col].notna().all():
                            self.df[col] = self.df[col].astype(dtype)
                    else:
                        self.df[col] = self.df[col].astype(dtype)
                    
                    self._update_metrics(cols_modified=1)
                    self._log_operation(f"convert_type[{col}]", f"Converted to {dtype}")
                    
                except Exception as e:
                    self.logger.error(f"Failed to convert '{col}' to {dtype}: {str(e)}")
                    if self.config.strict_mode:
                        raise
        
        return self
    
    # ========================================================================
    # ADVANCED FEATURES
    # ========================================================================
    
    def rollback(self, steps: int = 1) -> 'DataCleaningPipeline':
        """Rollback pipeline to previous state"""
        # Calculate target index: we want to go back 'steps' from the CURRENT state
        # Current state is at history[-1]
        # Previous state is at history[-2] (steps=1)
        # So we need history[-(steps + 1)]
        
        target_idx = -(steps + 1)
        if abs(target_idx) > len(self._history):
            raise PipelineError(f"Cannot rollback {steps} steps, only {len(self._history)-1} previous states available")
        
        self.df = self._history[target_idx].copy()
        self._log_operation("rollback", f"Rolled back {steps} step(s)")
        return self
    
    def apply_custom_function(
        self,
        column: str,
        func: Callable,
        new_column: Optional[str] = None,
        handle_errors: bool = True
    ) -> 'DataCleaningPipeline':
        """Apply custom function with error handling"""
        if not self.config.allow_custom_functions:
            raise ConfigurationError("Custom functions are disabled by configuration for security")
            
        with self._safe_operation("apply_custom_function"):
            if column not in self.df.columns:
                raise ValidationError(f"Column '{column}' not found")
            
            # SECURITY WARNING: Log every execution of arbitrary code
            func_name = getattr(func, '__name__', str(func))
            self.logger.warning(
                f"SECURITY: Executing custom function '{func_name}' on column '{column}'. "
                f"Ensure this function is from a trusted source."
            )
            
            target_col = new_column if new_column else column
            
            if handle_errors:
                def safe_func(x):
                    try:
                        return func(x)
                    except Exception as e:
                        self.logger.warning(f"Error applying function to value {x}: {str(e)}")
                        return x
                self.df[target_col] = self.df[column].apply(safe_func)
            else:
                self.df[target_col] = self.df[column].apply(func)
            
            self._update_metrics(cols_modified=1)
            self._log_operation("apply_custom_function", f"Applied custom function '{func_name}' to '{column}'")
        
        return self
    
    def reset_index(self) -> 'DataCleaningPipeline':
        """Reset DataFrame index"""
        self.df = self.df.reset_index(drop=True)
        return self
        
    def process_in_chunks(self, filepath: Union[str, Path], chunk_size: Optional[int] = None) -> Any:
        """
        Process a large file in chunks.
        Yields a DataCleaningPipeline instance for each chunk.
        """
        size = chunk_size or self.config.chunk_size or 10000
        
        # Determine file type
        fp = Path(filepath)
        if fp.suffix == '.csv':
            reader = pd.read_csv(fp, chunksize=size)
            for i, chunk in enumerate(reader):
                self.logger.info(f"Processing chunk {i+1} with {len(chunk)} rows")
                # Create a mini-pipeline for this chunk
                # Note: This yields a NEW pipeline instance, doesn't update self.df (streaming mode)
                chunk_pipeline = DataCleaningPipeline(chunk, config=self.config, logger=self.logger)
                yield chunk_pipeline
        else:
            raise ConfigurationError(f"Chunking currently only supported for CSV, got {fp.suffix}")
    
    # ========================================================================
    # REPORTING & EXPORT
    # ========================================================================
    
    def get_cleaning_report(self) -> List[str]:
        """Get detailed cleaning report"""
        return self.cleaning_report.copy()
    
    def get_summary(self) -> Dict[str, Any]:
        """Get comprehensive summary statistics"""
        return {
            'original_shape': self.original_df.shape,
            'cleaned_shape': self.df.shape,
            'rows_removed': self.metrics.rows_removed,
            'rows_remaining': len(self.df),
            'columns_modified': self.metrics.columns_modified,
            'missing_values_before': self.original_df.isna().sum().sum(),
            'missing_values_after': self.df.isna().sum().sum(),
            'duplicate_rows_remaining': self.df.duplicated().sum(),
            'errors_caught': self.metrics.errors_caught,
            'operations_performed': len(self.metrics.operations),
            'execution_time_seconds': (datetime.now(timezone.utc) - self._start_time).total_seconds(),
            'data_quality_score': self._calculate_quality_score()
        }
    
    def _calculate_quality_score(self) -> float:
        """Calculate data quality score (0-100)"""
        score = 100.0
        
        # Penalize for missing values
        total_cells = self.df.shape[0] * self.df.shape[1]
        if total_cells > 0:
            missing_ratio = self.df.isna().sum().sum() / total_cells
            score -= (missing_ratio * 30)
        
        # Penalize for duplicates
        if len(self.df) > 0:
            duplicate_ratio = self.df.duplicated().sum() / len(self.df)
            score -= (duplicate_ratio * 20)
        
        # Penalize for errors during processing
        if self.metrics.errors_caught > 0:
            score -= min(self.metrics.errors_caught * 5, 20)
        
        return max(0.0, min(100.0, score))
    
        return max(0.0, min(100.0, score))
    
    def validate_schema(self, schema_path: Union[str, Path], schema_name: str) -> 'DataCleaningPipeline':
        """Validate data against a schema file"""
        with self._safe_operation("validate_schema"):
            validator = SchemaValidator(schema_path)
            result = validator.validate(self.df, schema_name)
            
            if not result.is_valid:
                msg = f"Schema validation failed: {result.errors}"
                self.logger.error(msg)
                if self.config.strict_mode:
                    raise ValidationError(msg)
                else:
                    self.logger.warning("Continuing despite schema validation errors")
            
            self._log_operation("validate_schema", f"Validated against {schema_name}")
            
        return self

    def export_report(self, filepath: Union[str, Path]) -> None:
        """Export detailed cleaning report to JSON"""
        report = {
            'summary': self.get_summary(),
            'operations': self.cleaning_report,
            'data_profile': {
                'columns': list(self.df.columns),
                'dtypes': {col: str(dtype) for col, dtype in self.df.dtypes.items()},
                'missing_values': self.df.isna().sum().to_dict(),
                'shape': self.df.shape
            },
            'timestamp': datetime.now(timezone.utc).isoformat()
        }
        
        
        with open(filepath, 'w') as f:
            json.dump(report, f, indent=2, cls=NumpyEncoder)
        
        self.logger.info(f"Report exported to {filepath}")
    
    def get_cleaned_data(self) -> pd.DataFrame:
        """Return the cleaned DataFrame"""
        return self.df.copy()
    
    
    def to_csv(self, filepath: Union[str, Path], **kwargs) -> None:
        """Export cleaned data to CSV"""
        path = self._validate_path(filepath)
        self.df.to_csv(path, **kwargs)
        self.logger.info(f"Data exported to {path}")
        
    def to_parquet(self, filepath: Union[str, Path], **kwargs) -> None:
        """Export cleaned data to Parquet"""
        path = self._validate_path(filepath)
        # Ensure fastparquet or pyarrow is installed
        try:
            self.df.to_parquet(path, **kwargs)
            self.logger.info(f"Data exported to {path}")
        except ImportError:
            self.logger.error("pyarrow or fastparquet required for parquet export")
            raise

    def to_json(self, filepath: Union[str, Path], **kwargs) -> None:
        """Export cleaned data to JSON"""
        path = self._validate_path(filepath)
        # Use simple json export via pandas with date handling
        self.df.to_json(path, date_format='iso', default_handler=str, **kwargs)
        self.logger.info(f"Data exported to {path}")


# ============================================================================
# EXAMPLE USAGE
# ============================================================================

if __name__ == "__main__":
    # Create sample data
    data = {
        'name': ['John Doe', 'jane smith', '  Bob  ', 'Alice', 'John Doe', None],
        'age': [25, 30, 150, 28, 25, 35],
        'email': ['john@ex.com', 'bad-email', 'bob@test.com', 'alice@co.com', 'john@ex.com', 'test@test.com'],
        'salary': [50000, 60000, None, 70000, 50000, 80000],
    }
    
    df = pd.DataFrame(data)
    
    print("Original Data:")
    print(df)
    print("\n" + "="*80 + "\n")
    
    # Create pipeline with configuration
    # Create pipeline with configuration
    config = PipelineConfig(log_level="INFO", strict_mode=False, allow_custom_functions=True)
    pipeline = DataCleaningPipeline(df, config=config)
    
    # Execute cleaning
    cleaned_df = (pipeline
        .remove_duplicates()
        .handle_missing_values({'salary': 'median', 'name': 'drop'})
        .remove_outliers(['age'], method='iqr')
        .standardize_text(['name'], lowercase=True)
        .validate_email('email', remove_invalid=True)
        .reset_index()
        .get_cleaned_data())
    
    print("Cleaned Data:")
    print(cleaned_df)
    print("\n" + "="*80 + "\n")
    
    # Print report
    print("Cleaning Report:")
    for step in pipeline.get_cleaning_report():
        print(f"  • {step}")
    
    print("\n" + "="*80 + "\n")
    print("Summary:")
    for k, v in pipeline.get_summary().items():
        print(f"  {k}: {v}")