# test_pipeline_components.py
"""Tests for pipeline components: configurations, data classes, and utilities.

This module tests the new pipeline components including BatchConfig, ConsolidationConfig,
ShardInfo, PartitionResult, BatchResult, and utility functions.
"""

from unittest.mock import MagicMock

import geopandas as gpd
import pytest
from shapely.geometry import Point

from earthcatalog.pipeline import (
    BatchConfig,
    BatchResult,
    ConsolidationConfig,
    PartitionResult,
    ShardInfo,
    chunk_urls,
    group_shards_by_partition,
    merge_geodataframes,
)


class TestBatchConfig:
    """Tests for BatchConfig dataclass."""

    def test_default_values(self):
        """Test BatchConfig has sensible defaults."""
        config = BatchConfig()
        assert config.batch_size == 1000
        assert config.items_per_shard == 10000
        assert config.enable_concurrent_http is True
        assert config.concurrent_requests == 50
        assert config.connection_pool_size == 100
        assert config.request_timeout == 30
        assert config.retry_attempts == 3
        assert config.retry_delay == 1.0

    def test_custom_values(self):
        """Test BatchConfig with custom values."""
        config = BatchConfig(
            batch_size=500,
            items_per_shard=5000,
            enable_concurrent_http=False,
            concurrent_requests=25,
        )
        assert config.batch_size == 500
        assert config.items_per_shard == 5000
        assert config.enable_concurrent_http is False
        assert config.concurrent_requests == 25

    @pytest.mark.parametrize(
        "kwargs,error_pattern",
        [
            ({"batch_size": 0}, "batch_size must be positive"),
            ({"batch_size": -1}, "batch_size must be positive"),
            ({"items_per_shard": 0}, "items_per_shard must be positive"),
            ({"concurrent_requests": 0}, "concurrent_requests must be positive"),
            ({"connection_pool_size": -5}, "connection_pool_size must be positive"),
            ({"request_timeout": 0}, "request_timeout must be positive"),
            ({"retry_attempts": -1}, "retry_attempts must be non-negative"),
            ({"retry_delay": -0.5}, "retry_delay must be non-negative"),
        ],
        ids=[
            "batch_size_zero",
            "batch_size_negative",
            "items_per_shard_zero",
            "concurrent_requests_zero",
            "connection_pool_size_negative",
            "request_timeout_zero",
            "retry_attempts_negative",
            "retry_delay_negative",
        ],
    )
    def test_validation_rejects_invalid_values(self, kwargs, error_pattern):
        """Test validation rejects invalid configuration values."""
        with pytest.raises(ValueError, match=error_pattern):
            BatchConfig(**kwargs)

    def test_repr(self):
        """Test string representation."""
        config = BatchConfig(batch_size=100)
        repr_str = repr(config)
        assert "BatchConfig" in repr_str
        assert "batch_size=100" in repr_str

    def test_bool_valid(self):
        """Test __bool__ returns True for valid config."""
        config = BatchConfig()
        assert bool(config) is True

    def test_bool_would_be_invalid(self):
        """Test that a config that would be invalid fails at construction."""
        # Since validation happens in __post_init__, we can't create an invalid config
        # This test documents that behavior
        with pytest.raises(ValueError):
            BatchConfig(batch_size=-1)


class TestConsolidationConfig:
    """Tests for ConsolidationConfig dataclass."""

    def test_default_values(self):
        """Test ConsolidationConfig has sensible defaults."""
        config = ConsolidationConfig()
        assert config.strategy == "efficient"
        assert config.max_memory_per_partition_mb == 1024
        assert config.enable_streaming_merge is True
        assert config.s3_multipart_threshold_mb == 100
        assert config.sort_key == "datetime"
        assert config.sort_ascending is True
        assert config.deduplicate_key == "id"
        assert config.keep_duplicates == "last"

    def test_custom_values(self):
        """Test ConsolidationConfig with custom values."""
        config = ConsolidationConfig(
            strategy="legacy",
            max_memory_per_partition_mb=2048,
            enable_streaming_merge=False,
        )
        assert config.strategy == "legacy"
        assert config.max_memory_per_partition_mb == 2048
        assert config.enable_streaming_merge is False

    @pytest.mark.parametrize(
        "kwargs,error_pattern",
        [
            ({"strategy": "invalid"}, "strategy must be"),
            ({"max_memory_per_partition_mb": 0}, "max_memory_per_partition_mb must be positive"),
            ({"s3_multipart_threshold_mb": -10}, "s3_multipart_threshold_mb must be positive"),
            ({"keep_duplicates": "middle"}, "keep_duplicates must be"),
        ],
        ids=["invalid_strategy", "memory_zero", "s3_threshold_negative", "invalid_keep_duplicates"],
    )
    def test_validation_rejects_invalid_values(self, kwargs, error_pattern):
        """Test validation rejects invalid configuration values."""
        with pytest.raises(ValueError, match=error_pattern):
            ConsolidationConfig(**kwargs)

    def test_repr(self):
        """Test string representation."""
        config = ConsolidationConfig(strategy="legacy")
        repr_str = repr(config)
        assert "ConsolidationConfig" in repr_str
        assert "strategy='legacy'" in repr_str

    def test_bool_valid(self):
        """Test __bool__ returns True for valid config."""
        config = ConsolidationConfig()
        assert bool(config) is True


class TestShardInfo:
    """Tests for ShardInfo dataclass."""

    def test_basic_creation(self):
        """Test basic ShardInfo creation."""
        shard = ShardInfo(
            shard_path="/scratch/shards/shard.parquet",
            partition_key="dataset/partition=h3/level=2/82",
            item_count=100,
            worker_id="worker-0",
        )
        assert shard.shard_path == "/scratch/shards/shard.parquet"
        assert shard.partition_key == "dataset/partition=h3/level=2/82"
        assert shard.item_count == 100
        assert shard.worker_id == "worker-0"
        assert shard.shard_id == 0  # default

    def test_with_shard_id(self):
        """Test ShardInfo with custom shard_id."""
        shard = ShardInfo(
            shard_path="/path/shard.parquet",
            partition_key="p1",
            item_count=50,
            worker_id="w1",
            shard_id=5,
        )
        assert shard.shard_id == 5

    def test_repr(self):
        """Test string representation."""
        shard = ShardInfo(
            shard_path="/path/shard.parquet",
            partition_key="p1",
            item_count=50,
            worker_id="w1",
        )
        repr_str = repr(shard)
        assert "ShardInfo" in repr_str
        assert "items=50" in repr_str

    def test_bool_with_items(self):
        """Test __bool__ returns True when shard has items."""
        shard = ShardInfo(
            shard_path="/path/shard.parquet",
            partition_key="p1",
            item_count=10,
            worker_id="w1",
        )
        assert bool(shard) is True

    def test_bool_without_items(self):
        """Test __bool__ returns False when shard has no items."""
        shard = ShardInfo(
            shard_path="/path/shard.parquet",
            partition_key="p1",
            item_count=0,
            worker_id="w1",
        )
        assert bool(shard) is False

    def test_to_dict(self):
        """Test conversion to dictionary."""
        shard = ShardInfo(
            shard_path="/path/shard.parquet",
            partition_key="p1",
            item_count=50,
            worker_id="w1",
            shard_id=3,
        )
        d = shard.to_dict()
        assert d["shard_path"] == "/path/shard.parquet"
        assert d["partition_key"] == "p1"
        assert d["item_count"] == 50
        assert d["worker_id"] == "w1"
        assert d["shard_id"] == 3

    def test_from_dict(self):
        """Test creation from dictionary."""
        data = {
            "shard_path": "/path/shard.parquet",
            "partition_key": "p1",
            "item_count": 50,
            "worker_id": "w1",
            "shard_id": 3,
        }
        shard = ShardInfo.from_dict(data)
        assert shard.shard_path == "/path/shard.parquet"
        assert shard.partition_key == "p1"
        assert shard.item_count == 50
        assert shard.worker_id == "w1"
        assert shard.shard_id == 3

    def test_from_dict_minimal(self):
        """Test creation from minimal dictionary."""
        data = {
            "shard_path": "/path/shard.parquet",
            "item_count": 50,
            "worker_id": "w1",
        }
        shard = ShardInfo.from_dict(data)
        assert shard.shard_path == "/path/shard.parquet"
        assert shard.partition_key == ""  # default
        assert shard.shard_id == 0  # default


class TestPartitionResult:
    """Tests for PartitionResult dataclass."""

    def test_basic_creation(self):
        """Test basic PartitionResult creation."""
        result = PartitionResult(
            partition_key="dataset/partition=h3/level=2/82",
            item_count=1000,
            existing_count=500,
            new_count=550,
            duplicates_removed=50,
        )
        assert result.partition_key == "dataset/partition=h3/level=2/82"
        assert result.item_count == 1000
        assert result.existing_count == 500
        assert result.new_count == 550
        assert result.duplicates_removed == 50
        assert result.success is True
        assert result.error == ""

    def test_repr(self):
        """Test string representation."""
        result = PartitionResult(partition_key="p1", item_count=100)
        repr_str = repr(result)
        assert "PartitionResult" in repr_str
        assert "items=100" in repr_str
        assert "OK" in repr_str

    def test_repr_failed(self):
        """Test string representation for failed result."""
        result = PartitionResult(partition_key="p1", success=False, error="Disk full")
        repr_str = repr(result)
        assert "FAILED" in repr_str
        assert "Disk full" in repr_str

    @pytest.mark.parametrize(
        "item_count,success,expected_bool",
        [
            (100, True, True),  # success with items
            (0, True, False),  # success without items
            (100, False, False),  # failed with items
        ],
        ids=["success_with_items", "success_no_items", "failed"],
    )
    def test_bool_behavior(self, item_count, success, expected_bool):
        """Test __bool__ returns correct value based on success and item_count."""
        result = PartitionResult(partition_key="p1", item_count=item_count, success=success)
        assert bool(result) is expected_bool

    def test_to_dict(self):
        """Test conversion to dictionary."""
        result = PartitionResult(
            partition_key="p1",
            item_count=100,
            existing_count=50,
            new_count=60,
            duplicates_removed=10,
            final_path="/path/to/file.parquet",
        )
        d = result.to_dict()
        assert d["partition"] == "p1"
        assert d["item_count"] == 100
        assert d["existing_count"] == 50
        assert d["new_count"] == 60
        assert d["duplicates_removed"] == 10
        assert d["final_path"] == "/path/to/file.parquet"

    def test_from_dict(self):
        """Test creation from dictionary."""
        data = {
            "partition": "p1",
            "item_count": 100,
            "existing_count": 50,
            "new_count": 60,
            "duplicates_removed": 10,
        }
        result = PartitionResult.from_dict(data)
        assert result.partition_key == "p1"
        assert result.item_count == 100

    def test_empty_factory(self):
        """Test empty factory method."""
        result = PartitionResult.empty("p1")
        assert result.partition_key == "p1"
        assert result.item_count == 0
        assert result.success is True

    def test_failed_factory(self):
        """Test failed factory method."""
        result = PartitionResult.failed("p1", "Network error")
        assert result.partition_key == "p1"
        assert result.success is False
        assert result.error == "Network error"


class TestBatchResult:
    """Tests for BatchResult dataclass."""

    def test_basic_creation(self):
        """Test basic BatchResult creation."""
        result = BatchResult(
            worker_id="worker-0",
            urls_processed=100,
            urls_succeeded=95,
            urls_failed=5,
        )
        assert result.worker_id == "worker-0"
        assert result.urls_processed == 100
        assert result.urls_succeeded == 95
        assert result.urls_failed == 5
        assert result.shards == []

    def test_with_shards(self):
        """Test BatchResult with shards."""
        shards = [
            ShardInfo(shard_path="/a.parquet", partition_key="p1", item_count=50, worker_id="w1"),
            ShardInfo(shard_path="/b.parquet", partition_key="p2", item_count=45, worker_id="w1"),
        ]
        result = BatchResult(
            worker_id="w1",
            shards=shards,
            urls_processed=100,
            urls_succeeded=95,
            urls_failed=5,
        )
        assert len(result.shards) == 2
        assert result.total_items == 95

    def test_success_rate(self):
        """Test success_rate calculation."""
        result = BatchResult(
            worker_id="w1",
            urls_processed=100,
            urls_succeeded=80,
            urls_failed=20,
        )
        assert result.success_rate == 0.8

    def test_success_rate_zero_processed(self):
        """Test success_rate when no URLs processed."""
        result = BatchResult(worker_id="w1", urls_processed=0)
        assert result.success_rate == 0.0

    def test_total_items(self):
        """Test total_items property."""
        shards = [
            ShardInfo(shard_path="/a.parquet", partition_key="p1", item_count=100, worker_id="w1"),
            ShardInfo(shard_path="/b.parquet", partition_key="p2", item_count=200, worker_id="w1"),
            ShardInfo(shard_path="/c.parquet", partition_key="p3", item_count=300, worker_id="w1"),
        ]
        result = BatchResult(worker_id="w1", shards=shards)
        assert result.total_items == 600

    def test_repr(self):
        """Test string representation."""
        result = BatchResult(
            worker_id="w1",
            urls_processed=100,
            urls_succeeded=95,
            urls_failed=5,
        )
        repr_str = repr(result)
        assert "BatchResult" in repr_str
        assert "worker='w1'" in repr_str

    def test_bool_with_items(self):
        """Test __bool__ returns True when result has items."""
        shards = [ShardInfo(shard_path="/a.parquet", partition_key="p1", item_count=10, worker_id="w1")]
        result = BatchResult(worker_id="w1", shards=shards)
        assert bool(result) is True

    def test_bool_without_items(self):
        """Test __bool__ returns False when result has no items."""
        result = BatchResult(worker_id="w1")
        assert bool(result) is False

    def test_to_dict(self):
        """Test conversion to dictionary."""
        shards = [ShardInfo(shard_path="/a.parquet", partition_key="p1", item_count=50, worker_id="w1")]
        result = BatchResult(
            worker_id="w1",
            shards=shards,
            urls_processed=100,
            urls_succeeded=95,
            urls_failed=5,
        )
        d = result.to_dict()
        assert d["worker_id"] == "w1"
        assert d["urls_processed"] == 100
        assert d["total_items"] == 50
        assert d["success_rate"] == 0.95
        assert len(d["shards"]) == 1


class TestMergeGeoDataFrames:
    """Tests for merge_geodataframes function."""

    def test_merge_empty_list(self):
        """Test merging empty list returns empty GeoDataFrame."""
        result = merge_geodataframes([])
        assert isinstance(result, gpd.GeoDataFrame)
        assert len(result) == 0

    def test_merge_single_dataframe(self):
        """Test merging single dataframe returns copy."""
        gdf = gpd.GeoDataFrame({"id": ["a", "b"], "value": [1, 2], "geometry": [Point(0, 0), Point(1, 1)]})
        result = merge_geodataframes([gdf])
        assert len(result) == 2
        # Verify it's a copy, not the same object
        assert result is not gdf

    def test_merge_multiple_dataframes(self):
        """Test merging multiple dataframes."""
        gdf1 = gpd.GeoDataFrame({"id": ["a", "b"], "value": [1, 2], "geometry": [Point(0, 0), Point(1, 1)]})
        gdf2 = gpd.GeoDataFrame({"id": ["c", "d"], "value": [3, 4], "geometry": [Point(2, 2), Point(3, 3)]})
        result = merge_geodataframes([gdf1, gdf2])
        assert len(result) == 4

    def test_merge_with_deduplication(self):
        """Test merging with duplicate removal."""
        gdf1 = gpd.GeoDataFrame({"id": ["a", "b"], "value": [1, 2], "geometry": [Point(0, 0), Point(1, 1)]})
        gdf2 = gpd.GeoDataFrame({"id": ["b", "c"], "value": [20, 3], "geometry": [Point(1, 1), Point(2, 2)]})
        result = merge_geodataframes([gdf1, gdf2], deduplicate_key="id", keep="last")
        assert len(result) == 3
        # Check that "b" has the value from the second dataframe
        b_row = result[result["id"] == "b"]
        assert b_row["value"].iloc[0] == 20

    def test_merge_keep_first(self):
        """Test merging keeping first duplicate."""
        gdf1 = gpd.GeoDataFrame({"id": ["a", "b"], "value": [1, 2], "geometry": [Point(0, 0), Point(1, 1)]})
        gdf2 = gpd.GeoDataFrame({"id": ["b", "c"], "value": [20, 3], "geometry": [Point(1, 1), Point(2, 2)]})
        result = merge_geodataframes([gdf1, gdf2], deduplicate_key="id", keep="first")
        assert len(result) == 3
        # Check that "b" has the value from the first dataframe
        b_row = result[result["id"] == "b"]
        assert b_row["value"].iloc[0] == 2

    def test_merge_with_sort(self):
        """Test merging with sorting."""
        gdf1 = gpd.GeoDataFrame({"id": ["a", "c"], "value": [3, 1], "geometry": [Point(0, 0), Point(2, 2)]})
        gdf2 = gpd.GeoDataFrame({"id": ["b"], "value": [2], "geometry": [Point(1, 1)]})
        result = merge_geodataframes([gdf1, gdf2], sort_key="value", sort_ascending=True)
        assert list(result["value"]) == [1, 2, 3]

    def test_merge_sort_descending(self):
        """Test merging with descending sort."""
        gdf = gpd.GeoDataFrame(
            {"id": ["a", "b", "c"], "value": [1, 3, 2], "geometry": [Point(0, 0), Point(1, 1), Point(2, 2)]}
        )
        result = merge_geodataframes([gdf], sort_key="value", sort_ascending=False)
        assert list(result["value"]) == [3, 2, 1]


class TestGroupShardsByPartition:
    """Tests for group_shards_by_partition function."""

    def test_group_shard_info_objects(self):
        """Test grouping ShardInfo objects."""
        shards = [
            ShardInfo(shard_path="/a.parquet", partition_key="p1", item_count=10, worker_id="w1"),
            ShardInfo(shard_path="/b.parquet", partition_key="p1", item_count=20, worker_id="w2"),
            ShardInfo(shard_path="/c.parquet", partition_key="p2", item_count=15, worker_id="w1"),
        ]
        groups = group_shards_by_partition(shards)
        assert len(groups) == 2
        assert len(groups["p1"]) == 2
        assert len(groups["p2"]) == 1
        assert "/a.parquet" in groups["p1"]
        assert "/b.parquet" in groups["p1"]
        assert "/c.parquet" in groups["p2"]

    def test_group_dictionaries(self):
        """Test grouping dictionary-based shard info."""
        shards = [
            {"shard_path": "/a.parquet", "partition_key": "p1", "item_count": 10},
            {"shard_path": "/b.parquet", "partition_key": "p2", "item_count": 20},
        ]
        groups = group_shards_by_partition(shards)
        assert len(groups) == 2
        assert groups["p1"] == ["/a.parquet"]
        assert groups["p2"] == ["/b.parquet"]

    def test_group_empty_list(self):
        """Test grouping empty list."""
        groups = group_shards_by_partition([])
        assert groups == {}

    def test_group_skips_empty_partition_key(self):
        """Test that shards with empty partition keys are skipped."""
        shards = [
            ShardInfo(shard_path="/a.parquet", partition_key="", item_count=10, worker_id="w1"),
            ShardInfo(shard_path="/b.parquet", partition_key="p1", item_count=20, worker_id="w1"),
        ]
        groups = group_shards_by_partition(shards)
        assert len(groups) == 1
        assert "p1" in groups


class TestChunkUrls:
    """Tests for chunk_urls function."""

    @pytest.mark.parametrize(
        "urls,num_chunks,expected_total_urls,expected_chunks_range",
        [
            (["url1", "url2", "url3", "url4"], 2, 4, (2, 2)),  # even split
            (["url1", "url2", "url3", "url4", "url5"], 2, 5, (2, 3)),  # uneven split
            (["url1", "url2"], 10, 2, (1, 2)),  # more chunks than urls
            (["url1"], 3, 1, (1, 1)),  # single url
            ([], 5, 0, (0, 0)),  # empty list
        ],
        ids=["even_split", "uneven_split", "more_chunks_than_urls", "single_url", "empty_list"],
    )
    def test_chunk_behavior(self, urls, num_chunks, expected_total_urls, expected_chunks_range):
        """Test chunking behavior for various inputs."""
        chunks = chunk_urls(urls, num_chunks)
        total = sum(len(c) for c in chunks)
        assert total == expected_total_urls
        min_chunks, max_chunks = expected_chunks_range
        assert min_chunks <= len(chunks) <= max_chunks or (expected_total_urls == 0 and len(chunks) == 0)

    @pytest.mark.parametrize(
        "invalid_num_chunks",
        [0, -1],
        ids=["zero", "negative"],
    )
    def test_chunk_invalid_num_chunks(self, invalid_num_chunks):
        """Test that invalid num_chunks raises error."""
        with pytest.raises(ValueError, match="num_chunks must be positive"):
            chunk_urls(["url1"], invalid_num_chunks)


class TestIntegration:
    """Integration tests for pipeline components."""

    def test_shard_info_roundtrip(self):
        """Test ShardInfo can roundtrip through dict conversion."""
        original = ShardInfo(
            shard_path="/path/to/shard.parquet",
            partition_key="dataset/partition=h3/level=2/82",
            item_count=100,
            worker_id="worker-0-abc123",
            shard_id=5,
        )
        d = original.to_dict()
        restored = ShardInfo.from_dict(d)
        assert restored.shard_path == original.shard_path
        assert restored.partition_key == original.partition_key
        assert restored.item_count == original.item_count
        assert restored.worker_id == original.worker_id
        assert restored.shard_id == original.shard_id

    def test_partition_result_roundtrip(self):
        """Test PartitionResult can roundtrip through dict conversion."""
        original = PartitionResult(
            partition_key="p1",
            item_count=1000,
            existing_count=500,
            new_count=550,
            duplicates_removed=50,
            final_path="/path/to/partition.parquet",
        )
        d = original.to_dict()
        restored = PartitionResult.from_dict(d)
        assert restored.partition_key == original.partition_key
        assert restored.item_count == original.item_count
        assert restored.existing_count == original.existing_count
        assert restored.new_count == original.new_count
        assert restored.duplicates_removed == original.duplicates_removed

    def test_batch_result_with_stats(self):
        """Test BatchResult can hold stats object."""
        # Create a mock stats object
        mock_stats = MagicMock()
        mock_stats.total_items = 100

        result = BatchResult(
            worker_id="w1",
            urls_processed=100,
            urls_succeeded=95,
            urls_failed=5,
            stats=mock_stats,
        )
        assert result.stats is not None
        assert result.stats.total_items == 100

    def test_configs_work_together(self):
        """Test that BatchConfig and ConsolidationConfig can be used together."""
        batch_config = BatchConfig(
            batch_size=500,
            items_per_shard=5000,
            concurrent_requests=25,
        )
        consolidation_config = ConsolidationConfig(
            strategy="efficient",
            max_memory_per_partition_mb=512,
            enable_streaming_merge=True,
        )

        # Both should be valid
        assert bool(batch_config) is True
        assert bool(consolidation_config) is True

        # Can use values together
        assert batch_config.items_per_shard < consolidation_config.max_memory_per_partition_mb * 1000
