"""
Comprehensive Test Suite for Data Cleaning Pipeline
Tests cover unit, integration, edge cases, and performance
"""

import pytest
import pandas as pd
import numpy as np
from datetime import datetime
from DataCleaningPipeline.pipeline import (
    DataCleaningPipeline,
    ValidationError,
    PipelineError,
    DataValidator
)
from DataCleaningPipeline.config import (
    PipelineConfig,
    ImputationStrategy,
    OutlierMethod
)


# ============================================================================
# FIXTURES
# ============================================================================

@pytest.fixture
def sample_df():
    """Create a sample DataFrame for testing"""
    return pd.DataFrame({
        'id': [1, 2, 3, 4, 5],
        'name': ['Alice', 'Bob', 'Charlie', 'David', 'Eve'],
        'age': [25, 30, 35, 40, 45],
        'salary': [50000, 60000, 70000, 80000, 90000],
        'email': ['alice@test.com', 'bob@test.com', 'charlie@test.com', 
                  'david@test.com', 'eve@test.com']
    })


@pytest.fixture
def dirty_df():
    """Create a DataFrame with quality issues"""
    return pd.DataFrame({
        'name': ['John', 'Jane', '  Bob  ', 'Alice', 'John', None, 'Charlie'],
        'age': [25, 30, 150, 28, 25, 35, -5],
        'email': ['john@test.com', 'invalid', 'bob@test.com', 'alice@test.com',
                  'john@test.com', 'test@test.com', 'charlie@test.com'],
        'salary': [50000, 60000, None, 70000, 50000, 80000, 1000000],
        'city': ['NYC', 'LA', 'Chicago!', 'Houston', 'NYC', 'Phoenix', 'Seattle']
    })


@pytest.fixture
def config():
    """Create a test configuration"""
    return PipelineConfig(log_level="ERROR", strict_mode=False)


# ============================================================================
# VALIDATOR TESTS
# ============================================================================

class TestDataValidator:
    """Test DataValidator class"""
    
    def test_validate_valid_dataframe(self, sample_df):
        """Test validation of valid DataFrame"""
        result = DataValidator.validate_dataframe(sample_df)
        assert result.is_valid
        assert len(result.errors) == 0
    
    def test_validate_none_dataframe(self):
        """Test validation of None"""
        result = DataValidator.validate_dataframe(None)
        assert not result.is_valid
        assert len(result.errors) > 0
    
    def test_validate_empty_dataframe(self):
        """Test validation of empty DataFrame"""
        df = pd.DataFrame()
        result = DataValidator.validate_dataframe(df)
        assert not result.is_valid or len(result.warnings) > 0
    
    def test_validate_column_exists(self, sample_df):
        """Test column existence validation"""
        result = DataValidator.validate_column_exists(sample_df, 'name')
        assert result.is_valid
        
        result = DataValidator.validate_column_exists(sample_df, 'nonexistent')
        assert not result.is_valid
    
    def test_validate_numeric_column(self, sample_df):
        """Test numeric column validation"""
        result = DataValidator.validate_numeric_column(sample_df, 'age')
        assert result.is_valid
        
        result = DataValidator.validate_numeric_column(sample_df, 'name')
        assert not result.is_valid


# ============================================================================
# PIPELINE INITIALIZATION TESTS
# ============================================================================

class TestPipelineInitialization:
    """Test pipeline initialization"""
    
    def test_init_with_valid_df(self, sample_df, config):
        """Test initialization with valid DataFrame"""
        pipeline = DataCleaningPipeline(sample_df, config)
        assert pipeline.df.equals(sample_df)
        assert len(pipeline.cleaning_report) == 0
    
    def test_init_with_none_raises_error(self, config):
        """Test that None DataFrame raises error"""
        with pytest.raises(ValidationError):
            DataCleaningPipeline(None, config)
    
    def test_init_creates_copy(self, sample_df, config):
        """Test that initialization creates a copy"""
        pipeline = DataCleaningPipeline(sample_df, config)
        pipeline.df.loc[0, 'name'] = 'Modified'
        assert sample_df.loc[0, 'name'] != 'Modified'
    
    def test_init_without_config(self, sample_df):
        """Test initialization without config uses defaults"""
        pipeline = DataCleaningPipeline(sample_df)
        assert pipeline.config is not None
        assert isinstance(pipeline.config, PipelineConfig)


# ============================================================================
# DUPLICATE REMOVAL TESTS
# ============================================================================

class TestRemoveDuplicates:
    """Test duplicate removal functionality"""
    
    def test_remove_duplicates_basic(self, dirty_df, config):
        """Test basic duplicate removal"""
        pipeline = DataCleaningPipeline(dirty_df, config)
        initial_rows = len(dirty_df)
        
        result = pipeline.remove_duplicates()
        
        assert len(result.df) < initial_rows
        assert result.metrics.rows_removed > 0
    
    def test_remove_duplicates_with_subset(self, dirty_df, config):
        """Test duplicate removal with subset"""
        pipeline = DataCleaningPipeline(dirty_df, config)
        result = pipeline.remove_duplicates(subset=['name'])
        
        assert result.df['name'].duplicated().sum() == 0
    
    def test_remove_duplicates_keep_last(self, config):
        """Test duplicate removal keeping last occurrence"""
        df = pd.DataFrame({
            'id': [1, 1, 2],
            'value': ['a', 'b', 'c']
        })
        pipeline = DataCleaningPipeline(df, config)
        result = pipeline.remove_duplicates(subset=['id'], keep='last')
        
        assert len(result.df) == 2
        assert result.df.loc[result.df['id'] == 1, 'value'].values[0] == 'b'
    
    def test_remove_duplicates_invalid_column(self, sample_df):
        """Test error handling for invalid column"""
        # strict_mode=True needed to raise exception
        config = PipelineConfig(strict_mode=True)
        pipeline = DataCleaningPipeline(sample_df, config)
        
        with pytest.raises(PipelineError):
            # ValidationError is wrapped in PipelineError in safe_operation if strict_mode is True
            pipeline.remove_duplicates(subset=['nonexistent'])


# ============================================================================
# MISSING VALUE HANDLING TESTS
# ============================================================================

class TestHandleMissingValues:
    """Test missing value handling"""
    
    def test_handle_missing_drop(self, config):
        """Test dropping rows with missing values"""
        df = pd.DataFrame({
            'a': [1, 2, None, 4],
            'b': [5, 6, 7, 8]
        })
        pipeline = DataCleaningPipeline(df, config)
        result = pipeline.handle_missing_values({'a': 'drop'})
        
        assert result.df['a'].isna().sum() == 0
        assert len(result.df) == 3
    
    def test_handle_missing_mean(self, config):
        """Test mean imputation"""
        df = pd.DataFrame({
            'value': [1.0, 2.0, None, 4.0]
        })
        pipeline = DataCleaningPipeline(df, config)
        result = pipeline.handle_missing_values({'value': 'mean'})
        
        assert result.df['value'].isna().sum() == 0
        expected_mean = (1 + 2 + 4) / 3
        assert result.df['value'].iloc[2] == pytest.approx(expected_mean)
    
    def test_handle_missing_median(self, config):
        """Test median imputation"""
        df = pd.DataFrame({
            'value': [1, 2, None, 10]
        })
        pipeline = DataCleaningPipeline(df, config)
        result = pipeline.handle_missing_values({'value': 'median'})
        
        assert result.df['value'].isna().sum() == 0
        assert result.df['value'].iloc[2] == 2.0
    
    def test_handle_missing_mode(self, config):
        """Test mode imputation"""
        df = pd.DataFrame({
            'category': ['A', 'B', 'A', None, 'A']
        })
        pipeline = DataCleaningPipeline(df, config)
        result = pipeline.handle_missing_values({'category': 'mode'})
        
        assert result.df['category'].isna().sum() == 0
        assert result.df['category'].iloc[3] == 'A'
    
    def test_handle_missing_constant(self, config):
        """Test constant value imputation"""
        df = pd.DataFrame({
            'value': [1, None, 3]
        })
        pipeline = DataCleaningPipeline(df, config)
        result = pipeline.handle_missing_values({'value': ('constant', 999)})
        
        assert result.df['value'].iloc[1] == 999
    
    def test_handle_missing_forward_fill(self, config):
        """Test forward fill"""
        df = pd.DataFrame({
            'value': [1, None, None, 4]
        })
        pipeline = DataCleaningPipeline(df, config)
        result = pipeline.handle_missing_values({'value': 'ffill'})
        
        assert result.df['value'].iloc[1] == 1
        assert result.df['value'].iloc[2] == 1
    
    def test_handle_missing_backward_fill(self, config):
        """Test backward fill"""
        df = pd.DataFrame({
            'value': [1, None, None, 4]
        })
        pipeline = DataCleaningPipeline(df, config)
        result = pipeline.handle_missing_values({'value': 'bfill'})
        
        assert result.df['value'].iloc[1] == 4
        assert result.df['value'].iloc[2] == 4


# ============================================================================
# OUTLIER REMOVAL TESTS
# ============================================================================

class TestRemoveOutliers:
    """Test outlier removal functionality"""
    
    def test_remove_outliers_iqr(self, config):
        """Test IQR method for outlier removal"""
        df = pd.DataFrame({
            'value': [1, 2, 3, 4, 5, 100]
        })
        pipeline = DataCleaningPipeline(df, config)
        result = pipeline.remove_outliers(['value'], method='iqr', threshold=1.5)
        
        assert len(result.df) < 6
        assert 100 not in result.df['value'].values
    
    def test_remove_outliers_zscore(self, config):
        """Test Z-score method for outlier removal"""
        df = pd.DataFrame({
            'value': [1, 2, 3, 4, 5, 100]
        })
        pipeline = DataCleaningPipeline(df, config)
        result = pipeline.remove_outliers(['value'], method='zscore', threshold=2)
        
        assert len(result.df) < 6
    
    def test_remove_outliers_modified_zscore(self, config):
        """Test modified Z-score method"""
        df = pd.DataFrame({
            'value': [1, 2, 3, 4, 5, 100]
        })
        pipeline = DataCleaningPipeline(df, config)
        result = pipeline.remove_outliers(['value'], method='modified_zscore', threshold=3)
        
        assert len(result.df) < 6
    
    def test_remove_outliers_non_numeric_skipped(self, dirty_df, config):
        """Test that non-numeric columns are skipped"""
        pipeline = DataCleaningPipeline(dirty_df, config)
        result = pipeline.remove_outliers(['name'], method='iqr')
        
        # Should not crash, column should be skipped
        assert 'name' in result.df.columns


# ============================================================================
# TEXT STANDARDIZATION TESTS
# ============================================================================

class TestStandardizeText:
    """Test text standardization"""
    
    def test_standardize_lowercase(self, config):
        """Test lowercase transformation"""
        df = pd.DataFrame({
            'text': ['HELLO', 'World', 'TEST']
        })
        pipeline = DataCleaningPipeline(df, config)
        result = pipeline.standardize_text(['text'], lowercase=True)
        
        assert all(result.df['text'].str.islower())
    
    def test_standardize_strip(self, config):
        """Test stripping whitespace"""
        df = pd.DataFrame({
            'text': ['  hello  ', '  world  ']
        })
        pipeline = DataCleaningPipeline(df, config)
        result = pipeline.standardize_text(['text'], strip=True, lowercase=False)
        
        assert result.df['text'].iloc[0] == 'hello'
        assert result.df['text'].iloc[1] == 'world'
    
    def test_standardize_remove_special_chars(self, config):
        """Test removing special characters"""
        df = pd.DataFrame({
            'text': ['hello!', 'world@', 'test#123']
        })
        pipeline = DataCleaningPipeline(df, config)
        result = pipeline.standardize_text(
            ['text'],
            remove_special_chars=True,
            lowercase=False
        )
        
        assert result.df['text'].iloc[0] == 'hello'
        assert result.df['text'].iloc[1] == 'world'
    
    def test_standardize_remove_extra_spaces(self, config):
        """Test removing extra spaces"""
        df = pd.DataFrame({
            'text': ['hello    world', 'test  123']
        })
        pipeline = DataCleaningPipeline(df, config)
        result = pipeline.standardize_text(
            ['text'],
            remove_extra_spaces=True,
            lowercase=False
        )
        
        assert '    ' not in result.df['text'].iloc[0]


# ============================================================================
# EMAIL VALIDATION TESTS
# ============================================================================

class TestValidateEmail:
    """Test email validation"""
    
    def test_validate_email_remove_invalid(self, config):
        """Test removing invalid emails"""
        df = pd.DataFrame({
            'email': ['valid@test.com', 'invalid', 'another@test.com']
        })
        pipeline = DataCleaningPipeline(df, config)
        result = pipeline.validate_email('email', remove_invalid=True)
        
        assert len(result.df) == 2
        assert 'invalid' not in result.df['email'].values
    
    def test_validate_email_add_column(self, config):
        """Test adding validation column"""
        df = pd.DataFrame({
            'email': ['valid@test.com', 'invalid']
        })
        pipeline = DataCleaningPipeline(df, config)
        result = pipeline.validate_email('email', remove_invalid=False, add_validation_column=True)
        
        assert 'email_valid' in result.df.columns
        assert result.df['email_valid'].iloc[0] == True
        assert result.df['email_valid'].iloc[1] == False


# ============================================================================
# DATA TYPE CONVERSION TESTS
# ============================================================================

class TestConvertDataTypes:
    """Test data type conversion"""
    
    def test_convert_to_datetime(self, config):
        """Test datetime conversion"""
        df = pd.DataFrame({
            'date': ['2023-01-01', '2023-01-02', '2023-01-03']
        })
        pipeline = DataCleaningPipeline(df, config)
        result = pipeline.convert_data_types({'date': 'datetime'})
        
        assert pd.api.types.is_datetime64_any_dtype(result.df['date'])
    
    def test_convert_to_int(self, config):
        """Test integer conversion"""
        df = pd.DataFrame({
            'value': ['1', '2', '3']
        })
        pipeline = DataCleaningPipeline(df, config)
        result = pipeline.convert_data_types({'value': 'int'})
        
        assert pd.api.types.is_integer_dtype(result.df['value'])
    
    def test_convert_to_float(self, config):
        """Test float conversion"""
        df = pd.DataFrame({
            'value': ['1.5', '2.5', '3.5']
        })
        pipeline = DataCleaningPipeline(df, config)
        result = pipeline.convert_data_types({'value': 'float'})
        
        assert pd.api.types.is_float_dtype(result.df['value'])


# ============================================================================
# ROLLBACK TESTS
# ============================================================================

class TestRollback:
    """Test rollback functionality"""
    
    def test_rollback_one_step(self, sample_df, config):
        """Test rolling back one step"""
        pipeline = DataCleaningPipeline(sample_df, config)
        original_len = len(pipeline.df)
        
        pipeline.remove_duplicates()
        pipeline.rollback(steps=1)
        
        assert len(pipeline.df) == original_len
    
    def test_rollback_multiple_steps(self, dirty_df, config):
        """Test rolling back multiple steps"""
        pipeline = DataCleaningPipeline(dirty_df, config)
        original_shape = pipeline.df.shape
        
        pipeline.remove_duplicates()
        pipeline.handle_missing_values({'salary': 'median'})
        pipeline.rollback(steps=2)
        
        assert pipeline.df.shape == original_shape


# ============================================================================
# INTEGRATION TESTS
# ============================================================================

class TestIntegration:
    """Test complete pipeline workflows"""
    
    def test_full_cleaning_workflow(self, dirty_df, config):
        """Test complete cleaning workflow"""
        pipeline = DataCleaningPipeline(dirty_df, config)
        
        result = (pipeline
            .remove_duplicates()
            .handle_missing_values({'salary': 'median', 'name': 'drop'})
            .remove_outliers(['age', 'salary'], method='iqr')
            .standardize_text(['name', 'city'], lowercase=True)
            .validate_email('email', remove_invalid=True)
            .reset_index()
            .get_cleaned_data())
        
        # Verify cleaning happened
        assert len(result) < len(dirty_df)
        assert result.isna().sum().sum() < dirty_df.isna().sum().sum()
    
    def test_method_chaining(self, sample_df, config):
        """Test that method chaining works correctly"""
        pipeline = DataCleaningPipeline(sample_df, config)
        
        result = (pipeline
            .remove_duplicates()
            .reset_index())
        
        assert isinstance(result, DataCleaningPipeline)
        assert result.df is not None


# ============================================================================
# EDGE CASE TESTS
# ============================================================================

class TestEdgeCases:
    """Test edge cases and error handling"""
    
    def test_empty_dataframe(self, config):
        """Test handling of empty DataFrame"""
        df = pd.DataFrame()
        
        with pytest.raises(ValidationError):
            DataCleaningPipeline(df, config)
    
    def test_single_row_dataframe(self, config):
        """Test handling of single-row DataFrame"""
        df = pd.DataFrame({'a': [1]})
        pipeline = DataCleaningPipeline(df, config)
        
        result = pipeline.remove_duplicates()
        assert len(result.df) == 1
    
    def test_all_missing_values(self, config):
        """Test handling of column with all missing values"""
        df = pd.DataFrame({
            'a': [None, None, None]
        })
        pipeline = DataCleaningPipeline(df, config)
        
        # Should handle gracefully
        result = pipeline.handle_missing_values({'a': 'mode'})
        assert result is not None


# ============================================================================
# PERFORMANCE TESTS
# ============================================================================

class TestPerformance:
    """Test performance characteristics"""
    
    def test_large_dataset_handling(self, config):
        """Test handling of larger datasets"""
        # Create 100K row dataset
        df = pd.DataFrame({
            'id': range(100000),
            'value': np.random.randint(0, 1000, 100000)
        })
        
        pipeline = DataCleaningPipeline(df, config)
        start_time = datetime.now()
        
        result = pipeline.remove_duplicates()
        
        execution_time = (datetime.now() - start_time).total_seconds()
        assert execution_time < 5  # Should complete in under 5 seconds
    
    def test_memory_efficiency(self, config):
        """Test that pipeline doesn't create excessive copies"""
        df = pd.DataFrame({
            'value': range(10000)
        })
        
        pipeline = DataCleaningPipeline(df, config)
        initial_memory = pipeline.df.memory_usage(deep=True).sum()
        
        pipeline.remove_duplicates()
        
        # Memory shouldn't increase dramatically
        final_memory = pipeline.df.memory_usage(deep=True).sum()
        assert final_memory <= initial_memory * 1.1  # Allow 10% overhead


# ============================================================================
# RUN TESTS
# ============================================================================

if __name__ == "__main__":
    pytest.main([__file__, "-v", "--tb=short"])