Atomizer/optimization_engine/plugins/hooks.py

"""
Core Hook System for Atomizer

Defines hook points in the optimization lifecycle and hook registration mechanism.
"""

from enum import Enum
from typing import Callable, Dict, Any, Optional
from dataclasses import dataclass
import logging

logger = logging.getLogger(__name__)


class HookPoint(Enum):
    """Enumeration of available hook points in the optimization lifecycle."""

    PRE_MESH = "pre_mesh"           # Before meshing
    POST_MESH = "post_mesh"         # After meshing, before solve
    PRE_SOLVE = "pre_solve"         # Before solver execution
    POST_SOLVE = "post_solve"       # After solve, before extraction
    POST_EXTRACTION = "post_extraction"  # After result extraction
    POST_CALCULATION = "post_calculation"  # After inline calculations (Phase 2.8), before reporting
    CUSTOM_OBJECTIVE = "custom_objective"  # Custom objective functions


@dataclass
class Hook:
    """
    Represents a single hook function to be executed at a specific point.

    Attributes:
        name: Unique identifier for this hook
        hook_point: When this hook should execute (HookPoint enum)
        function: The callable to execute
        description: Human-readable description of what this hook does
        priority: Execution order (lower = earlier, default=100)
        enabled: Whether this hook is currently active
    """

    name: str
    hook_point: HookPoint
    function: Callable[[Dict[str, Any]], Optional[Dict[str, Any]]]
    description: str
    priority: int = 100
    enabled: bool = True

    def execute(self, context: Dict[str, Any]) -> Optional[Dict[str, Any]]:
        """
        Execute this hook with the given context.

        Args:
            context: Dictionary containing relevant data for this hook point
                Common keys:
                - trial_number: Current trial number
                - design_variables: Current design variable values
                - sim_file: Path to simulation file
                - working_dir: Current working directory

        Returns:
            Optional dictionary with results or modifications to context

        Raises:
            Exception: Any exception from the hook function is logged and re-raised
        """
        if not self.enabled:
            logger.debug(f"Hook '{self.name}' is disabled, skipping")
            return None

        try:
            logger.info(f"Executing hook '{self.name}' at {self.hook_point.value}")
            result = self.function(context)
            logger.debug(f"Hook '{self.name}' completed successfully")
            return result

        except Exception as e:
            logger.error(f"Hook '{self.name}' failed: {e}", exc_info=True)
            raise

    def __repr__(self) -> str:
        status = "enabled" if self.enabled else "disabled"
        return f"Hook(name='{self.name}', point={self.hook_point.value}, priority={self.priority}, {status})"


class HookContext:
    """
    Context object passed to hooks containing all relevant data.

    This is a convenience wrapper around a dictionary that provides
    both dict-like access and attribute access.
    """

    def __init__(self, **kwargs):
        self._data = kwargs

    def __getitem__(self, key: str) -> Any:
        return self._data[key]

    def __setitem__(self, key: str, value: Any):
        self._data[key] = value

    def __contains__(self, key: str) -> bool:
        return key in self._data

    def get(self, key: str, default: Any = None) -> Any:
        return self._data.get(key, default)

    def update(self, other: Dict[str, Any]):
        self._data.update(other)

    def to_dict(self) -> Dict[str, Any]:
        return self._data.copy()

    def __repr__(self) -> str:
        keys = list(self._data.keys())
        return f"HookContext({', '.join(keys)})"