"""
Pipeline hooks for pre/post save callbacks and custom transformation logic.

This module provides a hook system for executing custom logic at various
points in the data ingestion pipeline, configurable via YAML.
"""

import importlib
import logging
from abc import ABC, abstractmethod
from dataclasses import dataclass, field
from enum import Enum
from typing import Any, Callable, Dict, List, Optional, Type
from datetime import datetime

from atomsql import Model
from atomsql.orm import fields

from atomingest.observability.metrics import record_dlq_entry

logger = logging.getLogger(__name__)


class HookType(Enum):
    """Types of hooks that can be executed in the pipeline."""

    PRE_NORMALIZATION = "pre_normalization"
    POST_NORMALIZATION = "post_normalization"
    PRE_SAVE = "pre_save"
    POST_SAVE = "post_save"
    ON_ERROR = "on_error"
    ON_COMPLETE = "on_complete"


class HookExecutionError(Exception):
    """Raised when a hook execution fails."""

    def __init__(
        self,
        message: str,
        hook_name: str = None,
        hook_type: HookType = None,
        original_error: Exception = None,
    ):
        self.message = message
        self.hook_name = hook_name
        self.hook_type = hook_type
        self.original_error = original_error
        super().__init__(self.format_message())

    def format_message(self) -> str:
        parts = ["Hook execution failed"]
        if self.hook_name:
            parts.append(f"[{self.hook_name}]")
        if self.hook_type:
            parts.append(f"({self.hook_type.value})")
        parts.append(f": {self.message}")
        return " ".join(parts)


class Hook(ABC):
    """
    Base class for all pipeline hooks.

    Hooks are executed at specific points in the pipeline and can
    transform data, perform side effects, or handle errors.
    """

    def __init__(self, name: str = None, config: Dict[str, Any] = None):
        """
        Initialize the hook.

        Args:
            name: Optional name for the hook (for logging/debugging).
            config: Optional configuration dictionary for the hook.
        """
        self.name = name or self.__class__.__name__
        self.config = config or {}

    @abstractmethod
    def execute(self, context: "HookContext") -> "HookContext":
        """
        Execute the hook logic.

        Args:
            context: The hook execution context containing data and metadata.

        Returns:
            The modified context (or same context if no modifications).
        """
        pass

    def __repr__(self) -> str:
        return f"<{self.__class__.__name__}: {self.name}>"


@dataclass
class HookContext:
    """
    Context object passed to hooks during execution.

    Contains the data being processed and metadata about the current
    pipeline state.
    """

    # The current row data (dict for single row operations)
    row_data: Optional[Dict[str, Any]] = None

    # The model instance (for post-save hooks)
    instance: Optional[Any] = None

    # All rows (for batch operations)
    rows: Optional[List[Dict[str, Any]]] = None

    # The target table name
    target_table: Optional[str] = None

    # The schema definition
    schema: Optional[Dict[str, Any]] = None

    # Additional metadata
    metadata: Dict[str, Any] = field(default_factory=dict)

    # Error information (for on_error hooks)
    error: Optional[Exception] = None

    # Statistics
    processed_count: int = 0
    failed_count: int = 0
    total_count: int = 0

    # Flag to skip further processing
    skip_row: bool = False

    # Flag to abort the entire pipeline
    abort: bool = False


class PreSaveHook(Hook):
    """
    Hook executed before saving a row to the database.

    Can be used to:
    - Transform data before insertion
    - Add computed fields
    - Validate business rules
    - Skip rows based on conditions
    """

    hook_type = HookType.PRE_SAVE

    def execute(self, context: HookContext) -> HookContext:
        """Override in subclass to implement pre-save logic."""
        return context


class PostSaveHook(Hook):
    """
    Hook executed after successfully saving a row to the database.

    Can be used to:
    - Trigger side effects (notifications, audit logs)
    - Update related records
    - Cache data
    """

    hook_type = HookType.POST_SAVE

    def execute(self, context: HookContext) -> HookContext:
        """Override in subclass to implement post-save logic."""
        return context


class TransformHook(PreSaveHook):
    """
    Specialized pre-save hook for data transformation.

    Provides a simpler interface for transforming row data.
    """

    def execute(self, context: HookContext) -> HookContext:
        if context.row_data:
            context.row_data = self.transform(context.row_data, context)
        return context

    def transform(
        self, row_data: Dict[str, Any], context: HookContext
    ) -> Dict[str, Any]:
        """
        Transform the row data.

        Override this method in subclasses.

        Args:
            row_data: The current row data dictionary.
            context: The full hook context.

        Returns:
            The transformed row data dictionary.
        """
        return row_data


class OnErrorHook(Hook):
    """
    Hook executed when an error occurs during processing.

    Can be used to:
    - Log errors
    - Send alerts
    - Route to dead letter queue
    - Attempt recovery
    """

    hook_type = HookType.ON_ERROR

    def execute(self, context: HookContext) -> HookContext:
        """Override in subclass to implement error handling logic."""
        return context


class OnCompleteHook(Hook):
    """
    Hook executed when the pipeline completes (success or failure).

    Can be used to:
    - Generate reports
    - Send completion notifications
    - Clean up resources
    """

    hook_type = HookType.ON_COMPLETE

    def execute(self, context: HookContext) -> HookContext:
        """Override in subclass to implement completion logic."""
        return context


class DLQFailureModel(Model):
    """Internal model representing the Dead Letter Queue table.
    Stores failed rows for audit and replay.
    """

    row_data = fields.JSONField(nullable=False)
    error_message = fields.TextField(nullable=False)
    target_table = fields.CharField(max_length=255, nullable=False)
    source = fields.CharField(max_length=255, nullable=True)
    timestamp = fields.DateTimeField()


class DeadLetterQueueHook(OnErrorHook):
    """
    Hook that routes failed rows to a Dead Letter Queue.

    Persists failures via two complementary mechanisms:

    1. **Database** – saves a :class:`DLQFailureModel` row when a ``db``
       connection is present in ``context.metadata``.
    2. **File store** – writes to the file-based
       :class:`~atomingest.pipeline.dlq.DLQManager` when a ``dlq_manager``
       instance is present in ``context.metadata``.

    Either or both mechanisms can be active simultaneously.  Prometheus
    ``atomingest_dlq_records_total`` is always incremented regardless of
    which persistence paths are used.
    """

    def execute(self, context: HookContext) -> HookContext:
        table = context.metadata.get("table", "unknown")
        source = context.metadata.get("source", "unknown")
        run_id = context.metadata.get("run_id", "unknown")
        error_type = type(context.error).__name__ if context.error else "unknown"

        # --- Prometheus metric (always) ---
        record_dlq_entry(table_name=table, error_type=error_type)

        # --- Database persistence ---
        db = context.metadata.get("db")
        if db:
            try:
                db.register(DLQFailureModel)
                failure = DLQFailureModel(
                    row_data=context.row_data,
                    error_message=str(context.error),
                    target_table=table,
                    source=source,
                    timestamp=datetime.now(),
                )
                failure.save()
                logger.warning(
                    f"Row routed to DLQ DB (Failure ID: {failure.id}): {context.error}"
                )
            except Exception as e:
                logger.critical(f"CRITICAL: Failed to save to DLQ DB: {e}")
        else:
            logger.warning(
                "DeadLetterQueueHook: no 'db' in context.metadata — skipping DB persistence"
            )

        # --- File-based DLQManager persistence (optional) ---
        dlq_manager = context.metadata.get("dlq_manager")
        if dlq_manager is not None:
            try:
                dlq_manager.record_failure(
                    row_data=context.row_data,
                    error=context.error,
                    target_table=table,
                    source=source,
                    run_id=run_id,
                )
            except Exception as e:
                logger.error(f"DeadLetterQueueHook: failed to write to DLQManager: {e}")

        return context


class FunctionHook(Hook):
    """
    Hook that wraps a callable function.

    Allows using simple functions as hooks without creating a class.
    """

    def __init__(
        self,
        func: Callable[[HookContext], HookContext],
        name: str = None,
        config: Dict[str, Any] = None,
    ):
        super().__init__(name or func.__name__, config)
        self.func = func

    def execute(self, context: HookContext) -> HookContext:
        return self.func(context)


@dataclass
class HookConfig:
    """
    Configuration for a hook parsed from YAML.

    Supports both class-based hooks and function references.
    """

    hook_type: HookType
    # Module path for class-based hook (e.g., "myapp.hooks.MyHook")
    module_path: Optional[str] = None
    # Function reference (e.g., "myapp.transforms.uppercase_names")
    function_path: Optional[str] = None
    # Hook name for identification
    name: Optional[str] = None
    # Hook-specific configuration
    config: Dict[str, Any] = field(default_factory=dict)
    # Whether hook failures should be fatal
    fail_on_error: bool = True
    # Execution order (lower = earlier)
    order: int = 0

    def __post_init__(self):
        if not self.module_path and not self.function_path:
            raise ValueError(
                "Either 'module_path' or 'function_path' must be specified"
            )


class HookRegistry:
    """
    Registry for managing and executing hooks.

    Hooks can be registered by type and are executed in order.
    """

    # Built-in hook classes that can be referenced by name
    BUILTIN_HOOKS: Dict[str, Type[Hook]] = {
        "PreSaveHook": PreSaveHook,
        "PostSaveHook": PostSaveHook,
        "TransformHook": TransformHook,
        "OnErrorHook": OnErrorHook,
        "OnCompleteHook": OnCompleteHook,
        "DeadLetterQueueHook": DeadLetterQueueHook,
    }

    @classmethod
    def register_builtin_hook(cls, name: str, hook_class: Type[Hook]) -> None:
        """Register a new built-in hook class that can be referenced by name.

        This allows external modules to register their hooks as built-ins.

        Args:
            name: The name to reference the hook by in YAML configs.
            hook_class: The Hook subclass to register.
        """
        if not (isinstance(hook_class, type) and issubclass(hook_class, Hook)):
            raise ValueError(f"'{hook_class}' is not a Hook subclass")
        cls.BUILTIN_HOOKS[name] = hook_class
        logger.debug(f"Registered built-in hook: {name}")

    def __init__(self):
        self._hooks: Dict[HookType, List[tuple[int, Hook]]] = {
            hook_type: [] for hook_type in HookType
        }

    def register(self, hook_type: HookType, hook: Hook, order: int = 0) -> None:
        """
        Register a hook for a specific type.

        Args:
            hook_type: The type of hook (when it should be executed).
            hook: The hook instance.
            order: Execution order (lower = earlier).
        """
        self._hooks[hook_type].append((order, hook))
        # Sort by order
        self._hooks[hook_type].sort(key=lambda x: x[0])
        logger.debug(f"Registered hook '{hook.name}' for {hook_type.value}")

    def register_from_config(self, hook_config: HookConfig) -> None:
        """
        Register a hook from a HookConfig object.

        Handles loading hooks from module paths or function references.

        Args:
            hook_config: The hook configuration.
        """
        hook = self._create_hook_from_config(hook_config)
        self.register(hook_config.hook_type, hook, hook_config.order)

    def _create_hook_from_config(self, hook_config: HookConfig) -> Hook:
        """Create a Hook instance from configuration."""
        if hook_config.function_path:
            # Load function and wrap in FunctionHook
            func = self._load_callable(hook_config.function_path)
            return FunctionHook(
                func=func, name=hook_config.name, config=hook_config.config
            )
        elif hook_config.module_path:
            # Load hook class and instantiate
            # Check if it's a built-in hook name
            if hook_config.module_path in self.BUILTIN_HOOKS:
                hook_class = self.BUILTIN_HOOKS[hook_config.module_path]
            else:
                hook_class = self._load_class(hook_config.module_path)

            return hook_class(name=hook_config.name, config=hook_config.config)
        else:
            raise ValueError("No module_path or function_path specified")

    def _load_callable(self, path: str) -> Callable[..., Any]:
        """
        Load a callable from a dotted path.

        Args:
            path: Dotted path like "myapp.module.function_name"

        Returns:
            The loaded callable.
        """
        module_path, func_name = path.rsplit(".", 1)
        try:
            module = importlib.import_module(module_path)
            func = getattr(module, func_name)
            if not callable(func):
                raise ValueError(f"'{path}' is not callable")
            return func  # type: ignore[no-any-return]
        except (ImportError, AttributeError) as e:
            raise ImportError(f"Could not load callable '{path}': {e}")

    def _load_class(self, path: str) -> Type[Hook]:
        """
        Load a Hook class from a dotted path.

        Args:
            path: Dotted path like "myapp.hooks.MyHookClass"

        Returns:
            The loaded Hook class.
        """
        module_path, class_name = path.rsplit(".", 1)
        try:
            module = importlib.import_module(module_path)
            cls = getattr(module, class_name)
            if not (isinstance(cls, type) and issubclass(cls, Hook)):
                raise ValueError(f"'{path}' is not a Hook subclass")
            return cls
        except (ImportError, AttributeError) as e:
            raise ImportError(f"Could not load Hook class '{path}': {e}")

    def get_hooks(self, hook_type: HookType) -> List[Hook]:
        """
        Get all hooks registered for a type.

        Args:
            hook_type: The hook type.

        Returns:
            List of hooks in execution order.
        """
        return [hook for _, hook in self._hooks[hook_type]]

    def execute_hooks(
        self, hook_type: HookType, context: HookContext, fail_fast: bool = True
    ) -> HookContext:
        """
        Execute all hooks of a given type.

        Args:
            hook_type: The type of hooks to execute.
            context: The hook context.
            fail_fast: Whether to stop on first error.

        Returns:
            The modified context after all hooks have executed.

        Raises:
            HookExecutionError: If a hook fails and fail_fast is True.
        """
        hooks = self.get_hooks(hook_type)
        if not hooks:
            return context

        logger.debug(f"Executing {len(hooks)} {hook_type.value} hooks")

        for hook in hooks:
            try:
                context = hook.execute(context)
                if context.abort:
                    logger.warning(f"Hook '{hook.name}' requested pipeline abort")
                    break
            except Exception as e:
                logger.error(f"Hook '{hook.name}' failed: {e}")
                if fail_fast:
                    raise HookExecutionError(
                        message=str(e),
                        hook_name=hook.name,
                        hook_type=hook_type,
                        original_error=e,
                    )
                # Continue with next hook if not fail_fast
                context.metadata.setdefault("hook_errors", []).append(
                    {"hook": hook.name, "error": str(e)}
                )

        return context

    def clear(self, hook_type: HookType = None) -> None:
        """
        Clear registered hooks.

        Args:
            hook_type: Specific type to clear, or None to clear all.
        """
        if hook_type:
            self._hooks[hook_type] = []
        else:
            for ht in HookType:
                self._hooks[ht] = []


def parse_hooks_config(hooks_config: Dict[str, Any]) -> List[HookConfig]:
    """
    Parse hooks configuration from YAML.

    YAML Example:
        hooks:
          pre_save:
            - function: myapp.transforms.normalize_names
              name: normalize_names
              order: 1
            - module: myapp.hooks.ValidationHook
              name: custom_validation
              config:
                strict: true
              order: 2
          post_save:
            - function: myapp.audit.log_save
          on_error:
            - module: myapp.hooks.ErrorHandler
              fail_on_error: false

    Args:
        hooks_config: Dictionary from YAML hooks section.

    Returns:
        List of HookConfig objects.
    """
    configs = []

    for hook_type_str, hook_list in hooks_config.items():
        try:
            hook_type = HookType(hook_type_str)
        except ValueError:
            logger.warning(f"Unknown hook type '{hook_type_str}', skipping")
            continue

        if not isinstance(hook_list, list):
            hook_list = [hook_list]

        for idx, hook_def in enumerate(hook_list):
            if isinstance(hook_def, str):
                # Simple string reference (function path)
                hook_def = {"function": hook_def}

            config = HookConfig(
                hook_type=hook_type,
                module_path=hook_def.get("module"),
                function_path=hook_def.get("function"),
                name=hook_def.get("name"),
                config=hook_def.get("config", {}),
                fail_on_error=hook_def.get("fail_on_error", True),
                order=hook_def.get("order", idx),
            )
            configs.append(config)

    return configs


def create_registry_from_config(hooks_config: Dict[str, Any]) -> HookRegistry:
    """
    Create a HookRegistry populated from YAML configuration.

    Args:
        hooks_config: Dictionary from YAML hooks section.

    Returns:
        Configured HookRegistry.
    """
    registry = HookRegistry()
    hook_configs = parse_hooks_config(hooks_config)

    for config in hook_configs:
        try:
            registry.register_from_config(config)
        except (ImportError, ValueError) as e:
            logger.error(f"Failed to register hook: {e}")
            if config.fail_on_error:
                raise

    return registry
