"""
Integration Tests for ENAHOPY
============================

Tests that verify the complete workflow: download → read → merge

Author: ENAHOPY Team
Date: 2025-01-03
"""

import os
import tempfile
import unittest
from pathlib import Path
from unittest.mock import MagicMock, Mock, patch

import numpy as np
import pandas as pd

from enahopy.loader.io.local_reader import ENAHOLocalReader
from enahopy.merger import ENAHOGeoMerger
from enahopy.merger.config import GeoMergeConfiguration


def safe_read_data(reader):
    """Helper function to safely extract DataFrame from read_data() result."""
    result = reader.read_data()
    if isinstance(result, tuple):
        return result[0]  # Return DataFrame
    return result


class TestDownloadReadMergeWorkflow(unittest.TestCase):
    """Test the complete download → read → merge workflow."""

    def setUp(self):
        """Set up test environment with sample data."""
        # Sample ENAHO data
        self.sample_enaho_data = pd.DataFrame(
            {
                "conglome": ["001", "001", "002", "003"],
                "vivienda": ["01", "02", "01", "01"],
                "hogar": ["1", "1", "1", "1"],
                "ubigeo": ["150101", "150102", "150201", "150301"],
                "ingreso": [1000, 1200, 800, 1500],
                "gasto": [800, 900, 600, 1100],
            }
        )

        # Sample geographic data
        self.sample_geo_data = pd.DataFrame(
            {
                "ubigeo": ["150101", "150102", "150201", "150301", "150401"],
                "departamento": ["Lima", "Lima", "Lima", "Lima", "Lima"],
                "provincia": ["Lima", "Lima", "Huaral", "Canta", "Huaura"],
                "distrito": ["Cercado de Lima", "Ancón", "Huaral", "Canta", "Huacho"],
                "region": ["Costa", "Costa", "Costa", "Sierra", "Costa"],
            }
        )

        self.temp_dir = tempfile.mkdtemp()

    def tearDown(self):
        """Clean up test environment."""
        import shutil

        shutil.rmtree(self.temp_dir, ignore_errors=True)

    def test_local_file_read_and_merge_workflow(self):
        """Test reading local files and merging with geographic data."""
        # Step 1: Save test data to files
        enaho_file = Path(self.temp_dir) / "enaho_data.csv"
        geo_file = Path(self.temp_dir) / "geo_data.csv"

        self.sample_enaho_data.to_csv(enaho_file, index=False)
        self.sample_geo_data.to_csv(geo_file, index=False)

        # Step 2: Read ENAHO data
        reader = ENAHOLocalReader(file_path=str(enaho_file))
        enaho_df = safe_read_data(reader)

        # Verify data was read correctly
        self.assertIsInstance(enaho_df, pd.DataFrame)
        self.assertEqual(len(enaho_df), 4)
        self.assertIn("ubigeo", enaho_df.columns)
        self.assertIn("ingreso", enaho_df.columns)

        # Step 3: Read geographic data
        geo_reader = ENAHOLocalReader(file_path=str(geo_file))
        geo_df = safe_read_data(geo_reader)

        # Step 4: Merge with geographic data
        geo_config = GeoMergeConfiguration()
        merger = ENAHOGeoMerger(geo_config=geo_config, verbose=False)

        merged_df, validation = merger.merge_geographic_data(
            df_principal=enaho_df,
            df_geografia=geo_df,
            columnas_geograficas={
                "departamento": "departamento",
                "provincia": "provincia",
                "distrito": "distrito",
                "region": "region",
            },
        )

        # Verify merge results
        self.assertIsInstance(merged_df, pd.DataFrame)
        self.assertEqual(len(merged_df), 4)  # All ENAHO records should be preserved

        # Verify geographic columns were added
        expected_geo_cols = ["departamento", "provincia", "distrito", "region"]
        for col in expected_geo_cols:
            self.assertIn(col, merged_df.columns)

        # Verify data integrity
        self.assertEqual(merged_df["departamento"].iloc[0], "Lima")
        self.assertEqual(merged_df["distrito"].iloc[1], "Ancón")
        self.assertEqual(merged_df["region"].iloc[3], "Sierra")

        # Verify validation results
        self.assertTrue(validation.coverage_percentage > 0)
        self.assertEqual(validation.total_records, 4)

    def test_multi_module_merge_workflow(self):
        """Test merging multiple ENAHO modules and then adding geography."""
        # Create sample module data
        module_34 = pd.DataFrame(
            {  # Sumaria
                "conglome": ["001", "002", "003"],
                "vivienda": ["01", "01", "01"],
                "hogar": ["1", "1", "1"],
                "ubigeo": ["150101", "150201", "150301"],
                "ingreso_total": [2000, 1500, 2500],
            }
        )

        module_01 = pd.DataFrame(
            {  # Características de la vivienda
                "conglome": ["001", "002", "003"],
                "vivienda": ["01", "01", "01"],
                "hogar": ["1", "1", "1"],
                "tipo_vivienda": ["Casa", "Departamento", "Casa"],
                "material_pared": ["Ladrillo", "Concreto", "Adobe"],
            }
        )

        # Save to files
        file_34 = Path(self.temp_dir) / "module_34.csv"
        file_01 = Path(self.temp_dir) / "module_01.csv"
        geo_file = Path(self.temp_dir) / "geo_data.csv"

        module_34.to_csv(file_34, index=False)
        module_01.to_csv(file_01, index=False)
        self.sample_geo_data.to_csv(geo_file, index=False)

        # Read modules using helper function
        df_34 = safe_read_data(ENAHOLocalReader(file_path=str(file_34)))
        df_01 = safe_read_data(ENAHOLocalReader(file_path=str(file_01)))
        geo_df = safe_read_data(ENAHOLocalReader(file_path=str(geo_file)))

        # Merge modules and geography
        merger = ENAHOGeoMerger(verbose=False)

        modules_dict = {"34": df_34, "01": df_01}

        final_df, report = merger.merge_modules_with_geography(
            modules_dict=modules_dict, df_geografia=geo_df, base_module="34"
        )

        # Verify results
        self.assertIsInstance(final_df, pd.DataFrame)
        self.assertEqual(len(final_df), 3)  # Base module determines final count

        # Verify all columns are present
        expected_cols = [
            "ingreso_total",
            "tipo_vivienda",
            "material_pared",
            "departamento",
            "provincia",
            "distrito",
        ]
        for col in expected_cols:
            self.assertIn(col, final_df.columns)

        # Verify merge quality
        self.assertIn("module_merge", report)
        self.assertIn("geographic_merge", report)
        self.assertIn("overall_quality", report)

    @patch("requests.get")
    def test_mock_download_workflow(self, mock_get):
        """Test workflow with mocked INEI server downloads."""
        # Mock server response
        mock_response = Mock()
        mock_response.status_code = 200
        mock_response.content = (
            b"conglome,vivienda,hogar,ubigeo,ingreso\n001,01,1,150101,1000\n002,01,1,150201,1200"
        )
        mock_response.headers = {"content-type": "text/csv"}
        mock_get.return_value = mock_response

        # Create geographic data file
        geo_file = Path(self.temp_dir) / "geo_data.csv"
        self.sample_geo_data.to_csv(geo_file, index=False)

        # Simulate download workflow
        # Note: We would normally use the downloader here, but we'll simulate it
        downloaded_file = Path(self.temp_dir) / "downloaded_enaho.csv"
        with open(downloaded_file, "w") as f:
            f.write("conglome,vivienda,hogar,ubigeo,ingreso\n")
            f.write("001,01,1,150101,1000\n")
            f.write("002,01,1,150201,1200\n")

        # Read and merge
        reader = ENAHOLocalReader(file_path=str(downloaded_file))
        enaho_df = safe_read_data(reader)

        geo_reader = ENAHOLocalReader(file_path=str(geo_file))
        geo_df = safe_read_data(geo_reader)

        merger = ENAHOGeoMerger(verbose=False)
        merged_df, validation = merger.merge_geographic_data(
            df_principal=enaho_df, df_geografia=geo_df
        )

        # Verify workflow completed successfully
        self.assertEqual(len(merged_df), 2)
        self.assertIn("departamento", merged_df.columns)
        self.assertTrue(validation.coverage_percentage > 0)

        # Verify mock was called
        mock_get.assert_called()

    def test_error_handling_in_workflow(self):
        """Test error handling throughout the workflow."""
        # Test with invalid file
        with self.assertRaises(FileNotFoundError):
            reader = ENAHOLocalReader(file_path="nonexistent_file.csv")
            reader.read_data()

        # Test with empty DataFrames
        merger = ENAHOGeoMerger(verbose=False)

        with self.assertRaises(ValueError):
            merger.merge_geographic_data(
                df_principal=pd.DataFrame(), df_geografia=self.sample_geo_data  # Empty DataFrame
            )

        # Test with incompatible columns
        bad_enaho = pd.DataFrame(
            {"wrong_column": [1, 2, 3], "another_wrong_column": ["a", "b", "c"]}
        )

        with self.assertRaises(ValueError):
            merger.merge_geographic_data(
                df_principal=bad_enaho,
                df_geografia=self.sample_geo_data,
                columna_union="ubigeo",  # Column doesn't exist
            )

    def test_data_quality_validation_workflow(self):
        """Test data quality validation throughout the workflow."""
        # Create data with quality issues
        problematic_data = pd.DataFrame(
            {
                "conglome": ["001", "002", "003", "004"],
                "vivienda": ["01", "01", "01", "01"],
                "hogar": ["1", "1", "1", "1"],
                "ubigeo": ["150101", "999999", "150201", None],  # Invalid and missing UBIGEO
                "ingreso": [1000, 1200, 800, 1500],
            }
        )

        # Save to file
        problem_file = Path(self.temp_dir) / "problematic_data.csv"
        geo_file = Path(self.temp_dir) / "geo_data.csv"

        problematic_data.to_csv(problem_file, index=False)
        self.sample_geo_data.to_csv(geo_file, index=False)

        # Read data
        reader = ENAHOLocalReader(file_path=str(problem_file))
        enaho_df = safe_read_data(reader)

        geo_reader = ENAHOLocalReader(file_path=str(geo_file))
        geo_df = safe_read_data(geo_reader)

        # Merge with validation
        merger = ENAHOGeoMerger(verbose=False)
        merged_df, validation = merger.merge_geographic_data(
            df_principal=enaho_df, df_geografia=geo_df, validate_before_merge=True
        )

        # Verify validation detected issues
        self.assertLess(
            validation.coverage_percentage, 100
        )  # Should be < 100% due to invalid UBIGEO
        self.assertGreater(validation.invalid_ubigeos, 0)  # Should detect invalid UBIGEOs

        # Verify workflow still completed (with warnings)
        self.assertEqual(len(merged_df), 4)  # All records preserved

        # Verify geographic data is missing for problematic records
        # Record with invalid UBIGEO should have NaN for geographic columns
        invalid_row = merged_df[merged_df["ubigeo"] == "999999"]
        if not invalid_row.empty and "departamento" in merged_df.columns:
            self.assertTrue(pd.isna(invalid_row["departamento"].iloc[0]))


class TestPerformanceScenarios(unittest.TestCase):
    """Test performance-related scenarios and edge cases."""

    def setUp(self):
        """Set up performance test environment."""
        self.temp_dir = tempfile.mkdtemp()

    def tearDown(self):
        """Clean up test environment."""
        import shutil

        shutil.rmtree(self.temp_dir, ignore_errors=True)

    def test_large_dataset_simulation(self):
        """Test workflow with simulated large datasets."""
        # Create larger simulated datasets
        n_records = 1000  # Reasonable size for testing

        large_enaho = pd.DataFrame(
            {
                "conglome": [f"{i:03d}" for i in range(1, n_records + 1)],
                "vivienda": ["01"] * n_records,
                "hogar": ["1"] * n_records,
                "ubigeo": np.random.choice(["150101", "150102", "150201", "150301"], n_records),
                "ingreso": np.random.randint(500, 5000, n_records),
                "gasto": np.random.randint(400, 3000, n_records),
            }
        )

        # Geographic data (smaller, as typical)
        geo_data = pd.DataFrame(
            {
                "ubigeo": ["150101", "150102", "150201", "150301"],
                "departamento": ["Lima"] * 4,
                "provincia": ["Lima", "Lima", "Huaral", "Canta"],
                "distrito": ["Cercado", "Ancón", "Huaral", "Canta"],
            }
        )

        # Save to files
        enaho_file = Path(self.temp_dir) / "large_enaho.csv"
        geo_file = Path(self.temp_dir) / "geo_data.csv"

        large_enaho.to_csv(enaho_file, index=False)
        geo_data.to_csv(geo_file, index=False)

        # Test workflow
        reader = ENAHOLocalReader(file_path=str(enaho_file))
        enaho_df = safe_read_data(reader)

        geo_reader = ENAHOLocalReader(file_path=str(geo_file))
        geo_df = safe_read_data(geo_reader)

        merger = ENAHOGeoMerger(verbose=False)

        # Time the merge operation
        import time

        start_time = time.time()

        merged_df, validation = merger.merge_geographic_data(
            df_principal=enaho_df, df_geografia=geo_df
        )

        end_time = time.time()

        # Verify results
        self.assertEqual(len(merged_df), n_records)
        self.assertIn("departamento", merged_df.columns)

        # Performance should be reasonable (< 5 seconds for 1000 records)
        self.assertLess(end_time - start_time, 5.0)

        # Verify merge quality
        self.assertGreater(validation.coverage_percentage, 50)  # At least 50% should match

    def test_memory_efficient_chunked_processing(self):
        """Test memory-efficient chunked processing."""
        # This would test the chunked processing functionality
        # For now, we'll create a simple test

        # Create moderate-sized data
        n_records = 500
        chunk_size = 100

        enaho_data = pd.DataFrame(
            {
                "conglome": [f"{i:03d}" for i in range(1, n_records + 1)],
                "vivienda": ["01"] * n_records,
                "hogar": ["1"] * n_records,
                "ubigeo": "150101",  # All same UBIGEO for simplicity
                "data": np.random.randn(n_records),
            }
        )

        geo_data = pd.DataFrame(
            {"ubigeo": ["150101"], "departamento": ["Lima"], "provincia": ["Lima"]}
        )

        # Configure for chunked processing
        from enahopy.merger.config import GeoMergeConfiguration

        config = GeoMergeConfiguration(optimizar_memoria=True, chunk_size=chunk_size)

        merger = ENAHOGeoMerger(geo_config=config, verbose=False)

        merged_df, validation = merger.merge_geographic_data(
            df_principal=enaho_data, df_geografia=geo_data
        )

        # Verify chunked processing worked correctly
        self.assertEqual(len(merged_df), n_records)
        self.assertIn("departamento", merged_df.columns)


if __name__ == "__main__":
    unittest.main()
