"""Tests for core pipeline functionality."""

import shutil
import tempfile
from pathlib import Path
from typing import Any
from unittest.mock import Mock, patch

import pandas as pd
import pytest
import requests

# Import from earthcatalog package
from earthcatalog import ingestion_pipeline


# Module-level functions for multiprocessing tests (must be pickleable)
def mock_process_fn(chunk, worker_id):
    """Mock processing function for multiprocessing tests."""
    return f"processed_{len(chunk)}_items_by_worker_{worker_id}"


def mock_consolidate_fn(partition_key, shard_paths):
    """Mock consolidation function for multiprocessing tests."""
    return f"consolidated_{partition_key}_{len(shard_paths)}_shards"


class TestProcessingConfig:
    """Test ProcessingConfig dataclass."""

    def test_default_config(self):
        """Test default configuration values."""
        config = ingestion_pipeline.ProcessingConfig(
            input_file="test.parquet", output_catalog="./output", scratch_location="./scratch"
        )

        assert config.input_file == "test.parquet"
        assert config.output_catalog == "./output"
        assert config.scratch_location == "./scratch"
        assert config.grid_system == "h3"
        assert config.grid_resolution == 2
        assert config.temporal_bin == "month"
        assert config.sort_key == "datetime"
        assert config.sort_ascending is True
        assert config.items_per_shard == 10000
        assert config.max_workers == 8

    def test_custom_config(self):
        """Test custom configuration values."""
        config = ingestion_pipeline.ProcessingConfig(
            input_file="custom.parquet",
            output_catalog="./custom_output",
            scratch_location="./custom_scratch",
            grid_system="s2",
            grid_resolution=10,
            temporal_bin="day",
            sort_key="properties.datetime",
            sort_ascending=False,
            items_per_shard=5000,
            max_workers=4,
        )

        assert config.grid_system == "s2"
        assert config.grid_resolution == 10
        assert config.temporal_bin == "day"
        assert config.sort_key == "properties.datetime"
        assert config.sort_ascending is False
        assert config.items_per_shard == 5000
        assert config.max_workers == 4

    def test_to_dict_includes_all_fields(self):
        """to_dict should include all configuration fields."""
        config = ingestion_pipeline.ProcessingConfig(
            input_file="input.parquet",
            output_catalog="./catalog",
            scratch_location="./scratch",
        )

        data = config.to_dict()

        # Check required fields
        assert data["input_file"] == "input.parquet"
        assert data["output_catalog"] == "./catalog"
        assert data["scratch_location"] == "./scratch"

        # Check some default fields
        assert data["grid_system"] == "h3"
        assert data["grid_resolution"] == 2
        assert data["batch_size"] == 1000
        assert data["enable_concurrent_http"] is True

    def test_from_dict_restores_config(self):
        """from_dict should restore config from dictionary."""
        data = {
            "input_file": "test.parquet",
            "output_catalog": "s3://bucket/catalog",
            "scratch_location": "s3://bucket/scratch",
            "grid_resolution": 5,
            "batch_size": 2000,
            "concurrent_requests": 100,
        }

        config = ingestion_pipeline.ProcessingConfig.from_dict(data)

        assert config.input_file == "test.parquet"
        assert config.output_catalog == "s3://bucket/catalog"
        assert config.grid_resolution == 5
        assert config.batch_size == 2000
        assert config.concurrent_requests == 100
        # Defaults should still apply for missing fields
        assert config.grid_system == "h3"

    def test_to_dict_from_dict_roundtrip(self):
        """to_dict and from_dict should roundtrip correctly."""
        original = ingestion_pipeline.ProcessingConfig(
            input_file="input.parquet",
            output_catalog="./catalog",
            scratch_location="./scratch",
            grid_resolution=4,
            batch_size=500,
            concurrent_requests=75,
        )

        data = original.to_dict()
        restored = ingestion_pipeline.ProcessingConfig.from_dict(data)

        assert restored.input_file == original.input_file
        assert restored.grid_resolution == original.grid_resolution
        assert restored.batch_size == original.batch_size
        assert restored.concurrent_requests == original.concurrent_requests

    def test_config_hash_consistent(self):
        """config_hash should return same hash for same settings."""
        config1 = ingestion_pipeline.ProcessingConfig(
            input_file="a.parquet",
            output_catalog="./catalog1",
            scratch_location="./scratch1",
        )
        config2 = ingestion_pipeline.ProcessingConfig(
            input_file="b.parquet",  # Different input
            output_catalog="./catalog2",  # Different output
            scratch_location="./scratch2",  # Different scratch
        )

        # Same processing settings should give same hash
        assert config1.config_hash() == config2.config_hash()

    def test_config_hash_different_for_different_settings(self):
        """config_hash should differ when processing settings differ."""
        config1 = ingestion_pipeline.ProcessingConfig(
            input_file="input.parquet",
            output_catalog="./catalog",
            scratch_location="./scratch",
            grid_resolution=4,
        )
        config2 = ingestion_pipeline.ProcessingConfig(
            input_file="input.parquet",
            output_catalog="./catalog",
            scratch_location="./scratch",
            grid_resolution=5,  # Different resolution
        )

        assert config1.config_hash() != config2.config_hash()


class TestLocalProcessor:
    """Test LocalProcessor class."""

    def test_processor_initialization(self):
        """Test processor initialization."""
        processor = ingestion_pipeline.LocalProcessor(n_workers=4)
        assert processor.executor._max_workers == 4
        processor.close()

    def test_process_urls(self):
        """Test URL processing functionality."""
        processor = ingestion_pipeline.LocalProcessor(n_workers=2)

        url_chunks = [["url1", "url2"], ["url3", "url4"]]
        results = processor.process_urls(url_chunks, mock_process_fn)

        assert len(results) == 2
        assert "processed_2_items_by_worker_0" in results
        assert "processed_2_items_by_worker_1" in results

        processor.close()

    def test_consolidate_shards(self):
        """Test shard consolidation functionality."""
        processor = ingestion_pipeline.LocalProcessor(n_workers=2)

        partition_items = [("partition1", ["shard1", "shard2"]), ("partition2", ["shard3"])]
        results = processor.consolidate_shards(partition_items, mock_consolidate_fn)

        assert len(results) == 2
        assert "consolidated_partition1_2_shards" in results
        assert "consolidated_partition2_1_shards" in results

        processor.close()


class TestSTACIngestionPipeline:
    """Test STACIngestionPipeline class."""

    def setup_method(self):
        """Set up test fixtures."""
        self.temp_dir = tempfile.mkdtemp()
        self.config = ingestion_pipeline.ProcessingConfig(
            input_file=str(Path(self.temp_dir) / "input.parquet"),
            output_catalog=str(Path(self.temp_dir) / "catalog"),
            scratch_location=str(Path(self.temp_dir) / "scratch"),
            items_per_shard=10,
            max_workers=2,
        )

        # Create sample input parquet
        urls = ["https://example.com/item1.json", "https://example.com/item2.json"]
        df = pd.DataFrame({"url": urls})
        df.to_parquet(self.config.input_file, index=False)

        self.processor = ingestion_pipeline.LocalProcessor(n_workers=2)
        self.pipeline = ingestion_pipeline.STACIngestionPipeline(self.config, self.processor)

    def teardown_method(self):
        """Clean up test fixtures."""
        self.processor.close()
        shutil.rmtree(self.temp_dir, ignore_errors=True)

    def test_pipeline_initialization(self):
        """Test pipeline initialization."""
        assert self.pipeline.config == self.config
        assert self.pipeline.processor == self.processor
        assert self.pipeline.grid is not None
        assert self.pipeline.storage is not None
        assert self.pipeline.scratch_storage is not None

    def test_read_input_urls(self):
        """Test reading URLs from input parquet."""
        urls = self.pipeline._read_input_urls()
        assert len(urls) == 2
        assert "https://example.com/item1.json" in urls
        assert "https://example.com/item2.json" in urls

    def test_read_input_urls_missing_column(self):
        """Test error handling for missing URL column."""
        # Create parquet without URL column
        df = pd.DataFrame({"other_column": [1, 2]})
        bad_path = str(Path(self.temp_dir) / "bad.parquet")
        df.to_parquet(bad_path, index=False)

        self.config.input_file = bad_path
        pipeline = ingestion_pipeline.STACIngestionPipeline(self.config, self.processor)

        with pytest.raises(ValueError, match="must contain 'url' column"):
            pipeline._read_input_urls()

    def test_chunk_urls(self):
        """Test URL chunking functionality."""
        urls = [f"url_{i}" for i in range(10)]
        chunks = self.pipeline._chunk_urls(urls, 3)

        assert len(chunks) == 3
        # With current implementation: chunk_size = 10 // 3 = 3
        # chunks = [0:3], [3:6], [6:9], [9:10] -> last gets combined with previous
        assert len(chunks[0]) == 3
        assert len(chunks[1]) == 3
        assert len(chunks[2]) == 4  # Gets the remainder
        assert sum(len(chunk) for chunk in chunks) == 10

    def test_extract_temporal_hive_parts(self):
        """Test Hive-style temporal partition extraction."""
        # Test different temporal bins
        item_month = {"properties": {"datetime": "2024-01-15T10:30:00Z"}}
        item_year = {"properties": {"datetime": "2024-01-15T10:30:00Z"}}
        item_day = {"properties": {"datetime": "2024-01-15T10:30:00Z"}}
        item_no_datetime: dict[str, Any] = {"properties": {}}

        # Test month binning (Hive-style)
        self.config.temporal_bin = "month"
        result = self.pipeline._extract_temporal_hive_parts(item_month)
        assert result == "year=2024/month=01"

        # Test year binning (Hive-style)
        self.config.temporal_bin = "year"
        result = self.pipeline._extract_temporal_hive_parts(item_year)
        assert result == "year=2024"

        # Test day binning (Hive-style)
        self.config.temporal_bin = "day"
        result = self.pipeline._extract_temporal_hive_parts(item_day)
        assert result == "year=2024/month=01/day=15"

        # Test missing datetime
        result = self.pipeline._extract_temporal_hive_parts(item_no_datetime)
        assert result == "unknown"

    def test_compute_partition_key(self):
        """Test partition key computation with Hive-style temporal parts."""
        # Mock grid system
        mock_grid = Mock()
        mock_grid.tiles_for_geometry_with_spanning_detection.return_value = (["h3_cell_123"], False)
        mock_grid.get_global_partition_threshold.return_value = 10  # Return int instead of Mock
        self.pipeline.grid = mock_grid

        item = {
            "geometry": {"type": "Point", "coordinates": [-105.0, 40.0]},
            "properties": {"datetime": "2024-01-15T10:30:00Z"},
        }

        self.config.temporal_bin = "month"
        partition_key = self.pipeline._compute_partition_key(item)

        assert partition_key == "unknown_mission/partition=h3/level=2/h3_cell_123/year=2024/month=01"
        mock_grid.tiles_for_geometry_with_spanning_detection.assert_called_once()

    def test_get_final_partition_path(self):
        """Test final partition path generation with Hive-style."""
        partition_key = "mission/partition=h3/level=2/h3_cell_123/year=2024/month=01"
        expected_path = f"{self.config.output_catalog}/{partition_key}/items.parquet"

        result = self.pipeline._get_final_partition_path(partition_key)
        assert result == expected_path

    @patch("requests.get")
    def test_download_stac_item_http(self, mock_get):
        """Test downloading STAC item from HTTP URL."""
        # Mock successful HTTP response
        mock_response = Mock()
        mock_response.json.return_value = {"id": "test_item", "type": "Feature"}
        mock_response.raise_for_status.return_value = None
        mock_get.return_value = mock_response

        url = "https://example.com/item.json"
        result = self.pipeline._download_stac_item(url)

        assert result == {"id": "test_item", "type": "Feature"}
        mock_get.assert_called_once_with(url, timeout=30)

    @patch("fsspec.filesystem")
    def test_download_stac_item_s3(self, mock_fs):
        """Test downloading STAC item from S3."""
        # Mock S3 filesystem
        mock_s3_fs = Mock()
        mock_file = Mock()
        mock_file.__enter__ = Mock(return_value=mock_file)
        mock_file.__exit__ = Mock(return_value=None)
        mock_s3_fs.open.return_value = mock_file
        mock_fs.return_value = mock_s3_fs

        # Mock json.load
        with patch("json.load", return_value={"id": "s3_item"}):
            url = "s3://bucket/item.json"
            result = self.pipeline._download_stac_item(url)

            assert result == {"id": "s3_item"}
            mock_fs.assert_called_once_with("s3")

    def test_download_stac_item_error(self):
        """Test error handling in STAC item download."""
        url = "https://invalid-url.com/item.json"

        with patch("requests.get", side_effect=requests.exceptions.RequestException("Network error")):
            result = self.pipeline._download_stac_item(url)
            assert result is None


class TestDaskDistributedProcessor:
    """Test DaskDistributedProcessor class."""

    def test_dask_processor_import_error(self):
        """Test DaskDistributedProcessor raises ImportError when dask.distributed is not available."""
        # Mock the dask.distributed module to raise ImportError
        with patch.dict("sys.modules", {"dask.distributed": None}):
            with patch("builtins.__import__") as mock_import:

                def import_side_effect(name, *args, **kwargs):
                    if name == "dask.distributed":
                        raise ImportError("No module named 'dask.distributed'")
                    return __import__(name, *args, **kwargs)

                mock_import.side_effect = import_side_effect

                with pytest.raises(ImportError, match="Dask distributed required"):
                    ingestion_pipeline.DaskDistributedProcessor()


class TestPartitionAwareSharding:
    """Test partition-aware shard organization."""

    def setup_method(self):
        """Set up test fixtures."""
        self.temp_dir = tempfile.mkdtemp()
        self.config = ingestion_pipeline.ProcessingConfig(
            input_file=str(Path(self.temp_dir) / "input.parquet"),
            output_catalog=str(Path(self.temp_dir) / "catalog"),
            scratch_location=str(Path(self.temp_dir) / "scratch"),
            items_per_shard=2,  # Small shards for testing
            max_workers=2,
        )

        # Create sample input parquet
        urls = [
            "https://example.com/item1.json",
            "https://example.com/item2.json",
            "https://example.com/item3.json",
        ]
        df = pd.DataFrame({"url": urls})
        df.to_parquet(self.config.input_file, index=False)

        self.processor = ingestion_pipeline.LocalProcessor(n_workers=2)
        self.pipeline = ingestion_pipeline.STACIngestionPipeline(self.config, self.processor)

    def teardown_method(self):
        """Clean up test fixtures."""
        self.processor.close()
        shutil.rmtree(self.temp_dir, ignore_errors=True)

    def test_partition_aware_shard_path(self):
        """Test that shard paths include partition information."""
        # Mock sample STAC items
        items = [
            {
                "id": "item1",
                "geometry": {"type": "Point", "coordinates": [-105.0, 40.0]},
                "properties": {"datetime": "2024-01-15T10:30:00Z"},
            },
            {
                "id": "item2",
                "geometry": {"type": "Point", "coordinates": [-105.1, 40.1]},
                "properties": {"datetime": "2024-01-16T11:30:00Z"},
            },
        ]

        # Mock grid system to return predictable tile
        mock_grid = Mock()
        mock_grid.tiles_for_geometry_with_spanning_detection.return_value = (["h3_cell_123"], False)
        self.pipeline.grid = mock_grid

        # Test shard path generation with Hive-style temporal parts
        partition_key = self.pipeline._compute_partition_key(items[0])

        # The new implementation should include Hive-style temporal parts in partition key
        assert partition_key == "unknown_mission/partition=h3/level=2/h3_cell_123/year=2024/month=01"

    def test_consolidation_config(self):
        """Test consolidation configuration options."""
        # Test that config options are available
        config = ingestion_pipeline.ProcessingConfig(
            input_file="test.parquet",
            output_catalog="./output",
            scratch_location="./scratch",
            max_memory_per_partition_mb=512,
            enable_streaming_merge=True,
        )

        assert config.max_memory_per_partition_mb == 512
        assert config.enable_streaming_merge is True

    def test_shard_grouping_optimization(self):
        """Test that partition-aware sharding simplifies grouping with Hive-style paths."""
        # Test the new grouping logic structure with Hive-style temporal parts
        expected_groups = {
            "h3_cell_123/year=2024/month=01": [
                f"{self.config.scratch_location}/shards/h3_cell_123/year=2024/month=01/worker-1.parquet",
                f"{self.config.scratch_location}/shards/h3_cell_123/year=2024/month=01/worker-2.parquet",
            ],
            "h3_cell_456/year=2024/month=01": [
                f"{self.config.scratch_location}/shards/h3_cell_456/year=2024/month=01/worker-1.parquet"
            ],
        }

        # We'll implement the new grouping logic that uses path structure
        # For now, test the expected structure
        assert len(expected_groups) == 2
        assert len(expected_groups["h3_cell_123/year=2024/month=01"]) == 2
        assert len(expected_groups["h3_cell_456/year=2024/month=01"]) == 1


class TestS3StorageEnhancements:
    """Test enhanced S3 storage backend operations."""

    def test_atomic_s3_operations_config(self):
        """Test S3 atomic operations configuration."""
        config = ingestion_pipeline.ProcessingConfig(
            input_file="test.parquet",
            output_catalog="s3://bucket/catalog",
            scratch_location="s3://bucket/scratch",
            s3_multipart_threshold_mb=100,
            temp_dir_location="/tmp",
        )

        assert config.s3_multipart_threshold_mb == 100
        assert config.temp_dir_location == "/tmp"

    def test_streaming_merge_memory_limits(self):
        """Test memory-bounded streaming merge."""
        # Test configuration for memory limits
        config = ingestion_pipeline.ProcessingConfig(
            input_file="test.parquet",
            output_catalog="./output",
            scratch_location="./scratch",
            max_memory_per_partition_mb=256,
        )

        # Should use streaming when partition size exceeds limit
        assert config.max_memory_per_partition_mb == 256


class TestIntegratedPipelineWorkflow:
    """Test integrated pipeline workflow with improvements."""

    def setup_method(self):
        """Set up test fixtures."""
        self.temp_dir = tempfile.mkdtemp()

    def teardown_method(self):
        """Clean up test fixtures."""
        shutil.rmtree(self.temp_dir, ignore_errors=True)

    def test_end_to_end_partition_aware_workflow(self):
        """Test end-to-end workflow with partition-aware sharding."""
        # This will be a comprehensive test once we implement the changes
        config = ingestion_pipeline.ProcessingConfig(
            input_file=str(Path(self.temp_dir) / "input.parquet"),
            output_catalog=str(Path(self.temp_dir) / "catalog"),
            scratch_location=str(Path(self.temp_dir) / "scratch"),
            max_workers=2,
            items_per_shard=2,
        )

        # Create minimal test input
        urls = ["https://example.com/item1.json"]
        df = pd.DataFrame({"url": urls})
        df.to_parquet(config.input_file, index=False)

        processor = ingestion_pipeline.LocalProcessor(n_workers=2)
        pipeline = ingestion_pipeline.STACIngestionPipeline(config, processor)

        # Test that pipeline initializes with config
        assert pipeline.config.max_memory_per_partition_mb > 0

        processor.close()


class TestBatchModeConfig:
    """Test batch mode configuration and processing mode selection."""

    def setup_method(self):
        """Set up test fixtures."""
        self.temp_dir = tempfile.mkdtemp()

    def teardown_method(self):
        """Clean up test fixtures."""
        shutil.rmtree(self.temp_dir, ignore_errors=True)

    def test_batch_mode_config_defaults(self):
        """Test default batch mode configuration values."""
        config = ingestion_pipeline.ProcessingConfig(
            input_file="test.parquet",
            output_catalog="./output",
            scratch_location="./scratch",
        )

        assert config.batch_threshold == 10000
        assert config.distributed is None  # Auto mode by default
        assert config.large_batch_confirm_threshold == 20000

    def test_batch_mode_config_custom(self):
        """Test custom batch mode configuration values."""
        config = ingestion_pipeline.ProcessingConfig(
            input_file="test.parquet",
            output_catalog="./output",
            scratch_location="./scratch",
            batch_threshold=5000,
            distributed=True,
            large_batch_confirm_threshold=15000,
        )

        assert config.batch_threshold == 5000
        assert config.distributed is True
        assert config.large_batch_confirm_threshold == 15000

    def test_batch_mode_in_to_dict(self):
        """Test that batch mode options are included in to_dict."""
        config = ingestion_pipeline.ProcessingConfig(
            input_file="test.parquet",
            output_catalog="./output",
            scratch_location="./scratch",
            batch_threshold=7500,
            distributed=False,
        )

        data = config.to_dict()
        assert data["batch_threshold"] == 7500
        assert data["distributed"] is False
        assert "large_batch_confirm_threshold" in data

    def test_batch_mode_from_dict(self):
        """Test restoring batch mode options from dict."""
        data = {
            "input_file": "test.parquet",
            "output_catalog": "./output",
            "scratch_location": "./scratch",
            "batch_threshold": 8000,
            "distributed": True,
            "large_batch_confirm_threshold": 25000,
        }

        config = ingestion_pipeline.ProcessingConfig.from_dict(data)
        assert config.batch_threshold == 8000
        assert config.distributed is True
        assert config.large_batch_confirm_threshold == 25000

    def test_should_use_distributed_force_true(self):
        """Test _should_use_distributed with forced distributed mode."""
        config = ingestion_pipeline.ProcessingConfig(
            input_file=str(Path(self.temp_dir) / "input.parquet"),
            output_catalog=str(Path(self.temp_dir) / "catalog"),
            scratch_location=str(Path(self.temp_dir) / "scratch"),
            distributed=True,
        )

        # Create minimal input file
        df = pd.DataFrame({"url": ["url1"]})
        df.to_parquet(config.input_file, index=False)

        processor = ingestion_pipeline.LocalProcessor(n_workers=1)
        pipeline = ingestion_pipeline.STACIngestionPipeline(config, processor)

        # Should use distributed even with 1 URL (well below threshold)
        assert pipeline._should_use_distributed(1) is True
        assert pipeline._should_use_distributed(100) is True
        processor.close()

    def test_should_use_distributed_force_false(self):
        """Test _should_use_distributed with forced local mode."""
        config = ingestion_pipeline.ProcessingConfig(
            input_file=str(Path(self.temp_dir) / "input.parquet"),
            output_catalog=str(Path(self.temp_dir) / "catalog"),
            scratch_location=str(Path(self.temp_dir) / "scratch"),
            distributed=False,
        )

        # Create minimal input file
        df = pd.DataFrame({"url": ["url1"]})
        df.to_parquet(config.input_file, index=False)

        processor = ingestion_pipeline.LocalProcessor(n_workers=1)
        pipeline = ingestion_pipeline.STACIngestionPipeline(config, processor)

        # Should use local even with many URLs (above threshold)
        assert pipeline._should_use_distributed(1) is False
        assert pipeline._should_use_distributed(100000) is False
        processor.close()

    def test_should_use_distributed_auto_mode(self):
        """Test _should_use_distributed with auto mode."""
        config = ingestion_pipeline.ProcessingConfig(
            input_file=str(Path(self.temp_dir) / "input.parquet"),
            output_catalog=str(Path(self.temp_dir) / "catalog"),
            scratch_location=str(Path(self.temp_dir) / "scratch"),
            batch_threshold=1000,
            distributed=None,  # Auto mode
        )

        # Create minimal input file
        df = pd.DataFrame({"url": ["url1"]})
        df.to_parquet(config.input_file, index=False)

        processor = ingestion_pipeline.LocalProcessor(n_workers=1)
        pipeline = ingestion_pipeline.STACIngestionPipeline(config, processor)

        # Below threshold -> local
        assert pipeline._should_use_distributed(500) is False
        assert pipeline._should_use_distributed(999) is False

        # At or above threshold -> distributed
        assert pipeline._should_use_distributed(1000) is True
        assert pipeline._should_use_distributed(5000) is True
        processor.close()


if __name__ == "__main__":
    pytest.main([__file__])
