"""STAC I/O engine abstraction layer.

This module provides pluggable backends for STAC item conversion and GeoParquet I/O,
enabling EarthCatalog to use different underlying implementations without code changes.

Supported Engines:
    rustac: High-performance Rust-based engine using the rustac library.
        - Zero Python dependencies for core operations
        - Native async support for I/O operations
        - Built-in object store integration (S3, GCS, Azure)
        - Recommended for production use

    stac-geoparquet: Legacy Python-based engine using stac-geoparquet library.
        - Pure Python implementation
        - Well-tested and stable
        - Fallback option for compatibility

Engine Selection:
    The engine can be selected via ProcessingConfig.stac_engine parameter:
    - "rustac" (default): Use rustac engine
    - "stac-geoparquet": Use legacy stac-geoparquet engine
    - "auto": Auto-detect based on available libraries (prefers rustac)

Example:
    >>> from earthcatalog.engines import get_engine
    >>> engine = get_engine("rustac")
    >>> gdf = engine.items_to_geodataframe(items)
    >>> items = engine.geodataframe_to_items(gdf)

Performance:
    The rustac engine provides significant performance improvements:
    - Faster Arrow conversion through Rust implementation
    - Native async I/O for cloud storage operations
    - Reduced memory overhead for large datasets
"""

import io
import logging
from abc import ABC, abstractmethod
from typing import Any, Literal, cast

import geopandas as gpd
import pyarrow.parquet as pq

logger = logging.getLogger(__name__)

EngineType = Literal["rustac", "stac-geoparquet", "auto"]
CompressionType = Literal["snappy", "gzip", "brotli"]


class GeoParquetIOMixin:
    """Mixin providing shared GeoParquet I/O operations for STAC engines.

    This mixin contains the common implementation for reading and writing
    GeoParquet files that is shared between all engine implementations.
    Extracting this logic ensures bug fixes and improvements are applied
    consistently across all engines.
    """

    def write_geoparquet_sync(
        self,
        gdf: gpd.GeoDataFrame,
        path: str,
        storage: Any,
        compression: str = "snappy",
    ) -> None:
        """Write GeoDataFrame to GeoParquet synchronously.

        Uses GeoPandas' native to_parquet() method with storage backend
        support for cloud paths.

        Args:
            gdf: GeoDataFrame to write.
            path: Output path (local or s3://).
            storage: Storage backend for cloud operations.
            compression: Parquet compression (default: snappy).

        Raises:
            OSError: If write operation fails.
        """
        if gdf.empty:
            logger.warning(f"Skipping write of empty GeoDataFrame to {path}")
            return

        try:
            # Cast compression to the expected literal type
            comp = cast(CompressionType, compression) if compression in ("snappy", "gzip", "brotli") else "snappy"

            if path.startswith("s3://"):
                # Use storage backend for S3
                with storage.open(path, "wb") as f:
                    gdf.to_parquet(f, index=False, compression=comp)
            else:
                # Local filesystem - direct write
                gdf.to_parquet(path, index=False, compression=comp)

        except Exception as e:
            logger.error(f"Error writing GeoParquet to {path}: {e}")
            raise OSError(f"Failed to write GeoParquet to {path}: {e}") from e

    def read_geoparquet_sync(self, path: str, storage: Any) -> gpd.GeoDataFrame:
        """Read GeoParquet file to GeoDataFrame synchronously.

        Uses PyArrow for reading with storage backend support for cloud paths.

        Args:
            path: Input path (local or s3://).
            storage: Storage backend for cloud operations.

        Returns:
            GeoDataFrame containing the data.

        Raises:
            FileNotFoundError: If file does not exist.
            OSError: If read fails.
        """
        try:
            if path.startswith("s3://"):
                # Use storage backend for S3
                with storage.open(path, "rb") as f:
                    binary_data = f.read()
                    table = pq.read_table(io.BytesIO(binary_data))
            else:
                # Local filesystem
                table = pq.read_table(path)

            # Convert to GeoDataFrame (preserve geometry column type)
            gdf = gpd.GeoDataFrame.from_arrow(table)

            return gdf

        except FileNotFoundError:
            raise
        except Exception as e:
            logger.error(f"Error reading GeoParquet from {path}: {e}")
            raise OSError(f"Failed to read GeoParquet from {path}: {e}") from e


class STACEngine(ABC):
    """Abstract base class for STAC I/O operations.

    This class defines the interface that all STAC engines must implement.
    Engines handle the conversion between STAC items and GeoDataFrames,
    as well as reading/writing GeoParquet files.

    The interface is designed to be synchronous for GeoDataFrame operations
    (which require in-memory processing) and provides both sync and async
    options for I/O operations.

    Thread Safety:
        Engine instances should be thread-safe for read operations.
        Write operations may require external synchronization.
    """

    @property
    @abstractmethod
    def name(self) -> str:
        """Return the engine name identifier."""
        pass

    @abstractmethod
    def items_to_geodataframe(self, items: list[dict[str, Any]]) -> gpd.GeoDataFrame:
        """Convert STAC items to a GeoDataFrame for sorting and processing.

        This is a synchronous operation as it requires in-memory processing
        of the entire item set for operations like sorting and deduplication.

        Args:
            items: List of STAC item dictionaries.

        Returns:
            GeoDataFrame with STAC item data, geometry column, and all properties.

        Raises:
            ValueError: If items cannot be converted to GeoDataFrame.
        """
        pass

    @abstractmethod
    def geodataframe_to_items(self, gdf: gpd.GeoDataFrame) -> list[dict[str, Any]]:
        """Convert a GeoDataFrame back to STAC item dictionaries.

        Args:
            gdf: GeoDataFrame containing STAC item data.

        Returns:
            List of STAC item dictionaries.

        Raises:
            ValueError: If GeoDataFrame cannot be converted to STAC items.
        """
        pass

    @abstractmethod
    def write_geoparquet_sync(
        self,
        gdf: gpd.GeoDataFrame,
        path: str,
        storage: Any,
        compression: str = "snappy",
    ) -> None:
        """Write a GeoDataFrame to GeoParquet file synchronously.

        This method handles both local and cloud storage paths using the
        provided storage backend for cloud operations.

        Args:
            gdf: GeoDataFrame to write.
            path: Output path (local or cloud URL like s3://).
            storage: Storage backend instance for cloud operations.
            compression: Parquet compression algorithm (default: snappy).

        Raises:
            IOError: If write operation fails.
        """
        pass

    @abstractmethod
    def read_geoparquet_sync(self, path: str, storage: Any) -> gpd.GeoDataFrame:
        """Read a GeoParquet file to GeoDataFrame synchronously.

        Args:
            path: Input path (local or cloud URL).
            storage: Storage backend instance for cloud operations.

        Returns:
            GeoDataFrame containing the data.

        Raises:
            FileNotFoundError: If the file does not exist.
            IOError: If read operation fails.
        """
        pass


class EngineNotAvailableError(Exception):
    """Raised when the requested engine is not available."""

    pass


def get_engine(engine_type: EngineType = "rustac") -> STACEngine:
    """Factory function to get the appropriate STAC engine.

    Args:
        engine_type: Engine to use. Options:
            - "rustac": Use rustac engine (default, recommended)
            - "stac-geoparquet": Use legacy stac-geoparquet engine
            - "auto": Auto-detect best available engine (prefers rustac)

    Returns:
        STACEngine instance configured for the requested backend.

    Raises:
        EngineNotAvailableError: If the requested engine is not available.
        ValueError: If an invalid engine type is specified.

    Example:
        >>> engine = get_engine("rustac")
        >>> gdf = engine.items_to_geodataframe(items)
    """
    if engine_type == "auto":
        # Try rustac first, fall back to stac-geoparquet
        try:
            from .rustac_engine import RustacEngine

            return RustacEngine()
        except ImportError:
            try:
                from .stac_geoparquet_engine import StacGeoparquetEngine

                return StacGeoparquetEngine()
            except ImportError:
                raise EngineNotAvailableError(
                    "No STAC engine available. Install rustac[arrow] or stac-geoparquet."
                ) from None

    elif engine_type == "rustac":
        try:
            from .rustac_engine import RustacEngine

            return RustacEngine()
        except ImportError as e:
            raise EngineNotAvailableError(
                f"rustac engine not available: {e}. Install with: pip install 'rustac[arrow]'"
            ) from e

    elif engine_type == "stac-geoparquet":
        try:
            from .stac_geoparquet_engine import StacGeoparquetEngine

            return StacGeoparquetEngine()
        except ImportError as e:
            raise EngineNotAvailableError(
                f"stac-geoparquet engine not available: {e}. Install with: pip install stac-geoparquet"
            ) from e

    else:
        raise ValueError(f"Unknown engine type: {engine_type}. Valid options: rustac, stac-geoparquet, auto")


__all__ = [
    "STACEngine",
    "GeoParquetIOMixin",
    "EngineType",
    "EngineNotAvailableError",
    "get_engine",
]
