# test_workers.py
"""Tests for workers module - serializable worker functions.

This module tests:
- process_url_batch: URL batch processing
- consolidate_partition: Partition consolidation
- Result types: DownloadResult, ConsolidationResult
"""

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

import pytest

from earthcatalog.ingestion_pipeline import ProcessingConfig
from earthcatalog.statistics import IngestionStatistics
from earthcatalog.workers import (
    ConsolidationResult,
    DownloadResult,
    _download_stac_item,
    _get_partition_key,
    consolidate_partition,
    process_url_batch,
)


class TestDownloadResult:
    """Test DownloadResult dataclass."""

    def test_to_dict(self):
        """to_dict should serialize result."""
        stats = IngestionStatistics()
        stats.record_url_processed(success=True)
        result = DownloadResult(
            shards=[{"shard_path": "/path/to/shard.parquet"}],
            stats=stats,
            failed_urls=["http://failed.com"],
        )

        data = result.to_dict()

        assert data["shards"] == [{"shard_path": "/path/to/shard.parquet"}]
        assert data["failed_urls"] == ["http://failed.com"]
        # Stats contains url counts
        assert "urls_processed" in data["stats"]

    def test_from_dict(self):
        """from_dict should deserialize result."""
        data = {
            "shards": [{"shard_path": "/path/to/shard.parquet"}],
            "stats": {"items_processed": 10, "items_failed": 2},
            "failed_urls": ["http://failed.com"],
        }

        result = DownloadResult.from_dict(data)

        assert len(result.shards) == 1
        assert len(result.failed_urls) == 1


class TestConsolidationResult:
    """Test ConsolidationResult dataclass."""

    def test_to_dict(self):
        """to_dict should serialize result."""
        result = ConsolidationResult(
            partition_key="h3_82/2024/01",
            item_count=100,
            existing_count=50,
            new_count=50,
            duplicates_removed=5,
            final_path="/catalog/h3_82/2024/01/data.parquet",
        )

        data = result.to_dict()

        assert data["partition_key"] == "h3_82/2024/01"
        assert data["item_count"] == 100
        assert data["existing_count"] == 50
        assert data["new_count"] == 50
        assert data["duplicates_removed"] == 5


class TestDownloadStacItem:
    """Test _download_stac_item helper function."""

    def test_successful_download(self):
        """Should return parsed JSON on successful download."""
        mock_hook = MagicMock()
        mock_hook.fetch.return_value = {"id": "test-item", "type": "Feature"}

        with patch("earthcatalog.workers._get_stac_hook", return_value=mock_hook):
            result = _download_stac_item("http://example.com/item.json")

        assert result is not None
        assert result["id"] == "test-item"
        mock_hook.fetch.assert_called_once()

    def test_failed_download_returns_none(self):
        """Should return None after all retries fail."""
        mock_hook = MagicMock()
        mock_hook.fetch.return_value = None

        with patch("earthcatalog.workers._get_stac_hook", return_value=mock_hook):
            result = _download_stac_item("http://example.com/item.json", retry_attempts=1)

        assert result is None


class TestGetPartitionKey:
    """Test _get_partition_key helper function."""

    def test_partition_key_with_datetime(self):
        """Should generate correct partition key with datetime."""
        from earthcatalog.grid_systems import get_grid_system

        item = {
            "id": "test-item",
            "geometry": {"type": "Point", "coordinates": [-122.0, 37.0]},
            "properties": {"datetime": "2024-06-15T12:00:00Z"},
        }
        grid = get_grid_system("h3", resolution=2)

        key = _get_partition_key(
            item,
            grid_resolver=grid,
            temporal_bin="month",
            enable_global=False,
            global_threshold=50,
        )

        assert "2024" in key
        assert "06" in key

    def test_partition_key_with_year_bin(self):
        """Should generate year-level partition key."""
        from earthcatalog.grid_systems import get_grid_system

        item = {
            "id": "test-item",
            "geometry": {"type": "Point", "coordinates": [-122.0, 37.0]},
            "properties": {"datetime": "2024-06-15T12:00:00Z"},
        }
        grid = get_grid_system("h3", resolution=2)

        key = _get_partition_key(
            item,
            grid_resolver=grid,
            temporal_bin="year",
            enable_global=False,
            global_threshold=50,
        )

        assert "2024" in key
        # Month should not be in year-level binning
        assert key.count("/") == 1  # grid_cell/year


class TestProcessUrlBatch:
    """Test process_url_batch function."""

    @pytest.fixture
    def temp_dir(self):
        """Create a temporary directory."""
        with tempfile.TemporaryDirectory() as tmpdir:
            yield tmpdir

    @pytest.fixture
    def config_dict(self, temp_dir):
        """Create a config dict for testing."""
        config = ProcessingConfig(
            input_file="test.parquet",
            output_catalog=f"{temp_dir}/catalog",
            scratch_location=f"{temp_dir}/scratch",
            enable_concurrent_http=False,  # Disable async for testing
            batch_size=10,
        )
        return config.to_dict()

    def test_process_empty_batch(self, config_dict):
        """Should handle empty URL list."""
        result = process_url_batch(
            urls=[],
            worker_id=0,
            config_dict=config_dict,
            job_id="test-job",
        )

        assert result["shards"] == []
        assert result["failed_urls"] == []

    def test_process_batch_with_failures(self, config_dict):
        """Should record failed URLs."""
        mock_hook = MagicMock()
        mock_hook.fetch.return_value = None  # Simulate failed fetch
        mock_hook.fetch_batch.return_value = [None, None]  # Batch also fails

        with patch("earthcatalog.workers._get_stac_hook", return_value=mock_hook):
            result = process_url_batch(
                urls=["http://bad1.com", "http://bad2.com"],
                worker_id=0,
                config_dict=config_dict,
                job_id="test-job",
            )

        assert len(result["failed_urls"]) == 2

    @pytest.mark.skip(reason="Requires complex mocking of STAC engine serialization")
    def test_process_batch_successful(self, config_dict):
        """Should process URLs and write shards."""
        mock_item = {
            "id": "test-item",
            "type": "Feature",
            "geometry": {"type": "Point", "coordinates": [-122.0, 37.0]},
            "bbox": [-122.0, 37.0, -122.0, 37.0],
            "properties": {"datetime": "2024-06-15T12:00:00Z"},
            "stac_version": "1.0.0",
            "links": [],
        }

        mock_response = MagicMock()
        mock_response.json.return_value = mock_item
        mock_response.raise_for_status = MagicMock()

        with patch("earthcatalog.workers.requests.get", return_value=mock_response):
            result = process_url_batch(
                urls=["http://example.com/item1.json"],
                worker_id=0,
                config_dict=config_dict,
                job_id="test-job",
            )

        assert len(result["shards"]) >= 1
        assert result["failed_urls"] == []


class TestConsolidatePartition:
    """Test consolidate_partition function."""

    @pytest.fixture
    def temp_dir(self):
        """Create a temporary directory."""
        with tempfile.TemporaryDirectory() as tmpdir:
            yield tmpdir

    @pytest.fixture
    def config_dict(self, temp_dir):
        """Create a config dict for testing."""
        config = ProcessingConfig(
            input_file="test.parquet",
            output_catalog=f"{temp_dir}/catalog",
            scratch_location=f"{temp_dir}/scratch",
        )
        return config.to_dict()

    def test_consolidate_empty_shards(self, config_dict):
        """Should handle empty shard list."""
        result = consolidate_partition(
            partition_key="h3_82/2024/01",
            shard_paths=[],
            config_dict=config_dict,
        )

        assert result["partition_key"] == "h3_82/2024/01"
        assert result["item_count"] == 0

    def test_consolidate_with_shards(self, config_dict, temp_dir):
        """Should consolidate shards into final partition."""
        import geopandas as gpd
        from shapely.geometry import Point

        # Create a test shard
        scratch_dir = Path(temp_dir) / "scratch" / "shards"
        scratch_dir.mkdir(parents=True)

        shard_path = str(scratch_dir / "test_shard.parquet")
        gdf = gpd.GeoDataFrame(
            {
                "id": ["item1", "item2"],
                "datetime": ["2024-06-15T12:00:00Z", "2024-06-16T12:00:00Z"],
            },
            geometry=gpd.GeoSeries([Point(-122, 37), Point(-122, 37)]),
        )
        gdf.to_parquet(shard_path)

        result = consolidate_partition(
            partition_key="h3_82/2024/06",
            shard_paths=[shard_path],
            config_dict=config_dict,
        )

        assert result["partition_key"] == "h3_82/2024/06"
        assert result["item_count"] == 2
        assert Path(result["final_path"]).exists()

    def test_consolidate_deduplicates_by_datetime(self, config_dict, temp_dir):
        """Should keep newer item when deduplicating by ID."""
        import geopandas as gpd
        from shapely.geometry import Point

        scratch_dir = Path(temp_dir) / "scratch" / "shards"
        scratch_dir.mkdir(parents=True)

        # Create two shards with same ID but different datetimes
        shard1_path = str(scratch_dir / "shard1.parquet")
        gdf1 = gpd.GeoDataFrame(
            {
                "id": ["item1"],
                "datetime": ["2024-06-15T12:00:00Z"],  # Older
            },
            geometry=gpd.GeoSeries([Point(-122, 37)]),
        )
        gdf1.to_parquet(shard1_path)

        shard2_path = str(scratch_dir / "shard2.parquet")
        gdf2 = gpd.GeoDataFrame(
            {
                "id": ["item1"],
                "datetime": ["2024-06-16T12:00:00Z"],  # Newer
            },
            geometry=gpd.GeoSeries([Point(-122, 37)]),
        )
        gdf2.to_parquet(shard2_path)

        result = consolidate_partition(
            partition_key="h3_82/2024/06",
            shard_paths=[shard1_path, shard2_path],
            config_dict=config_dict,
        )

        assert result["item_count"] == 1  # Deduplicated to 1
        assert result["duplicates_removed"] == 1

        # Verify the newer item was kept
        final_gdf = gpd.read_parquet(result["final_path"])
        assert len(final_gdf) == 1
        assert "2024-06-16" in final_gdf.iloc[0]["datetime"]


class TestPickleability:
    """Test that worker functions can be pickled for Dask."""

    def test_process_url_batch_is_pickleable(self):
        """process_url_batch should be pickleable."""
        import pickle

        # Should not raise
        pickled = pickle.dumps(process_url_batch)
        restored = pickle.loads(pickled)
        assert callable(restored)

    def test_consolidate_partition_is_pickleable(self):
        """consolidate_partition should be pickleable."""
        import pickle

        pickled = pickle.dumps(consolidate_partition)
        restored = pickle.loads(pickled)
        assert callable(restored)


class TestDaskIntegration:
    """Tests for Dask distributed integration with serializable workers.

    These tests verify that the worker functions can be properly used
    with Dask's distributed computing model.
    """

    @pytest.fixture
    def temp_dir(self):
        """Create a temporary directory."""
        with tempfile.TemporaryDirectory() as tmpdir:
            yield tmpdir

    @pytest.fixture
    def config_dict(self, temp_dir):
        """Create a config dict for testing."""
        config = ProcessingConfig(
            input_file="test.parquet",
            output_catalog=f"{temp_dir}/catalog",
            scratch_location=f"{temp_dir}/scratch",
            enable_concurrent_http=False,
            batch_size=10,
        )
        return config.to_dict()

    def test_config_dict_roundtrip(self, config_dict):
        """Config dict should serialize and deserialize correctly."""
        restored = ProcessingConfig.from_dict(config_dict)

        assert restored.output_catalog == config_dict["output_catalog"]
        assert restored.scratch_location == config_dict["scratch_location"]
        assert restored.enable_concurrent_http == config_dict["enable_concurrent_http"]

    def test_process_url_batch_accepts_config_dict(self, config_dict):
        """process_url_batch should accept config_dict parameter."""
        # Should not raise
        result = process_url_batch(
            urls=[],
            worker_id=0,
            config_dict=config_dict,
            job_id="test-job-123",
        )

        assert "shards" in result
        assert "stats" in result
        assert "failed_urls" in result

    def test_consolidate_partition_accepts_config_dict(self, config_dict):
        """consolidate_partition should accept config_dict parameter."""
        result = consolidate_partition(
            partition_key="test/2024/01",
            shard_paths=[],
            config_dict=config_dict,
        )

        assert "partition_key" in result
        assert result["partition_key"] == "test/2024/01"
        assert result["item_count"] == 0

    def test_worker_functions_are_module_level(self):
        """Worker functions should be importable at module level for Dask."""
        from earthcatalog import workers

        # These should be module-level functions, not methods
        assert hasattr(workers, "process_url_batch")
        assert hasattr(workers, "consolidate_partition")
        assert callable(workers.process_url_batch)
        assert callable(workers.consolidate_partition)

    def test_worker_functions_have_correct_signature(self):
        """Worker functions should have correct signature for Dask usage."""
        import inspect

        # process_url_batch signature
        sig = inspect.signature(process_url_batch)
        params = list(sig.parameters.keys())
        assert "urls" in params
        assert "worker_id" in params
        assert "config_dict" in params
        assert "job_id" in params

        # consolidate_partition signature
        sig = inspect.signature(consolidate_partition)
        params = list(sig.parameters.keys())
        assert "partition_key" in params
        assert "shard_paths" in params
        assert "config_dict" in params

    def test_result_serialization_for_dask(self, config_dict):
        """Results should be serializable for Dask transmission."""
        import pickle

        result = process_url_batch(
            urls=[],
            worker_id=0,
            config_dict=config_dict,
            job_id="test",
        )

        # Should be pickleable (required for Dask)
        pickled = pickle.dumps(result)
        restored = pickle.loads(pickled)

        assert restored["shards"] == result["shards"]
        assert restored["failed_urls"] == result["failed_urls"]

    def test_consolidation_result_serialization(self, config_dict):
        """Consolidation results should be serializable."""
        import pickle

        result = consolidate_partition(
            partition_key="h3_82/2024/01",
            shard_paths=[],
            config_dict=config_dict,
        )

        pickled = pickle.dumps(result)
        restored = pickle.loads(pickled)

        assert restored["partition_key"] == result["partition_key"]
        assert restored["item_count"] == result["item_count"]


class TestDaskDistributedProcessorIntegration:
    """Tests for DaskDistributedProcessor using workers.py functions.

    These tests mock Dask to verify the correct functions are called.
    """

    @pytest.fixture
    def mock_dask_client(self):
        """Create a mock Dask client."""
        mock_client = MagicMock()
        mock_client.dashboard_link = "http://localhost:8787"
        mock_client.submit.return_value = MagicMock()
        mock_client.gather.return_value = [{"shards": [], "stats": {}, "failed_urls": []}]
        return mock_client

    def test_dask_processor_uses_workers_process_url_batch(self, mock_dask_client):
        """DaskDistributedProcessor should use workers.process_url_batch when config_dict provided."""
        # Create mock dd module
        mock_dd = MagicMock()
        mock_dd.Client.return_value = mock_dask_client

        with patch.dict("sys.modules", {"dask.distributed": mock_dd}):
            # Need to reimport to pick up mocked module
            from importlib import reload

            from earthcatalog import ingestion_pipeline

            reload(ingestion_pipeline)

            processor = ingestion_pipeline.DaskDistributedProcessor()

            config_dict = {
                "input_file": "test.parquet",
                "output_catalog": "./out",
                "scratch_location": "./scratch",
            }

            # Call process_urls with config_dict
            processor.process_urls(
                url_chunks=[["url1", "url2"]],
                process_fn=lambda x, y: None,  # Should be ignored
                config_dict=config_dict,
                job_id="test-job",
            )

            # Verify the submit call used process_url_batch
            assert mock_dask_client.submit.called
            call_args = mock_dask_client.submit.call_args
            # First positional arg should be the function
            called_func = call_args[0][0] if call_args[0] else call_args[1].get("func")
            # The function name should be process_url_batch
            assert called_func.__name__ == "process_url_batch"

    def test_dask_processor_uses_workers_consolidate_partition(self, mock_dask_client):
        """DaskDistributedProcessor should use workers.consolidate_partition when config_dict provided."""
        mock_dask_client.gather.return_value = [{"partition_key": "test", "item_count": 0}]
        mock_dd = MagicMock()
        mock_dd.Client.return_value = mock_dask_client

        with patch.dict("sys.modules", {"dask.distributed": mock_dd}):
            from importlib import reload

            from earthcatalog import ingestion_pipeline

            reload(ingestion_pipeline)

            processor = ingestion_pipeline.DaskDistributedProcessor()

            config_dict = {
                "input_file": "test.parquet",
                "output_catalog": "./out",
                "scratch_location": "./scratch",
            }

            processor.consolidate_shards(
                partition_items=[("test/2024/01", ["/path/shard1.parquet"])],
                consolidate_fn=lambda x, y: None,  # Should be ignored
                config_dict=config_dict,
            )

            assert mock_dask_client.submit.called
            call_args = mock_dask_client.submit.call_args
            called_func = call_args[0][0] if call_args[0] else call_args[1].get("func")
            assert called_func.__name__ == "consolidate_partition"

    def test_dask_processor_fallback_without_config_dict(self, mock_dask_client):
        """DaskDistributedProcessor should fall back to old behavior without config_dict."""
        mock_dask_client.gather.return_value = ["result1"]
        mock_dd = MagicMock()
        mock_dd.Client.return_value = mock_dask_client

        with patch.dict("sys.modules", {"dask.distributed": mock_dd}):
            from importlib import reload

            from earthcatalog import ingestion_pipeline

            reload(ingestion_pipeline)

            processor = ingestion_pipeline.DaskDistributedProcessor()

            custom_fn = MagicMock(return_value="custom_result")

            # Call without config_dict
            processor.process_urls(
                url_chunks=[["url1"]],
                process_fn=custom_fn,
                # No config_dict - should use custom_fn
            )

            assert mock_dask_client.submit.called
            call_args = mock_dask_client.submit.call_args
            called_func = call_args[0][0]
            # Should use the provided custom function, not process_url_batch
            assert called_func == custom_fn


class TestStacHookIntegrationWithWorkers:
    """Test STAC hook integration with workers.py."""

    @pytest.fixture
    def temp_dir(self):
        """Create a temporary directory."""
        with tempfile.TemporaryDirectory() as tmpdir:
            yield tmpdir

    def test_hook_config_in_config_dict(self, temp_dir):
        """stac_hook should be included in config_dict."""
        config = ProcessingConfig(
            input_file="test.parquet",
            output_catalog=f"{temp_dir}/catalog",
            scratch_location=f"{temp_dir}/scratch",
            stac_hook="module:my_module:my_func",
        )

        config_dict = config.to_dict()
        assert "stac_hook" in config_dict
        assert config_dict["stac_hook"] == "module:my_module:my_func"

    def test_hook_config_restored_from_dict(self, temp_dir):
        """stac_hook should be restored from config_dict."""
        config_dict = {
            "input_file": "test.parquet",
            "output_catalog": f"{temp_dir}/catalog",
            "scratch_location": f"{temp_dir}/scratch",
            "stac_hook": "script:/path/to/script.py",
        }

        config = ProcessingConfig.from_dict(config_dict)
        assert config.stac_hook == "script:/path/to/script.py"

    def test_default_hook_in_workers(self, temp_dir):
        """Workers should use default hook when not specified."""
        config = ProcessingConfig(
            input_file="test.parquet",
            output_catalog=f"{temp_dir}/catalog",
            scratch_location=f"{temp_dir}/scratch",
        )

        config_dict = config.to_dict()

        # process_url_batch should work with default hook
        result = process_url_batch(
            urls=[],
            worker_id=0,
            config_dict=config_dict,
            job_id="test",
        )

        assert result["failed_urls"] == []
