backend/core/format_handlers.py

"""
Format Handlers for Multi-Format Data Ingestion

Provides a unified interface for loading various geospatial data formats
into GeoDataFrames for ingestion into the Perch catalog.
"""

import logging
from abc import ABC, abstractmethod
from pathlib import Path
from typing import Dict, Any, Optional, List, Tuple
import geopandas as gpd
import pandas as pd

logger = logging.getLogger(__name__)


class FormatHandler(ABC):
    """Base class for format-specific data loaders."""
    
    @property
    @abstractmethod
    def extensions(self) -> List[str]:
        """File extensions this handler supports (lowercase, with dot)."""
        pass
    
    @property
    @abstractmethod
    def format_name(self) -> str:
        """Human-readable format name."""
        pass
    
    def can_handle(self, path: str) -> bool:
        """Check if this handler can process the given path."""
        path_lower = path.lower()
        return any(path_lower.endswith(ext) for ext in self.extensions)
    
    @abstractmethod
    def load(self, path: str, **kwargs) -> pd.DataFrame:
        """Load the data file into a GeoDataFrame or DataFrame."""
        pass
    
    def get_metadata(self, path: str, gdf: pd.DataFrame) -> Dict[str, Any]:
        """Extract metadata from loaded data."""
        is_spatial = isinstance(gdf, gpd.GeoDataFrame)
        geometry_type = self._detect_geometry_type(gdf) if is_spatial else None
        
        return {
            "format": self.format_name,
            "row_count": len(gdf),
            "columns": list(gdf.columns),
            "geometry_type": geometry_type,
            "crs": str(gdf.crs) if is_spatial and gdf.crs else None,
            "bbox": list(gdf.total_bounds) if is_spatial and not gdf.empty else None,
        }
    
    def _detect_geometry_type(self, gdf: gpd.GeoDataFrame) -> Optional[str]:
        """Detect the predominant geometry type."""
        if gdf.empty or 'geometry' not in gdf.columns:
            return None
        geom_types = gdf.geometry.geom_type.value_counts()
        return geom_types.index[0] if not geom_types.empty else None


class GeoJSONHandler(FormatHandler):
    """Handler for GeoJSON files."""
    
    @property
    def extensions(self) -> List[str]:
        return ['.geojson', '.geojson.gz', '.json']
    
    @property
    def format_name(self) -> str:
        return "geojson"
    
    def load(self, path: str, **kwargs) -> gpd.GeoDataFrame:
        logger.info(f"Loading GeoJSON: {path}")
        gdf = gpd.read_file(path)
        
        # Ensure WGS84 CRS
        if gdf.crs is None:
            gdf = gdf.set_crs("EPSG:4326")
        elif gdf.crs.to_epsg() != 4326:
            gdf = gdf.to_crs("EPSG:4326")
            
        return gdf


class ShapefileHandler(FormatHandler):
    """Handler for ESRI Shapefiles."""
    
    @property
    def extensions(self) -> List[str]:
        return ['.shp', '.zip']  # .zip for zipped shapefiles
    
    @property
    def format_name(self) -> str:
        return "shapefile"
    
    def load(self, path: str, **kwargs) -> gpd.GeoDataFrame:
        logger.info(f"Loading Shapefile: {path}")
        gdf = gpd.read_file(path)
        
        # Ensure WGS84 CRS
        if gdf.crs is None:
            logger.warning(f"No CRS found for {path}, assuming EPSG:4326")
            gdf = gdf.set_crs("EPSG:4326")
        elif gdf.crs.to_epsg() != 4326:
            gdf = gdf.to_crs("EPSG:4326")
            
        return gdf


class GeoParquetHandler(FormatHandler):
    """Handler for Parquet files, spatial (GeoParquet) or plain tabular."""

    @property
    def extensions(self) -> List[str]:
        return ['.parquet', '.geoparquet']

    @property
    def format_name(self) -> str:
        return "geoparquet"

    def load(self, path: str, **kwargs) -> pd.DataFrame:
        logger.info(f"Loading Parquet: {path}")
        try:
            gdf = gpd.read_parquet(path)
        except (ValueError, AttributeError) as e:
            # Not a GeoParquet file (no geometry column). Attribute tables are a
            # first-class dataset type here, so fall back to plain pandas rather
            # than failing the ingest.
            logger.info(f"No geometry column found ({e}); loading as tabular Parquet.")
            return pd.read_parquet(path)

        # Ensure WGS84 CRS
        if gdf.crs is None:
            gdf = gdf.set_crs("EPSG:4326")
        elif gdf.crs.to_epsg() != 4326:
            gdf = gdf.to_crs("EPSG:4326")

        return gdf


class CSVHandler(FormatHandler):
    """Handler for CSV files (Spatial or Non-Spatial)."""
    
    # Common column names for coordinates
    LAT_COLUMNS = ['lat', 'latitude', 'y', 'lat_y', 'center_lat', 'centroid_lat']
    LON_COLUMNS = ['lon', 'lng', 'longitude', 'x', 'lon_x', 'center_lon', 'centroid_lon']
    
    @property
    def extensions(self) -> List[str]:
        return ['.csv']
    
    @property
    def format_name(self) -> str:
        return "csv"
    
    def _detect_coord_columns(self, df: pd.DataFrame) -> Tuple[Optional[str], Optional[str]]:
        """Auto-detect latitude and longitude columns."""
        columns_lower = {c.lower(): c for c in df.columns}
        
        lat_col = None
        lon_col = None
        
        for lat_name in self.LAT_COLUMNS:
            if lat_name in columns_lower:
                lat_col = columns_lower[lat_name]
                break
                
        for lon_name in self.LON_COLUMNS:
            if lon_name in columns_lower:
                lon_col = columns_lower[lon_name]
                break
                
        return lat_col, lon_col
    
    def load(self, path: str, lat_col: str = None, lon_col: str = None, **kwargs) -> pd.DataFrame:
        logger.info(f"Loading CSV: {path}")
        df = pd.read_csv(path)
        
        # Try to detect coordinates
        if lat_col is None or lon_col is None:
            lat_col, lon_col = self._detect_coord_columns(df)
            
        if lat_col and lon_col:
            logger.info(f"Detected spatial coordinates: lat={lat_col}, lon={lon_col}")
            # Create geometry
            gdf = gpd.GeoDataFrame(
                df,
                geometry=gpd.points_from_xy(df[lon_col], df[lat_col]),
                crs="EPSG:4326"
            )
            return gdf
        else:
            logger.info("No coordinate columns detected. Loading as non-spatial DataFrame (Attribute Table).")
            return df


class GeoPackageHandler(FormatHandler):
    """Handler for GeoPackage files."""
    
    @property
    def extensions(self) -> List[str]:
        return ['.gpkg']
    
    @property
    def format_name(self) -> str:
        return "geopackage"
    
    def load(self, path: str, layer: str = None, **kwargs) -> gpd.GeoDataFrame:
        logger.info(f"Loading GeoPackage: {path}")

        if layer is None:
            # Pick the first layer. Use pyogrio (geopandas' default engine) rather
            # than fiona, which is not a hard dependency of geopandas 1.x.
            import pyogrio
            layers = [row[0] for row in pyogrio.list_layers(path)]
            if not layers:
                raise ValueError(f"No layers found in GeoPackage: {path}")
            layer = layers[0]
            logger.info(f"Loading layer: {layer} (available: {layers})")

        gdf = gpd.read_file(path, layer=layer)

        # Ensure WGS84 CRS
        if gdf.crs is None:
            gdf = gdf.set_crs("EPSG:4326")
        elif gdf.crs.to_epsg() != 4326:
            gdf = gdf.to_crs("EPSG:4326")

        return gdf





class RasterHandler(FormatHandler):
    """Handler for Raster files (GeoTIFF), aggregating to H3 Hexagon grid."""
    
    @property
    def extensions(self) -> List[str]:
        return ['.tif', '.tiff']
    
    @property
    def format_name(self) -> str:
        return "raster"
    
    def load(self, path: str, **kwargs) -> gpd.GeoDataFrame:
        """
        Aggregate one raster band onto an H3 hexagon grid.

        `band` selects which band to read (1-based, default 1). Multi-band cubes —
        such as eBird's per-season abundance rasters — are handled by calling this
        once per band and concatenating, so each band keeps its own identity.
        """
        try:
            import rasterio
            import h3
            from shapely.geometry import Polygon
            import numpy as np
        except ImportError:
            raise ImportError("rasterio and h3 are required. Install with 'pip install rasterio h3'")

        band_index = int(kwargs.get('band', 1))
        logger.info(f"Loading Raster: {path} (band {band_index})")

        with rasterio.open(path) as src:
            logger.info(f"Raster size: {src.width}x{src.height}, Bands: {src.count}")

            if not 1 <= band_index <= src.count:
                raise ValueError(
                    f"Band {band_index} out of range for {path} (has {src.count} band(s))"
                )

            band1 = src.read(band_index)
            
            # H3 Resolution Strategy
            # Res 7 is approx 5km^2 (fine but manageable, ~15k cells for Panama)
            # Res 6 is approx 36km^2 (coarse, ~2k cells)
            # User request: "finer-grained" + "H3"
            h3_res = kwargs.get('h3_res', 7)
            
            # Subsampling Strategy for Aggregation
            # We need enough sample points to hit every H3 cell.
            # H3 Edge Lengths (approx km): 7=1.22, 8=0.46, 9=0.17, 10=0.066
            h3_edge_km = h3.average_hexagon_edge_length(h3_res, unit='km')
            pixel_size_x = src.transform[0] 
            # Assuming meters/degrees. If degrees (e.g. 0.00027), need conversion.
            # If coordinates are lat/lon (WGS84), pixel size is in degrees.
            # 1 deg ~ 111km.
            
            if src.crs.to_epsg() == 4326:
                pixel_km = abs(pixel_size_x) * 111
            else:
                 # Assume meters if projected (e.g. UTM)
                 pixel_km = abs(pixel_size_x) / 1000.0

            # Target step size: 1/3 of edge length to ensure hits
            target_step_km = h3_edge_km / 3.0
            step = max(1, int(target_step_km / pixel_km))
            
            logger.info(f"H3 Res {h3_res} edge={h3_edge_km:.3f}km, Pixel={pixel_km:.5f}km. Calc step={step}") 
            
            logger.info(f"Aggregating raster to H3 (Res {h3_res}) with subsample step {step}...")
            
            rows, cols = np.indices(band1.shape)
            rows = rows[::step, ::step].flatten()
            cols = cols[::step, ::step].flatten()
            values = band1[::step, ::step].flatten()

            # Filter nodata. NaN must be tested with isnan, not `!=`: for a NaN
            # nodata value (which eBird Status and Trends rasters use) `values !=
            # nodata` is True everywhere and would let every NaN cell through.
            mask = ~np.isnan(values) if np.issubdtype(values.dtype, np.floating) else np.ones(values.shape, bool)
            if src.nodata is not None and not (isinstance(src.nodata, float) and np.isnan(src.nodata)):
                mask &= (values != src.nodata)

            rows = rows[mask]
            cols = cols[mask]
            values = values[mask]
            
            # Transform pixel coordinates to source CRS coordinates
            xs, ys = rasterio.transform.xy(src.transform, rows, cols)
            xs = np.array(xs)
            ys = np.array(ys)
            
            # Reproject to WGS84 (EPSG:4326) if needed - H3 requires lat/lon
            if src.crs and src.crs.to_epsg() != 4326:
                from pyproj import Transformer
                transformer = Transformer.from_crs(src.crs, "EPSG:4326", always_xy=True)
                lons, lats = transformer.transform(xs, ys)
                logger.info(f"Reprojected {len(xs)} points from {src.crs} to EPSG:4326")
            else:
                lons, lats = xs, ys
            
            # Aggregate by H3
            hex_values = {} # h3_index -> [values]
            
            for lon, lat, v in zip(lons, lats, values):
                try:
                    # H3 v4 API - takes (lat, lon)
                    h3_idx = h3.latlng_to_cell(lat, lon, h3_res)
                except AttributeError:
                    # Fallback to v3 API if needed (though v4 is installed)
                    h3_idx = h3.geo_to_h3(lat, lon, h3_res)
                    
                if h3_idx not in hex_values:
                    hex_values[h3_idx] = []
                hex_values[h3_idx].append(v)
            
            logger.info(f"Aggregated into {len(hex_values)} H3 hexagons.")
            
            # Create Geometries and Compute Means
            geometries = []
            mean_values = []
            
            for h3_idx, vals in hex_values.items():
                mean_val = float(np.mean(vals))
                
                try:
                    # H3 v4 API
                    boundary = h3.cell_to_boundary(h3_idx)
                except AttributeError:
                    # v3 API
                    boundary = h3.h3_to_geo_boundary(h3_idx)
                
                # H3 returns (lat, lon) tuples. Shapely expects (lon, lat).
                coords = [(lon, lat) for lat, lon in boundary]

                # Un-wrap cells that straddle the antimeridian. Their vertices come
                # back as a mix of ~+180 and ~-180, which Shapely reads as a polygon
                # spanning the whole globe and renders as a streak across the map.
                # Shifting the negative vertices by +360 restores a contiguous ring.
                lons = [c[0] for c in coords]
                if max(lons) - min(lons) > 180:
                    coords = [(lon + 360 if lon < 0 else lon, lat) for lon, lat in coords]

                geometries.append(Polygon(coords))
                mean_values.append(mean_val)
                
            # Create GeoDataFrame
            df = pd.DataFrame({'value': mean_values, 'h3_index': list(hex_values.keys())})
            gdf = gpd.GeoDataFrame(
                df,
                geometry=geometries,
                crs="EPSG:4326"
            )
                
            return gdf


class ExcelHandler(FormatHandler):
    """Handler for Excel files."""
    
    @property
    def extensions(self) -> List[str]:
        return ['.xlsx', '.xls']
    
    @property
    def format_name(self) -> str:
        return "excel"
    
    def load(self, path: str, lat_col: str = None, lon_col: str = None, **kwargs) -> pd.DataFrame:
        logger.info(f"Loading Excel: {path}")
        # Use pandas to read excel
        # kwargs might contain sheet_name if we wanted to support it, but for now default
        # Force string to avoid mixed types in Parquet
        df = pd.read_excel(path, dtype=str)
        
        # Try to detect coordinates (logic shared with CSVHandler?)
        # Reuse logic by creating a dummy CSVHandler or extracting method?
        # Duplication is fine for now, or better: extract mixin.
        # But for speed, I'll copy the detection logic or call CSVHandler helper if it was static.
        # It's an instance method.
        # I'll just look for columns.
        
        CSVHandler_defaults = CSVHandler()
        if lat_col is None or lon_col is None:
            lat_col, lon_col = CSVHandler_defaults._detect_coord_columns(df)
            
        if lat_col and lon_col:
            logger.info(f"Detected spatial coordinates: lat={lat_col}, lon={lon_col}")
            # Create geometry
            gdf = gpd.GeoDataFrame(
                df,
                geometry=gpd.points_from_xy(df[lon_col], df[lat_col]),
                crs="EPSG:4326"
            )
            return gdf
        else:
            logger.info("No coordinate columns detected. Loading as non-spatial DataFrame.")
            return df


class GMLHandler(FormatHandler):
    """Handler for GML files (e.g. Catastro)."""
    
    @property
    def extensions(self) -> List[str]:
        return ['.gml']
    
    @property
    def format_name(self) -> str:
        return "gml"
    
    def load(self, path: str, **kwargs) -> pd.DataFrame:
        logger.info(f"Loading GML: {path}")
        return gpd.read_file(path)


# Registry of all handlers
HANDLERS: List[FormatHandler] = [
    GeoJSONHandler(),
    ShapefileHandler(),
    GeoParquetHandler(),
    CSVHandler(),
    ExcelHandler(),
    GMLHandler(),
    GeoPackageHandler(),
    RasterHandler(),
]


def get_handler(path: str) -> Optional[FormatHandler]:
    """Get the appropriate handler for a file path."""
    for handler in HANDLERS:
        if handler.can_handle(path):
            return handler
    return None


def get_supported_formats() -> List[str]:
    """Get list of supported file extensions."""
    extensions = []
    for handler in HANDLERS:
        extensions.extend(handler.extensions)
    return extensions