# stac_hooks.py
"""STAC fetch hooks for custom item generation.

This module provides a flexible hook system for generating STAC items from URLs
when the URL doesn't point directly to a STAC JSON. This is useful when:

1. URLs point to raw data files (COG, NetCDF, HDF5) that need metadata extraction
2. URLs are identifiers that need to be transformed via an external API
3. Custom STAC generation logic is required for specific data providers

Architecture:
    Hooks are designed to be serializable for Dask distributed processing.
    Instead of passing callable objects (which can't be pickled), hooks are
    configured via module paths that can be imported on remote workers.

Hook Types:
    - default: Standard STAC JSON fetch from URL (current behavior)
    - callable: Python callable (local processing only, not serializable)
    - module: Import path to Python function (serializable for Dask)
    - script: External executable (serializable, runs as subprocess)

Usage:
    >>> # Module-based hook (serializable, recommended for Dask)
    >>> config = ProcessingConfig(
    ...     input_file='urls.parquet',
    ...     output_catalog='./catalog',
    ...     scratch_location='./scratch',
    ...     stac_hook='module:mypackage.hooks:url_to_stac_item',
    ... )

    >>> # Script-based hook (serializable)
    >>> config = ProcessingConfig(
    ...     stac_hook='script:/path/to/generate_stac.py',
    ... )

    >>> # Callable hook (local processing only)
    >>> def my_hook(url: str) -> dict:
    ...     return {...}
    >>> config = ProcessingConfig(
    ...     stac_hook=my_hook,  # Only works with LocalProcessor
    ... )

Hook Function Signature:
    All hook functions must implement the following signature:

    def hook_function(url: str, **kwargs) -> dict[str, Any] | pystac.Item | None:
        '''Generate a STAC item from a URL.

        Args:
            url: The input URL/identifier to process.
            **kwargs: Additional context (timeout, retry_attempts, etc.)

        Returns:
            A valid STAC item dictionary, pystac.Item object, or None if generation failed.
            pystac.Item objects are automatically converted to dictionaries.
        '''
        pass

Batch Hook Signature (Optional):
    For efficiency, hooks can optionally support batch processing:

    def batch_hook_function(urls: list[str], **kwargs) -> list[dict[str, Any] | pystac.Item | None]:
        '''Generate STAC items from multiple URLs.

        Args:
            urls: List of URLs/identifiers to process.
            **kwargs: Additional context.

        Returns:
            List of STAC items (dict or pystac.Item) or None for failures, same order as input.
            pystac.Item objects are automatically converted to dictionaries.
        '''
        pass
"""

from __future__ import annotations

import importlib
import json
import logging
import subprocess  # nosec B404 - subprocess used for external tool integration
import tempfile
from abc import ABC, abstractmethod
from collections.abc import Callable
from typing import Any, Protocol, runtime_checkable

import requests

logger = logging.getLogger(__name__)


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


def _normalize_stac_result(result: Any) -> dict[str, Any] | None:
    """Convert hook result to a dictionary, handling pystac.Item objects.

    Args:
        result: The result from a hook function. Can be:
            - dict: Returned as-is
            - pystac.Item: Converted to dict via to_dict()
            - None: Returned as-is
            - Other: Returns None

    Returns:
        STAC item dictionary or None.
    """
    if result is None:
        return None
    if isinstance(result, dict):
        return result
    # Check for pystac.Item or any object with to_dict method
    if hasattr(result, "to_dict") and callable(result.to_dict):
        try:
            converted = result.to_dict()
            if isinstance(converted, dict):
                return converted
            return None
        except (ValueError, TypeError, AttributeError, RuntimeError) as e:
            logger.warning(f"Failed to convert pystac.Item to dict: {e}")
            return None
    return None


def _normalize_stac_results(results: list[Any]) -> list[dict[str, Any] | None]:
    """Convert a list of hook results to dictionaries.

    Args:
        results: List of results from a batch hook function.

    Returns:
        List of STAC item dictionaries (or None for failures).
    """
    return [_normalize_stac_result(r) for r in results]


# =============================================================================
# Protocol Definitions
# =============================================================================


@runtime_checkable
class STACHookProtocol(Protocol):
    """Protocol for STAC fetch hooks."""

    def fetch(self, url: str, **kwargs) -> dict[str, Any] | None:
        """Fetch/generate a STAC item from a URL.

        Args:
            url: The URL or identifier to process.
            **kwargs: Additional context like timeout, retry_attempts.

        Returns:
            STAC item dictionary or None if fetch failed.
        """
        ...


@runtime_checkable
class BatchSTACHookProtocol(Protocol):
    """Protocol for batch STAC fetch hooks."""

    def fetch_batch(self, urls: list[str], **kwargs) -> list[dict[str, Any] | None]:
        """Fetch/generate STAC items from multiple URLs.

        Args:
            urls: List of URLs or identifiers to process.
            **kwargs: Additional context.

        Returns:
            List of STAC items (or None for failures), in same order as input.
        """
        ...


# =============================================================================
# Hook Implementations
# =============================================================================


class BaseSTACHook(ABC):
    """Base class for STAC fetch hooks."""

    @abstractmethod
    def fetch(self, url: str, **kwargs) -> dict[str, Any] | None:
        """Fetch a single STAC item."""
        pass

    def fetch_batch(self, urls: list[str], **kwargs) -> list[dict[str, Any] | None]:
        """Fetch multiple STAC items. Default: sequential fetch."""
        return [self.fetch(url, **kwargs) for url in urls]

    def to_config(self) -> str:
        """Serialize hook to config string for storage/transmission."""
        raise NotImplementedError("Subclass must implement to_config()")


class DefaultSTACHook(BaseSTACHook):
    """Default hook: fetch STAC JSON directly from URL.

    This is the standard behavior - assumes the URL points to a valid STAC item JSON.
    """

    def fetch(
        self,
        url: str,
        timeout: int = 30,
        retry_attempts: int = 3,
        **kwargs,
    ) -> dict[str, Any] | None:
        """Download and parse STAC item from URL.

        Args:
            url: URL to STAC item JSON.
            timeout: Request timeout in seconds.
            retry_attempts: Number of retry attempts.

        Returns:
            Parsed STAC item or None if failed.
        """
        import time

        for attempt in range(retry_attempts):
            try:
                response = requests.get(url, timeout=timeout)
                response.raise_for_status()
                return response.json()
            except requests.RequestException as e:
                if attempt == retry_attempts - 1:
                    logger.warning(f"Failed to download {url} after {retry_attempts} attempts: {e}")
                    return None
                time.sleep(2**attempt)  # Exponential backoff
        return None

    def to_config(self) -> str:
        """Return config string for default hook."""
        return "default"


class ModuleSTACHook(BaseSTACHook):
    """Hook that calls a Python function specified by module path.

    This is the recommended approach for Dask distributed processing because
    the module path is a string that can be serialized and the function is
    imported on each worker.

    The module path format is: 'package.module:function_name'

    Example:
        >>> hook = ModuleSTACHook('mypackage.stac_generator:url_to_item')
        >>> item = hook.fetch('https://example.com/data.tif')
    """

    def __init__(self, module_path: str):
        """Initialize with module path.

        Args:
            module_path: Import path in format 'package.module:function_name'
        """
        self.module_path = module_path
        self._func: Callable | None = None
        self._batch_func: Callable | None = None

    def _load_function(self) -> Callable[..., Any]:
        """Import and return the hook function."""
        if self._func is not None:
            return self._func

        try:
            module_name, func_name = self.module_path.rsplit(":", 1)
            module = importlib.import_module(module_name)
            func = getattr(module, func_name)
            self._func = func

            # Check for batch function (convention: func_name + '_batch')
            batch_func_name = f"{func_name}_batch"
            if hasattr(module, batch_func_name):
                self._batch_func = getattr(module, batch_func_name)

            return func
        except (ValueError, ImportError, AttributeError) as e:
            raise ImportError(f"Failed to import hook function from '{self.module_path}': {e}") from e

    def fetch(self, url: str, **kwargs) -> dict[str, Any] | None:
        """Fetch STAC item using the imported function.

        Args:
            url: URL or identifier to process.
            **kwargs: Additional context passed to the function.

        Returns:
            STAC item or None. Supports both dict and pystac.Item returns.
        """
        func = self._load_function()
        try:
            result = func(url, **kwargs)
            return _normalize_stac_result(result)
        except (ValueError, TypeError, AttributeError, RuntimeError, OSError) as e:
            logger.warning(f"Hook function failed for {url}: {e}")
            return None

    def fetch_batch(self, urls: list[str], **kwargs) -> list[dict[str, Any] | None]:
        """Fetch multiple STAC items, using batch function if available.

        Args:
            urls: List of URLs to process.
            **kwargs: Additional context.

        Returns:
            List of STAC items (or None for failures). Supports pystac.Item returns.
        """
        self._load_function()  # Ensure functions are loaded

        if self._batch_func is not None:
            try:
                results = self._batch_func(urls, **kwargs)
                if isinstance(results, list) and len(results) == len(urls):
                    return _normalize_stac_results(results)
                logger.warning("Batch hook returned invalid result, falling back to sequential")
            except (ValueError, TypeError, AttributeError, RuntimeError, OSError) as e:
                logger.warning(f"Batch hook failed: {e}, falling back to sequential")

        # Fall back to sequential
        return super().fetch_batch(urls, **kwargs)

    def to_config(self) -> str:
        """Return config string for module hook."""
        return f"module:{self.module_path}"


class PassthroughSTACHook(BaseSTACHook):
    """Hook that treats input URLs as pre-fetched STAC item JSON.

    This hook is designed for processing NDJSON files that already contain
    complete STAC items, eliminating the need for HTTP fetching. The "URL"
    in the input file is actually the full STAC item JSON as a string.

    Use Cases:
        - ITS_LIVE bulk data: NDJSON files with pre-aggregated STAC items
        - Local STAC collections: STAC items already downloaded
        - Performance optimization: Skip HTTP fetch for cached data

    Example:
        >>> hook = PassthroughSTACHook()
        >>> # url is actually a JSON string
        >>> item = hook.fetch('{"type": "Feature", "id": "abc", ...}')
        >>> print(item['id'])
        'abc'

    Performance:
        - Eliminates HTTP overhead (no network calls)
        - Faster processing for local/cached STAC items
        - Reduces load on STAC catalog servers
    """

    def fetch(self, url: str, **kwargs) -> dict[str, Any] | None:
        """Parse input URL as STAC item JSON directly.

        Args:
            url: JSON string containing a complete STAC item.
            **kwargs: Additional context (timeout, retry_attempts) - ignored for passthrough.

        Returns:
            Parsed STAC item dictionary or None if JSON is invalid.
        """
        try:
            # Input "url" is actually a JSON string
            item = json.loads(url)

            # Validate it's a dictionary with STAC item structure
            if not isinstance(item, dict):
                logger.warning(f"Passthrough hook input is not a dict: {type(url)}")
                return None

            # Basic STAC item validation - should have at minimum 'type': 'Feature'
            if item.get("type") != "Feature":
                logger.warning(f"Passthrough hook input missing 'type': 'Feature': {url[:100]}...")
                return None

            # Should have geometry and properties
            if "geometry" not in item:
                logger.warning(f"Passthrough hook input missing 'geometry': {url[:100]}...")
                return None

            if "properties" not in item:
                logger.warning(f"Passthrough hook input missing 'properties': {url[:100]}...")
                return None

            return item

        except json.JSONDecodeError as e:
            logger.warning(f"Passthrough hook failed to parse JSON: {e}")
            return None
        except (ValueError, TypeError, OSError) as e:
            logger.warning(f"Passthrough hook error: {e}")
            return None

    def fetch_batch(self, urls: list[str], **kwargs) -> list[dict[str, Any] | None]:
        """Parse multiple JSON strings as STAC items.

        Args:
            urls: List of JSON strings containing STAC items.

        Returns:
            List of STAC item dictionaries (or None for failures), in same order as input.
        """
        return [self.fetch(url, **kwargs) for url in urls]

    def to_config(self) -> str:
        """Return config string for passthrough hook."""
        return "passthrough"


class ScriptSTACHook(BaseSTACHook):
    """Hook that calls an external script/executable.

    The script receives the URL as an argument and should output valid STAC JSON
    to stdout. For batch processing, the script can accept multiple URLs.

    Script invocation:
        Single: script_path URL
        Batch:  script_path --batch < urls.txt  (one URL per line)

    The script should exit with code 0 on success, non-zero on failure.
    For batch mode, output should be NDJSON (one JSON object per line).

    Example:
        >>> hook = ScriptSTACHook('/path/to/generate_stac.py')
        >>> item = hook.fetch('https://example.com/data.tif')
    """

    def __init__(self, script_path: str, interpreter: str | None = None):
        """Initialize with script path.

        Args:
            script_path: Path to the script/executable.
            interpreter: Optional interpreter (e.g., 'python', 'python3').
                        If None, script is executed directly.
        """
        self.script_path = script_path
        self.interpreter = interpreter

    def fetch(
        self,
        url: str,
        timeout: int = 60,
        **kwargs,
    ) -> dict[str, Any] | None:
        """Run script to generate STAC item.

        Args:
            url: URL to pass to the script.
            timeout: Script execution timeout in seconds.

        Returns:
            Parsed STAC item from script stdout, or None if failed.
        """
        cmd = self._build_command([url])

        try:
            result = subprocess.run(  # nosec B603 - command built from validated config
                cmd,
                capture_output=True,
                text=True,
                timeout=timeout,
                check=True,
            )
            return json.loads(result.stdout.strip())
        except subprocess.TimeoutExpired:
            logger.warning(f"Script timed out for {url}")
            return None
        except subprocess.CalledProcessError as e:
            logger.warning(f"Script failed for {url}: {e.stderr}")
            return None
        except json.JSONDecodeError as e:
            logger.warning(f"Script output is not valid JSON for {url}: {e}")
            return None

    def fetch_batch(
        self,
        urls: list[str],
        timeout: int = 300,
        **kwargs,
    ) -> list[dict[str, Any] | None]:
        """Run script in batch mode for multiple URLs.

        Args:
            urls: List of URLs to process.
            timeout: Script execution timeout in seconds.

        Returns:
            List of STAC items (or None for failures).
        """
        if len(urls) == 0:
            return []

        if len(urls) == 1:
            return [self.fetch(urls[0], timeout=timeout, **kwargs)]

        # Try batch mode first
        cmd = self._build_command(["--batch"])

        try:
            # Write URLs to temp file for stdin
            with tempfile.NamedTemporaryFile(mode="w", suffix=".txt", delete=False) as f:
                f.write("\n".join(urls))
                temp_path = f.name

            with open(temp_path) as stdin_file:
                result = subprocess.run(  # nosec B603 - command built from validated config
                    cmd,
                    stdin=stdin_file,
                    capture_output=True,
                    text=True,
                    timeout=timeout,
                )

            # Clean up temp file
            import os

            os.unlink(temp_path)

            if result.returncode == 0:
                # Parse NDJSON output
                items: list[dict[str, Any] | None] = []
                for line in result.stdout.strip().split("\n"):
                    if line.strip():
                        try:
                            items.append(json.loads(line))
                        except json.JSONDecodeError:
                            items.append(None)
                    else:
                        items.append(None)

                if len(items) == len(urls):
                    return items

                logger.warning("Batch script returned wrong number of items, falling back to sequential")

        except subprocess.TimeoutExpired:
            logger.warning("Batch script timed out, falling back to sequential")
        except (subprocess.CalledProcessError, OSError, ValueError, FileNotFoundError) as e:
            logger.warning(f"Batch script failed: {e}, falling back to sequential")

        # Fall back to sequential processing
        return super().fetch_batch(urls, **kwargs)

    def _build_command(self, args: list[str]) -> list[str]:
        """Build command list for subprocess."""
        if self.interpreter:
            return [self.interpreter, self.script_path, *args]
        return [self.script_path, *args]

    def to_config(self) -> str:
        """Return config string for script hook."""
        if self.interpreter:
            return f"script:{self.interpreter}:{self.script_path}"
        return f"script:{self.script_path}"


class CallableSTACHook(BaseSTACHook):
    """Hook that wraps a Python callable.

    WARNING: This hook is NOT serializable and will not work with Dask
    distributed processing. Use ModuleSTACHook for distributed workloads.

    This hook is useful for local processing or testing where you want to
    pass a function directly without creating a module.

    Example:
        >>> def my_hook(url, **kwargs):
        ...     return {'type': 'Feature', 'id': url, ...}
        >>> hook = CallableSTACHook(my_hook)
        >>> item = hook.fetch('https://example.com/data.tif')
    """

    def __init__(
        self,
        func: Callable[[str], dict[str, Any] | None],
        batch_func: Callable[[list[str]], list[dict[str, Any] | None]] | None = None,
    ):
        """Initialize with callable.

        Args:
            func: Function that takes URL and returns STAC item or None.
            batch_func: Optional batch function for efficient batch processing.
        """
        self.func = func
        self.batch_func = batch_func

    def fetch(self, url: str, **kwargs) -> dict[str, Any] | None:
        """Invoke the callable with the URL.

        Args:
            url: URL to process.
            **kwargs: Additional context.

        Returns:
            STAC item or None. Supports both dict and pystac.Item returns.
        """
        try:
            result = self.func(url, **kwargs)
            return _normalize_stac_result(result)
        except (ValueError, TypeError, AttributeError, RuntimeError, OSError) as e:
            logger.warning(f"Callable hook failed for {url}: {e}")
            return None

    def fetch_batch(self, urls: list[str], **kwargs) -> list[dict[str, Any] | None]:
        """Invoke batch callable if available, otherwise sequential.

        Args:
            urls: List of URLs to process.
            **kwargs: Additional context.

        Returns:
            List of STAC items (or None). Supports pystac.Item returns.
        """
        if self.batch_func is not None:
            try:
                results = self.batch_func(urls, **kwargs)
                if isinstance(results, list) and len(results) == len(urls):
                    return _normalize_stac_results(results)
            except (ValueError, TypeError, AttributeError, RuntimeError, OSError) as e:
                logger.warning(f"Batch callable failed: {e}, falling back to sequential")

        return super().fetch_batch(urls, **kwargs)

    def to_config(self) -> str:
        """Return config string - not meaningful for callables."""
        func_name = getattr(self.func, "__name__", "anonymous")
        module = getattr(self.func, "__module__", "unknown")
        return f"callable:{module}.{func_name}"


# =============================================================================
# Hook Factory
# =============================================================================


def parse_hook_config(config: str | Callable | BaseSTACHook | None) -> BaseSTACHook:
    """Parse hook configuration and return appropriate hook instance.

    Args:
        config: Hook configuration, can be:
            - None or "default": Use DefaultSTACHook
            - "passthrough": Use PassthroughSTACHook (URLs are pre-fetched STAC JSON)
            - "module:path.to.module:function": Use ModuleSTACHook
            - "script:/path/to/script": Use ScriptSTACHook
            - "script:python:/path/to/script.py": Use ScriptSTACHook with interpreter
            - Callable: Use CallableSTACHook (local only)
            - BaseSTACHook instance: Use as-is

    Returns:
        Configured hook instance.

    Raises:
        ValueError: If config format is invalid.
    """
    if config is None or config == "default":
        return DefaultSTACHook()

    if isinstance(config, BaseSTACHook):
        return config

    if callable(config):
        return CallableSTACHook(config)

    if not isinstance(config, str):
        raise ValueError(f"Invalid hook config type: {type(config)}")

    # Parse string config
    if config == "passthrough":
        return PassthroughSTACHook()

    if config.startswith("module:"):
        module_path = config[7:]  # Remove 'module:' prefix
        return ModuleSTACHook(module_path)

    if config.startswith("script:"):
        script_config = config[7:]  # Remove 'script:' prefix
        # Check for interpreter:path format
        if ":" in script_config and not script_config.startswith("/"):
            interpreter, script_path = script_config.split(":", 1)
            return ScriptSTACHook(script_path, interpreter=interpreter)
        return ScriptSTACHook(script_config)

    raise ValueError(
        f"Invalid hook config: {config}. Expected 'default', 'passthrough', 'module:path:func', "
        "'script:/path', or a callable."
    )


def get_hook(config: str | Callable | BaseSTACHook | None = None) -> BaseSTACHook:
    """Get a STAC hook instance from configuration.

    This is the main entry point for getting hooks. It's an alias for
    parse_hook_config with a friendlier name.

    Args:
        config: Hook configuration (see parse_hook_config for details).

    Returns:
        Configured hook instance.
    """
    return parse_hook_config(config)


def serialize_hook(hook: BaseSTACHook | Callable | None) -> str:
    """Serialize a hook to a config string for transmission.

    Args:
        hook: Hook instance or callable.

    Returns:
        Config string that can be used with parse_hook_config.
    """
    if hook is None:
        return "default"

    if isinstance(hook, BaseSTACHook):
        return hook.to_config()

    if callable(hook):
        # Try to create a module path for the callable
        func_name = getattr(hook, "__name__", None)
        module = getattr(hook, "__module__", None)
        if func_name and module and module != "__main__":
            return f"module:{module}:{func_name}"
        # Can't serialize anonymous or __main__ functions
        raise ValueError(
            f"Cannot serialize callable {hook}. Use ModuleSTACHook with an importable "
            "function path for distributed processing."
        )

    raise ValueError(f"Cannot serialize hook: {hook}")


__all__ = [
    "STACHookProtocol",
    "BatchSTACHookProtocol",
    "BaseSTACHook",
    "DefaultSTACHook",
    "PassthroughSTACHook",
    "ModuleSTACHook",
    "ScriptSTACHook",
    "CallableSTACHook",
    "parse_hook_config",
    "get_hook",
    "serialize_hook",
    "_normalize_stac_result",
    "_normalize_stac_results",
]
