# workers.py
"""Serializable worker functions for distributed STAC ingestion.

This module provides **module-level functions** that can be pickled and sent
to remote Dask workers. Unlike the methods in `ingestion_pipeline.py`, these
functions do not capture `self` and receive all dependencies as explicit parameters.

Architecture:
    These functions are designed to be called from both:
    1. ThreadPoolExecutor (local processing)
    2. Dask distributed workers (cluster processing)

    Each function receives a serialized config dict and storage paths,
    avoiding closure captures that would prevent pickling.

Key Functions:
    process_url_batch: Download STAC items from URLs and write shards
    consolidate_partition: Merge shards for a partition with existing data

Usage:
    >>> from earthcatalog.workers import process_url_batch
    >>> from earthcatalog.ingestion_pipeline import ProcessingConfig
    >>>
    >>> config = ProcessingConfig(...)
    >>> config_dict = config.to_dict()
    >>>
    >>> # Can be called directly or via Dask
    >>> result = process_url_batch(
    ...     urls=["http://example.com/item1.json"],
    ...     worker_id=0,
    ...     config_dict=config_dict,
    ...     job_id="abc-123",
    ... )
"""

from __future__ import annotations

import asyncio
import logging
import tempfile
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path
from typing import TYPE_CHECKING, Any

import geopandas as gpd
import pandas as pd
from tqdm import tqdm

# Conditional async HTTP imports
try:
    from .async_http_client import HAS_ASYNC_HTTP, download_stac_items_async
except ImportError:
    HAS_ASYNC_HTTP = False

    async def download_stac_items_async(*args, **kwargs) -> list[dict[str, Any]]:  # type: ignore
        """Dummy async function when async HTTP is not available."""
        raise ImportError("Async HTTP client not available")


from .engines import get_engine
from .grid_systems import get_grid_system
from .statistics import IngestionStatistics
from .storage_backends import get_storage_backend

if TYPE_CHECKING:
    from .ingestion_pipeline import ProcessingConfig

logger = logging.getLogger(__name__)


# =============================================================================
# Worker Result Types
# =============================================================================


class DownloadResult:
    """Result of processing a batch of URLs.

    Attributes:
        shards: List of shard info dictionaries with paths and counts.
        stats: Statistics from this worker's processing.
        failed_urls: URLs that failed after all retries.
    """

    def __init__(
        self,
        shards: list[dict[str, Any]],
        stats: IngestionStatistics,
        failed_urls: list[str],
    ):
        self.shards = shards
        self.stats = stats
        self.failed_urls = failed_urls

    def to_dict(self) -> dict[str, Any]:
        """Serialize for transmission."""
        return {
            "shards": self.shards,
            "stats": {
                "urls_processed": self.stats.urls_processed,
                "urls_failed": self.stats.urls_failed,
            },
            "failed_urls": self.failed_urls,
        }

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> DownloadResult:
        """Deserialize from dictionary."""
        stats = IngestionStatistics()
        stats_data = data.get("stats", {})
        # Reconstruct stats from dict
        if "urls_processed" in stats_data:
            stats.urls_processed = stats_data["urls_processed"]
        if "urls_failed" in stats_data:
            stats.urls_failed = stats_data["urls_failed"]
        return cls(
            shards=data.get("shards", []),
            stats=stats,
            failed_urls=data.get("failed_urls", []),
        )


class ConsolidationResult:
    """Result of consolidating a partition.

    Attributes:
        partition_key: The partition that was consolidated.
        item_count: Total items in the final partition.
        existing_count: Items that existed before.
        new_count: New items added.
        duplicates_removed: Duplicates that were removed.
        final_path: Path to the final partition file.
    """

    def __init__(
        self,
        partition_key: str,
        item_count: int,
        existing_count: int,
        new_count: int,
        duplicates_removed: int = 0,
        final_path: str = "",
    ):
        self.partition_key = partition_key
        self.item_count = item_count
        self.existing_count = existing_count
        self.new_count = new_count
        self.duplicates_removed = duplicates_removed
        self.final_path = final_path

    def to_dict(self) -> dict[str, Any]:
        """Serialize for transmission."""
        return {
            "partition_key": self.partition_key,
            "item_count": self.item_count,
            "existing_count": self.existing_count,
            "new_count": self.new_count,
            "duplicates_removed": self.duplicates_removed,
            "final_path": self.final_path,
        }


# =============================================================================
# Helper Functions
# =============================================================================


def _get_stac_hook(hook_config: str) -> Any:
    """Get a STAC hook instance from configuration string.

    This is called on workers to instantiate the hook from the serialized config.

    Args:
        hook_config: Hook configuration string (e.g., "default", "module:pkg:func").

    Returns:
        Configured hook instance (BaseSTACHook).
    """
    from .stac_hooks import get_hook

    return get_hook(hook_config)


def _download_stac_item(
    url: str,
    timeout: int = 30,
    retry_attempts: int = 3,
    hook_config: str = "default",
) -> dict[str, Any] | None:
    """Download/generate a single STAC item from URL using configured hook.

    Args:
        url: URL to the STAC item JSON or data file.
        timeout: Request timeout in seconds.
        retry_attempts: Number of retry attempts.
        hook_config: STAC hook configuration string.

    Returns:
        Parsed STAC item dictionary, or None if download failed.
    """
    hook = _get_stac_hook(hook_config)
    return hook.fetch(url, timeout=timeout, retry_attempts=retry_attempts)


def _download_stac_items_batch(
    urls: list[str],
    timeout: int = 30,
    retry_attempts: int = 3,
    hook_config: str = "default",
) -> list[dict[str, Any] | None]:
    """Download/generate multiple STAC items using configured hook.

    This function leverages the hook's batch processing capability if available,
    falling back to sequential processing otherwise.

    Args:
        urls: List of URLs to process.
        timeout: Request timeout in seconds.
        retry_attempts: Number of retry attempts.
        hook_config: STAC hook configuration string.

    Returns:
        List of STAC items (or None for failures), same order as input.
    """
    hook = _get_stac_hook(hook_config)
    return hook.fetch_batch(urls, timeout=timeout, retry_attempts=retry_attempts)


def _get_partition_key(
    item: dict[str, Any],
    grid_resolver: Any,
    temporal_bin: str,
    enable_global: bool,
    global_threshold: int,
) -> str:
    """Compute the partition key for a STAC item.

    Args:
        item: STAC item dictionary.
        grid_resolver: Grid resolver instance.
        temporal_bin: Temporal binning level ("year", "month", "day").
        enable_global: Whether to enable global partitioning.
        global_threshold: Threshold for global partition routing.

    Returns:
        Partition key string like "h3_82/2024/01" or "global/2024/01".
    """
    # Get geometry for grid cell resolution
    geometry = item.get("geometry")

    # Get grid cells for this item
    if geometry:
        tiles, is_spanning = grid_resolver.tiles_for_geometry_with_spanning_detection(geometry)
        if not tiles:
            # Fallback to bbox center
            bbox = item.get("bbox", [0, 0, 0, 0])
            center_geom = {"type": "Point", "coordinates": [(bbox[0] + bbox[2]) / 2, (bbox[1] + bbox[3]) / 2]}
            tiles = grid_resolver.tiles_for_geometry(center_geom)
        grid_cell = tiles[0] if tiles else "unknown"

        # Check for global partition routing
        if enable_global and len(tiles) > global_threshold:
            grid_cell = "global"
    else:
        # Use bbox center if no geometry
        bbox = item.get("bbox", [0, 0, 0, 0])
        center_geom = {"type": "Point", "coordinates": [(bbox[0] + bbox[2]) / 2, (bbox[1] + bbox[3]) / 2]}
        tiles = grid_resolver.tiles_for_geometry(center_geom)
        grid_cell = tiles[0] if tiles else "unknown"

    # Get temporal component
    datetime_str = item.get("properties", {}).get("datetime", "")
    if datetime_str:
        try:
            from datetime import datetime

            dt = datetime.fromisoformat(datetime_str.replace("Z", "+00:00"))
            if temporal_bin == "year":
                temporal = str(dt.year)
            elif temporal_bin == "month":
                temporal = f"{dt.year}/{dt.month:02d}"
            else:  # day
                temporal = f"{dt.year}/{dt.month:02d}/{dt.day:02d}"
        except (ValueError, AttributeError):
            temporal = "unknown"
    else:
        temporal = "unknown"

    return f"{grid_cell}/{temporal}"


def _write_shard_to_scratch(
    items: list[dict[str, Any]],
    partition_key: str,
    scratch_path: str,
    worker_tag: str,
    shard_id: int,
    engine_type: str = "rustac",
) -> dict[str, Any]:
    """Write items to a shard file in scratch storage.

    Args:
        items: STAC items to write.
        partition_key: Partition key for directory organization.
        scratch_path: Base scratch directory path.
        worker_tag: Unique worker identifier.
        shard_id: Shard number for this worker.
        engine_type: STAC engine to use.

    Returns:
        Shard info dictionary with path and metadata.
    """
    storage = get_storage_backend(scratch_path)

    # Create shard path
    safe_partition = partition_key.replace("/", "_")
    shard_filename = f"{safe_partition}_{worker_tag}_{shard_id}.parquet"
    shard_path = f"{scratch_path}/shards/{shard_filename}"

    # Ensure directory exists
    storage.makedirs(Path(shard_path).parent)

    # Convert items to GeoDataFrame using engine
    engine = get_engine(engine_type)  # type: ignore
    gdf = engine.items_to_geodataframe(items)

    # Write to scratch
    with storage.open(shard_path, "wb") as f:
        gdf.to_parquet(f, index=False, compression="snappy")

    return {
        "shard_path": shard_path,
        "partition_key": partition_key,
        "item_count": len(items),
        "worker_tag": worker_tag,
        "shard_id": shard_id,
    }


# =============================================================================
# Main Worker Functions (Module-Level, Serializable)
# =============================================================================


def process_url_batch(
    urls: list[str],
    worker_id: int,
    config_dict: dict[str, Any],
    job_id: str,
    batch_idx: int = 0,
) -> dict[str, Any]:
    """Process a batch of URLs, downloading STAC items and writing shards.

    This is a **module-level function** that can be pickled and sent to remote
    Dask workers. All dependencies are passed as parameters.

    Args:
        urls: List of STAC item URLs to download.
        worker_id: Unique worker identifier.
        config_dict: Serialized ProcessingConfig dictionary.
        job_id: Job ID for tracking.
        batch_idx: Batch index for this worker.

    Returns:
        Dictionary with:
        - shards: List of shard info dictionaries
        - stats: Worker statistics as dict
        - failed_urls: List of URLs that failed
    """
    # Reconstruct config from dict
    from .ingestion_pipeline import ProcessingConfig

    config = ProcessingConfig.from_dict(config_dict)

    # Create worker-local stats
    worker_stats = IngestionStatistics()
    failed_urls: list[str] = []

    # Create grid resolver
    grid_resolver = get_grid_system(
        config.grid_system,
        resolution=config.grid_resolution,
        geojson_path=config.geojson_path if config.grid_system == "geojson" else None,
    )

    # Unique tag for this worker's shards
    worker_tag = f"w{worker_id}-{job_id[:8]}-{batch_idx}"

    # Group items by partition as we process
    partition_items: dict[str, list[dict[str, Any]]] = {}
    shards: list[dict[str, Any]] = []

    # Choose processing method
    use_async = config.enable_concurrent_http and HAS_ASYNC_HTTP and len(urls) >= config.batch_size

    if use_async:
        # Async HTTP processing
        items = _download_batch_async(urls, config)
        for item in items:
            if item:
                worker_stats.record_url_processed(success=True)
                partition_key = _get_partition_key(
                    item,
                    grid_resolver,
                    config.temporal_bin,
                    config.enable_global_partitioning,
                    config.global_partition_threshold,
                )
                if partition_key not in partition_items:
                    partition_items[partition_key] = []
                partition_items[partition_key].append(item)
    else:
        # Synchronous processing with progress bar
        # Get hook config from ProcessingConfig
        hook_config = config_dict.get("stac_hook", "default")
        with tqdm(total=len(urls), desc=f"Worker {worker_id}", unit="urls", leave=False) as pbar:
            for url in urls:
                fetched_item = _download_stac_item(
                    url,
                    timeout=config.request_timeout,
                    retry_attempts=config.retry_attempts,
                    hook_config=hook_config,
                )
                if fetched_item:
                    worker_stats.record_url_processed(success=True)
                    partition_key = _get_partition_key(
                        fetched_item,
                        grid_resolver,
                        config.temporal_bin,
                        config.enable_global_partitioning,
                        config.global_partition_threshold,
                    )
                    if partition_key not in partition_items:
                        partition_items[partition_key] = []
                    partition_items[partition_key].append(fetched_item)
                else:
                    worker_stats.record_url_processed(success=False)
                    failed_urls.append(url)
                pbar.update(1)

    # Write shards for each partition
    shard_counter = 0
    for partition_key, items in partition_items.items():
        # Split into shard-sized chunks
        for i in range(0, len(items), config.items_per_shard):
            chunk = items[i : i + config.items_per_shard]
            shard_info = _write_shard_to_scratch(
                items=chunk,
                partition_key=partition_key,
                scratch_path=config.scratch_location,
                worker_tag=worker_tag,
                shard_id=shard_counter,
                engine_type=config.stac_engine,
            )
            shards.append(shard_info)
            shard_counter += 1

    logger.info(
        f"Worker {worker_id}: processed {len(urls)} URLs, wrote {len(shards)} shards, {len(failed_urls)} failures"
    )

    # Serialize stats as simple dict with basic counters
    stats_dict = {
        "urls_processed": worker_stats.urls_processed,
        "urls_failed": worker_stats.urls_failed,
    }

    return {
        "shards": shards,
        "stats": stats_dict,
        "failed_urls": failed_urls,
    }


def _download_batch_async(urls: list[str], config: ProcessingConfig) -> list[dict[str, Any]]:
    """Download URLs using async HTTP."""
    try:
        # Try to get or create event loop
        try:
            loop = asyncio.get_event_loop()
            if loop.is_running():
                # Running in async context (Jupyter, etc.)
                with ThreadPoolExecutor(max_workers=1) as _:
                    try:
                        import nest_asyncio

                        nest_asyncio.apply()
                    except ImportError:
                        pass  # nest_asyncio not available
                    return loop.run_until_complete(
                        download_stac_items_async(
                            urls,
                            concurrent_requests=config.concurrent_requests,
                            connection_pool_size=config.connection_pool_size,
                            request_timeout=config.request_timeout,
                            retry_attempts=config.retry_attempts,
                            retry_delay=config.retry_delay,
                            batch_size=config.batch_size,
                        )
                    )
        except RuntimeError:
            pass

        # Create new event loop
        loop = asyncio.new_event_loop()
        asyncio.set_event_loop(loop)
        try:
            return loop.run_until_complete(
                download_stac_items_async(
                    urls,
                    concurrent_requests=config.concurrent_requests,
                    connection_pool_size=config.connection_pool_size,
                    request_timeout=config.request_timeout,
                    retry_attempts=config.retry_attempts,
                    retry_delay=config.retry_delay,
                    batch_size=config.batch_size,
                )
            )
        finally:
            loop.close()
    except (ValueError, TypeError, RuntimeError, ConnectionError) as e:
        logger.error(f"Async download failed: {e}, falling back to sync")
        # Fallback to sync
        items = []
        for url in urls:
            item = _download_stac_item(url, config.request_timeout, config.retry_attempts)
            if item:
                items.append(item)
        return items


def consolidate_partition(
    partition_key: str,
    shard_paths: list[str],
    config_dict: dict[str, Any],
) -> dict[str, Any]:
    """Consolidate shards for a partition, merging with existing data.

    This is a **module-level function** that can be pickled and sent to remote
    Dask workers. All dependencies are passed as parameters.

    Deduplication uses newer datetime wins strategy: when items have the same ID,
    the one with the more recent datetime property is kept.

    Args:
        partition_key: Partition key (e.g., "h3_82/2024/01").
        shard_paths: List of shard file paths to merge.
        config_dict: Serialized ProcessingConfig dictionary.

    Returns:
        Dictionary with consolidation results:
        - partition_key: The partition that was consolidated
        - item_count: Total items in final partition
        - existing_count: Items that existed before
        - new_count: New items added
        - duplicates_removed: Number of duplicates removed
        - final_path: Path to final partition file
    """
    from .ingestion_pipeline import ProcessingConfig

    config = ProcessingConfig.from_dict(config_dict)

    # Get storage backends
    catalog_storage = get_storage_backend(config.output_catalog)
    scratch_storage = get_storage_backend(config.scratch_location)

    # Compute final path
    final_path = f"{config.output_catalog}/{partition_key}/data.parquet"

    existing_count = 0
    all_gdfs: list[gpd.GeoDataFrame] = []

    # Use temp directory for staging
    with tempfile.TemporaryDirectory(dir=config.temp_dir_location) as temp_dir:
        temp_existing_path = Path(temp_dir) / "existing.parquet"
        temp_merged_path = Path(temp_dir) / "merged.parquet"

        # Step 1: Download existing partition if it exists
        if catalog_storage.exists(final_path):
            try:
                logger.debug(f"Partition {partition_key}: downloading existing data")
                with catalog_storage.open(final_path, "rb") as f:
                    binary_data = f.read()
                    # Write to temp file so we can use gpd.read_parquet
                    with open(temp_existing_path, "wb") as temp_f:
                        temp_f.write(binary_data)
                    existing_gdf = gpd.read_parquet(temp_existing_path)
                    existing_count = len(existing_gdf)
                    all_gdfs.append(existing_gdf)
            except (OSError, ValueError, TypeError) as e:
                logger.error(f"Error reading existing partition {final_path}: {e}")

        # Step 2: Read all new shards
        for idx, shard_path in enumerate(shard_paths):
            try:
                with scratch_storage.open(shard_path, "rb") as f:
                    binary_data = f.read()
                    # Write to temp file so we can use gpd.read_parquet
                    temp_shard_path = Path(temp_dir) / f"shard_{idx}.parquet"
                    with open(temp_shard_path, "wb") as temp_f:
                        temp_f.write(binary_data)
                    gdf = gpd.read_parquet(temp_shard_path)
                    all_gdfs.append(gdf)
            except (OSError, ValueError, TypeError) as e:
                logger.error(f"Error reading shard {shard_path}: {e}")
                continue

        if not all_gdfs:
            return {
                "partition_key": partition_key,
                "item_count": 0,
                "existing_count": 0,
                "new_count": 0,
                "duplicates_removed": 0,
                "final_path": final_path,
            }

        # Step 3: Merge all data
        merged = pd.concat(all_gdfs, ignore_index=True)
        original_count = len(merged)

        # Step 4: Deduplicate - newer datetime wins
        if "datetime" in merged.columns and "id" in merged.columns:
            # Sort by datetime descending so newer items come first
            try:
                merged["_dt_sort"] = pd.to_datetime(merged["datetime"], errors="coerce")
                merged = merged.sort_values("_dt_sort", ascending=False, na_position="last")
                merged = merged.drop_duplicates(subset=["id"], keep="first")
                merged = merged.drop(columns=["_dt_sort"])
            except (ValueError, TypeError):
                # Fallback to simple dedup
                merged = merged.drop_duplicates(subset=["id"], keep="last")
        elif "id" in merged.columns:
            merged = merged.drop_duplicates(subset=["id"], keep="last")

        duplicates_removed = original_count - len(merged)

        # Step 5: Sort by configured key
        if config.sort_key in merged.columns:
            merged = merged.sort_values(config.sort_key, ascending=config.sort_ascending)

        # Step 6: Write to temp file
        # Ensure geometry column is properly set
        if "geometry" in merged.columns:
            merged_gdf = gpd.GeoDataFrame(merged, geometry="geometry")
        else:
            merged_gdf = gpd.GeoDataFrame(merged)
        merged_gdf.to_parquet(temp_merged_path, index=False, compression="snappy")

        # Step 7: Upload to final location
        catalog_storage.makedirs(Path(final_path).parent)
        if final_path.startswith("s3://"):
            catalog_storage.upload(str(temp_merged_path), final_path)
        else:
            # Local storage - atomic move
            import shutil

            shutil.copy2(temp_merged_path, final_path)

        new_count = len(merged) - existing_count

    logger.info(
        f"Partition {partition_key}: {existing_count} existing + {new_count} new = "
        f"{len(merged)} total ({duplicates_removed} duplicates removed)"
    )

    return {
        "partition_key": partition_key,
        "item_count": len(merged),
        "existing_count": existing_count,
        "new_count": new_count,
        "duplicates_removed": duplicates_removed,
        "final_path": final_path,
    }


__all__ = [
    "DownloadResult",
    "ConsolidationResult",
    "process_url_batch",
    "consolidate_partition",
]
