"""Rustac-based STAC I/O engine for high-performance operations.

This module provides the rustac engine implementation using the Rust-based
rustac library for STAC item conversion and GeoParquet I/O. It offers significant
performance improvements over pure Python implementations.

Key Features:
    - High-performance Arrow conversion through Rust
    - Native object store support for S3, GCS, Azure
    - Memory-efficient processing for large datasets
    - Zero Python dependencies for core operations

Requirements:
    - rustac[arrow]>=0.9.0

Example:
    >>> from earthcatalog.engines.rustac_engine import RustacEngine
    >>> engine = RustacEngine()
    >>> gdf = engine.items_to_geodataframe(items)
    >>> engine.write_geoparquet_sync(gdf, "output.parquet", storage)
"""

import logging
from typing import Any

import geopandas as gpd

from . import GeoParquetIOMixin, STACEngine

logger = logging.getLogger(__name__)

# Import rustac - will raise ImportError if not available
try:
    import rustac

    HAS_RUSTAC = True
except ImportError:
    HAS_RUSTAC = False
    rustac = None  # type: ignore


class RustacEngine(GeoParquetIOMixin, STACEngine):
    """High-performance STAC engine using rustac library.

    This engine leverages rustac's Rust implementation for fast Arrow conversion
    and provides native support for cloud object stores.

    The engine uses synchronous wrappers around rustac's async functions for
    compatibility with the existing pipeline architecture. For direct async usage,
    use the rustac library directly.

    GeoParquet I/O operations are inherited from GeoParquetIOMixin to ensure
    consistent behavior across all engine implementations.

    Attributes:
        name: Engine identifier string ("rustac").

    Example:
        >>> engine = RustacEngine()
        >>> items = [{"type": "Feature", "id": "item1", ...}]
        >>> gdf = engine.items_to_geodataframe(items)
        >>> # Process GeoDataFrame (sort, filter, etc.)
        >>> items_out = engine.geodataframe_to_items(gdf)
    """

    def __init__(self) -> None:
        """Initialize the rustac engine.

        Raises:
            ImportError: If rustac library is not installed.
        """
        if not HAS_RUSTAC:
            raise ImportError("rustac library not available. Install with: pip install 'rustac[arrow]'")

    @property
    def name(self) -> str:
        """Return the engine name."""
        return "rustac"

    def items_to_geodataframe(self, items: list[dict[str, Any]]) -> gpd.GeoDataFrame:
        """Convert STAC items to GeoDataFrame using rustac.to_arrow().

        Uses rustac's Rust-based Arrow conversion for better performance
        compared to pure Python implementations.

        Args:
            items: List of STAC item dictionaries.

        Returns:
            GeoDataFrame with geometry and all STAC properties.

        Raises:
            ValueError: If items cannot be converted.
        """
        if not items:
            return gpd.GeoDataFrame()

        try:
            # rustac.to_arrow() accepts a list of items or an item collection
            table = rustac.to_arrow(items)  # type: ignore[union-attr]

            # Convert Arrow table to GeoDataFrame
            # GeoDataFrame.from_arrow() handles the geometry column correctly
            gdf = gpd.GeoDataFrame.from_arrow(table)

            return gdf

        except Exception as e:
            logger.error(f"Error converting items to GeoDataFrame with rustac: {e}")
            raise ValueError(f"Failed to convert STAC items to GeoDataFrame: {e}") from e

    def geodataframe_to_items(self, gdf: gpd.GeoDataFrame) -> list[dict[str, Any]]:
        """Convert GeoDataFrame back to STAC items using rustac.from_arrow().

        Args:
            gdf: GeoDataFrame containing STAC item data.

        Returns:
            List of STAC item dictionaries.

        Raises:
            ValueError: If GeoDataFrame cannot be converted.
        """
        if gdf.empty:
            return []

        try:
            # Convert GeoDataFrame to Arrow table
            table = gdf.to_arrow()

            # Use rustac to convert Arrow table to STAC item collection
            item_collection = rustac.from_arrow(table)  # type: ignore[union-attr, arg-type]

            # Extract features from the item collection
            if isinstance(item_collection, dict):
                return item_collection.get("features", [])
            else:
                return []

        except Exception as e:
            logger.error(f"Error converting GeoDataFrame to items with rustac: {e}")
            raise ValueError(f"Failed to convert GeoDataFrame to STAC items: {e}") from e
