"""Tests for spatial partition resolver functionality."""

import json
import shutil
import tempfile
from pathlib import Path

import pytest
from shapely.geometry import Point, box

# Import from earthcatalog package
from earthcatalog.spatial_resolver import (
    SpatialPartitionResolver,
    resolve_and_query,
    spatial_resolver,
)


class TestSpatialPartitionResolver:
    """Test SpatialPartitionResolver class."""

    def setup_method(self):
        """Set up test fixtures."""
        self.temp_dir = tempfile.mkdtemp()
        self.catalog_path = Path(self.temp_dir) / "catalog"
        self.catalog_path.mkdir(parents=True, exist_ok=True)

        # Create mock schema configurations
        self.h3_schema = {
            "spatial_partitioning": {
                "grid_system": "h3",
                "resolution": 6,
                "coordinate_system": "EPSG:4326",
                "example_paths": [
                    "testmission/partition=h3/level=6/86283082fffffff/year=2024/month=01/items.parquet",
                    "testmission/partition=h3/level=6/global/year=2024/month=01/items.parquet",
                ],
            },
            "global_partitioning": {
                "enabled": True,
                "threshold": 50,
                "description": "Items spanning >50 H3 cells go to /global/",
            },
        }

        self.geojson_schema = {
            "spatial_partitioning": {
                "grid_system": "geojson",
                "custom_grid": True,
                "geojson_source": str(Path(__file__).parent.parent / "examples" / "custom_tiles.geojson"),
                "custom_tiles": {
                    "total_tiles": 3,
                    "tile_ids": ["region_a", "region_b", "region_c"],
                },
            },
            "global_partitioning": {"enabled": False},
        }

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

    def test_resolver_initialization(self):
        """Test resolver initialization."""
        resolver = SpatialPartitionResolver(self.h3_schema, str(self.catalog_path))

        assert resolver.grid_system == "h3"
        assert resolver.global_enabled is True
        assert resolver.global_threshold == 50
        assert resolver.catalog_path == self.catalog_path

    def test_h3_point_resolution(self):
        """Test H3 point resolution."""
        resolver = SpatialPartitionResolver(self.h3_schema, str(self.catalog_path))

        # San Francisco coordinates
        sf_point = Point(-122.4194, 37.7749)
        partitions = resolver.resolve_partitions(sf_point)

        assert len(partitions) == 1
        assert isinstance(partitions[0], str)
        assert len(partitions[0]) == 15  # H3 cell ID length
        assert "global" not in partitions  # Small query shouldn't include global

    def test_h3_polygon_resolution(self):
        """Test H3 polygon resolution."""
        resolver = SpatialPartitionResolver(self.h3_schema, str(self.catalog_path))

        # San Francisco Bay Area
        sf_bay = box(-122.5, 37.5, -122.0, 38.0)
        partitions = resolver.resolve_partitions(sf_bay, overlap=True)

        assert len(partitions) > 10  # Should cover multiple cells
        assert all(isinstance(p, str) for p in partitions if p != "global")
        # Moderate area might include global depending on exact cell count

    def test_h3_buffer_cells(self):
        """Test H3 resolution with buffer cells."""
        resolver = SpatialPartitionResolver(self.h3_schema, str(self.catalog_path))

        sf_point = Point(-122.4194, 37.7749)

        # No buffer
        partitions_no_buffer = resolver.resolve_partitions(sf_point, buffer_cells=0)

        # With buffer
        partitions_with_buffer = resolver.resolve_partitions(sf_point, buffer_cells=2)

        assert len(partitions_no_buffer) == 1
        assert len(partitions_with_buffer) > len(partitions_no_buffer)
        assert partitions_no_buffer[0] in partitions_with_buffer

    def test_global_partition_threshold_logic(self):
        """Test global partition inclusion based on threshold."""
        resolver = SpatialPartitionResolver(self.h3_schema, str(self.catalog_path))

        # Small area - should not include global
        small_area = Point(-122.4194, 37.7749).buffer(0.01)
        small_partitions = resolver.resolve_partitions(small_area)
        assert "global" not in small_partitions

        # Large area - should include global due to threshold
        large_area = box(-125.0, 32.0, -115.0, 42.0)  # Entire California
        large_partitions = resolver.resolve_partitions(large_area)
        assert "global" in large_partitions
        assert len([p for p in large_partitions if p != "global"]) > 50

    def test_global_partition_geometry_size_logic(self):
        """Test global partition inclusion based on geometry size."""
        resolver = SpatialPartitionResolver(self.h3_schema, str(self.catalog_path))

        # Very large geometry (continental scale)
        continental = box(-130.0, 25.0, -65.0, 50.0)
        partitions = resolver.resolve_partitions(continental)

        assert "global" in partitions
        # Should trigger large geometry threshold even if cell count logic doesn't

    def test_global_partition_manual_override(self):
        """Test manual global partition inclusion/exclusion."""
        resolver = SpatialPartitionResolver(self.h3_schema, str(self.catalog_path))

        small_area = Point(-122.4194, 37.7749).buffer(0.01)
        large_area = box(-125.0, 32.0, -115.0, 42.0)

        # Force global on small area
        partitions_force_global = resolver.resolve_partitions(small_area, include_global=True)
        assert "global" in partitions_force_global

        # Disable global on large area
        partitions_no_global = resolver.resolve_partitions(large_area, include_global=False)
        assert "global" not in partitions_no_global

    def test_global_partition_disabled(self):
        """Test behavior when global partitioning is disabled."""
        schema_no_global = {
            "spatial_partitioning": {"grid_system": "h3", "resolution": 6},
            "global_partitioning": {"enabled": False},
        }

        resolver = SpatialPartitionResolver(schema_no_global, str(self.catalog_path))

        # Even large areas shouldn't include global when disabled
        large_area = box(-125.0, 32.0, -115.0, 42.0)
        partitions = resolver.resolve_partitions(large_area)

        assert "global" not in partitions
        assert resolver.global_enabled is False

    def test_geojson_resolution(self):
        """Test GeoJSON grid resolution."""
        resolver = SpatialPartitionResolver(self.geojson_schema, str(self.catalog_path))

        # Area that should intersect region_a (Northern region: -120 to -110, 40 to 50)
        northern_area = box(-119, 41, -111, 49)
        partitions = resolver.resolve_partitions(northern_area)

        assert "region_a" in partitions
        assert len(partitions) >= 1

        # Large area spanning multiple regions
        large_area = box(-119, 32, -105, 48)
        spanning_partitions = resolver.resolve_partitions(large_area)

        assert len(spanning_partitions) > 1
        assert "region_a" in spanning_partitions
        assert "region_b" in spanning_partitions

    def test_path_existence_checking(self):
        """Test existing partition path filtering."""
        resolver = SpatialPartitionResolver(self.h3_schema, str(self.catalog_path))

        # Create some test directories in file structure
        test_partitions = ["86283082fffffff", "862830827ffffff", "global"]
        existing_partitions = test_partitions[:2]  # Only create first two

        for partition_id in existing_partitions:
            partition_path = self.catalog_path / "testmission" / "partition=h3" / "level=6" / partition_id
            partition_path.mkdir(parents=True, exist_ok=True)

        # Test path filtering
        existing_paths = resolver.get_existing_partition_paths(test_partitions)

        assert len(existing_paths) == len(existing_partitions)
        for partition_id in existing_partitions:
            assert any(partition_id in path for path in existing_paths)

    def test_query_path_generation(self):
        """Test query path pattern generation with Hive-style temporal partitioning."""
        resolver = SpatialPartitionResolver(self.h3_schema, str(self.catalog_path))

        # Create test directories in file structure
        test_partitions = ["86283082fffffff", "global"]
        for partition_id in test_partitions:
            partition_path = self.catalog_path / "testmission" / "partition=h3" / "level=6" / partition_id
            partition_path.mkdir(parents=True, exist_ok=True)

        # Test without temporal filter - should glob all items files
        query_patterns = resolver.generate_query_paths(test_partitions)
        assert len(query_patterns) == len(test_partitions)
        assert all("/**/items.parquet" in pattern for pattern in query_patterns)

        # Test with temporal filter - should convert to Hive-style path
        temporal_patterns = resolver.generate_query_paths(test_partitions, "2024-01")
        assert len(temporal_patterns) == len(test_partitions)
        assert all("year=2024/month=01/items.parquet" in pattern for pattern in temporal_patterns)

    def test_different_grid_systems(self):
        """Test resolution across different grid systems."""
        test_point = Point(-122.4194, 37.7749)

        grid_schemas = {
            "h3": {"spatial_partitioning": {"grid_system": "h3", "resolution": 6}},
            "s2": {"spatial_partitioning": {"grid_system": "s2", "level": 13}},
            "utm": {"spatial_partitioning": {"grid_system": "utm", "precision": 1}},
            "latlon": {"spatial_partitioning": {"grid_system": "latlon", "cell_size_degrees": 0.1}},
        }

        for grid_name, schema in grid_schemas.items():
            schema["global_partitioning"] = {"enabled": False}  # Disable for simplicity

            try:
                resolver = SpatialPartitionResolver(schema, str(self.catalog_path))
                partitions = resolver.resolve_partitions(test_point)

                assert len(partitions) >= 1
                assert all(isinstance(p, str) for p in partitions)

            except ImportError:
                # Skip if optional dependencies not available
                pytest.skip(f"Optional dependency for {grid_name} grid not available")

    def test_effective_global_threshold(self):
        """Test effective global threshold calculation."""
        resolver = SpatialPartitionResolver(self.h3_schema, str(self.catalog_path))

        # Should use configured threshold
        threshold = resolver._get_effective_global_threshold()
        assert threshold == 50

        # Test with custom thresholds in schema
        schema_with_custom = self.h3_schema.copy()
        schema_with_custom["custom_thresholds"] = {"h3": {"6": 100, "7": 200}}

        resolver_custom = SpatialPartitionResolver(schema_with_custom, str(self.catalog_path))
        custom_threshold = resolver_custom._get_effective_global_threshold()
        assert custom_threshold == 100  # Should use custom threshold for H3 resolution 6

    def test_parse_temporal_filter_to_hive(self):
        """Test temporal filter to Hive path conversion."""
        resolver = SpatialPartitionResolver(self.h3_schema, str(self.catalog_path))

        # Test year only
        assert resolver._parse_temporal_filter_to_hive("2024") == "year=2024"

        # Test year-month
        assert resolver._parse_temporal_filter_to_hive("2024-01") == "year=2024/month=01"

        # Test year-month-day
        assert resolver._parse_temporal_filter_to_hive("2024-01-15") == "year=2024/month=01/day=15"

        # Test year with wildcard
        assert resolver._parse_temporal_filter_to_hive("2024-*") == "year=2024/*"

        # Test year-month with wildcard
        assert resolver._parse_temporal_filter_to_hive("2024-01-*") == "year=2024/month=01/*"

        # Test empty or wildcard only
        assert resolver._parse_temporal_filter_to_hive("*") == "*"
        assert resolver._parse_temporal_filter_to_hive("") == "*"

        # Test single-digit month padding
        assert resolver._parse_temporal_filter_to_hive("2024-1") == "year=2024/month=01"


class TestConvenienceFunctions:
    """Test convenience functions."""

    def setup_method(self):
        """Set up test fixtures."""
        self.temp_dir = tempfile.mkdtemp()
        self.schema_file = Path(self.temp_dir) / "schema.json"
        self.catalog_path = Path(self.temp_dir) / "catalog"

        # Create test schema
        schema = {
            "spatial_partitioning": {"grid_system": "h3", "resolution": 7},
            "global_partitioning": {"enabled": True, "threshold": 50},
        }

        with open(self.schema_file, "w") as f:
            json.dump(schema, f)

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

    def test_create_resolver_from_schema_file(self):
        """Test creating resolver from schema file (using modern spatial_resolver function)."""
        resolver = spatial_resolver(str(self.schema_file), str(self.catalog_path))

        assert isinstance(resolver, SpatialPartitionResolver)
        assert resolver.grid_system == "h3"
        assert resolver.spatial_config["resolution"] == 7

    def test_spatial_resolver_function(self):
        """Test the new spatial_resolver function."""
        resolver = spatial_resolver(str(self.schema_file), str(self.catalog_path))

        assert isinstance(resolver, SpatialPartitionResolver)
        assert resolver.grid_system == "h3"
        assert resolver.spatial_config["resolution"] == 7

        # Test with catalog_path defaulting to schema directory
        resolver2 = spatial_resolver(str(self.schema_file))
        assert isinstance(resolver2, SpatialPartitionResolver)

    def test_remote_schema_support(self):
        """Test remote URL support in spatial_resolver function."""
        # Test error when catalog_path not provided for various remote URLs
        remote_urls = [
            "s3://bucket/schema.json",
            "s3a://bucket/schema.json",
            "s3n://bucket/schema.json",
            "gcs://bucket/schema.json",
            "gs://bucket/schema.json",
            "https://example.com/schema.json",
            "abfs://container/schema.json",
        ]

        for url in remote_urls:
            with pytest.raises(ValueError, match="catalog_path must be explicitly provided"):
                spatial_resolver(url)

    def test_resolve_and_query_function(self):
        """Test the convenience resolve_and_query function."""
        # Create a test geometry
        test_geometry = box(-122.5, 37.7, -122.0, 38.0)

        partition_ids, query = resolve_and_query(
            schema_path=str(self.schema_file),
            catalog_path=str(self.catalog_path),
            aoi_geometry=test_geometry,
            temporal_filter="2024-*",
            overlap=True,
        )

        assert isinstance(partition_ids, list)
        assert len(partition_ids) > 0
        assert isinstance(query, str)
        # Query will be empty if no existing partition paths, which is expected in test


class TestErrorHandling:
    """Test error handling and edge cases."""

    def test_unsupported_grid_system(self):
        """Test error handling for unsupported grid systems."""
        schema = {
            "spatial_partitioning": {"grid_system": "invalid_grid"},
            "global_partitioning": {"enabled": False},
        }

        resolver = SpatialPartitionResolver(schema, "/tmp")

        with pytest.raises(ValueError, match="Unsupported grid system"):
            resolver.resolve_partitions(Point(0, 0))

    def test_missing_dependencies(self):
        """Test handling of missing optional dependencies."""
        # This test would need to mock the import failures
        # Implementation depends on how you want to handle missing dependencies
        pass

    def test_invalid_geometry(self):
        """Test handling of invalid geometries."""
        from shapely.errors import GeometryTypeError

        schema = {
            "spatial_partitioning": {"grid_system": "h3", "resolution": 6},
            "global_partitioning": {"enabled": False},
        }

        resolver = SpatialPartitionResolver(schema, "/tmp")

        # Test with invalid geometry dict
        with pytest.raises(GeometryTypeError):
            resolver.resolve_partitions({"type": "Invalid"})


if __name__ == "__main__":
    # Allow running tests directly
    pytest.main([__file__])
