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

class TestSchemaValidation:
    """Test schema validation functionality"""
    
    def test_validate_schema_valid(self, sample_df, tmp_path):
        """Test validation against a valid schema"""
        schema_path = tmp_path / "schema.yaml"
        # Using simple indentation-based yaml content
        schema_content = """
test_schema:
  columns:
    age:
      type: int64
      min: 0
      max: 120
    salary:
      type: int64
"""
        schema_path.write_text(schema_content)
        
        config = PipelineConfig(strict_mode=True)
        pipeline = DataCleaningPipeline(sample_df, config)
        result = pipeline.validate_schema(schema_path, "test_schema")
        
        assert isinstance(result, DataCleaningPipeline)

    def test_validate_schema_invalid_type(self, sample_df, tmp_path):
        """Test validation failure on type mismatch"""
        schema_path = tmp_path / "schema.yaml"
        schema_content = """
test_schema:
  columns:
    name:
      type: int64
"""
        schema_path.write_text(schema_content)
        
        config = PipelineConfig(strict_mode=True)
        pipeline = DataCleaningPipeline(sample_df, config)
        
        with pytest.raises(PipelineError):
            pipeline.validate_schema(schema_path, "test_schema")


class TestConfig:
    """Test configuration loading"""
    
    def test_load_from_yaml(self, tmp_path):
        """Test loading config from YAML"""
        config_path = tmp_path / "config.yaml"
        config_content = """
production:
  log_level: DEBUG
  max_checkpoints: 5
"""
        config_path.write_text(config_content)
        
        config = PipelineConfig.from_yaml(config_path, "production")
        assert config.log_level == "DEBUG"
        assert config.max_checkpoints == 5


class TestChunking:
    """Test chunked processing"""
    
    def test_process_in_chunks(self, tmp_path):
        """Test processing CSV in chunks"""
        csv_path = tmp_path / "large.csv"
        df = pd.DataFrame({'a': range(100)})
        df.to_csv(csv_path, index=False)
        
        config = PipelineConfig()
        pipeline = DataCleaningPipeline(pd.DataFrame({'a': []}), config)
        
        chunks = list(pipeline.process_in_chunks(csv_path, chunk_size=20))
        assert len(chunks) == 5
        assert len(chunks[0].df) == 20


class TestExport:
    """Test export functionalities"""
    
    def test_to_json(self, sample_df, tmp_path):
        """Test JSON export with NumpyEncoder"""
        export_path = tmp_path / "export.json"
        config = PipelineConfig(allow_file_operations=True)
        pipeline = DataCleaningPipeline(sample_df, config)
        
        pipeline.to_json(export_path)
        assert export_path.exists()
        
    def test_to_parquet(self, sample_df, tmp_path):
        """Test Parquet export"""
        export_path = tmp_path / "export.parquet"
        config = PipelineConfig(allow_file_operations=True)
        pipeline = DataCleaningPipeline(sample_df, config)
        
        try:
            import pyarrow
            import fastparquet
        except ImportError:
            pytest.skip("pyarrow or fastparquet not installed")
            
        try:
            pipeline.to_parquet(export_path)
            assert export_path.exists()
        except Exception as e:
            pytest.fail(f"Parquet export failed: {e}")


class TestRollbackConfig:
    """Test configurable rollback"""

    def test_max_checkpoints(self, sample_df):
        """Test that history respects max_checkpoints"""
        config = PipelineConfig(max_checkpoints=2)
        pipeline = DataCleaningPipeline(sample_df, config)
        
        # Initial state is checkpoint 1 (from init calls to validation?) 
        # Actually init calls _checkpoint_data once.
        # History: [Init State]
        assert len(pipeline._history) == 1
        
        pipeline.remove_duplicates()
        # History: [Init State, After Dedup] (Len 2)
        
        pipeline.reset_index()
        # History: [After Dedup, After Reset] (Len 2, Init dropped)
        
        assert len(pipeline._history) == 2
        # Verify first item is not original if logical shift happened, 
        # or size is just capped.
