"""Hooks system for observing and intercepting agent behavior."""

from __future__ import annotations

from abc import ABC, abstractmethod
from dataclasses import dataclass, field
from enum import Enum
from typing import TYPE_CHECKING, Any, Callable

if TYPE_CHECKING:
    from agenai.agent.context import AgentContext
    from agenai.agent.tools.base import ToolResult


class HookEvent(str, Enum):
    """Events that can trigger hooks."""

    # Agent lifecycle
    AGENT_START = "agent_start"
    AGENT_END = "agent_end"
    AGENT_ERROR = "agent_error"

    # Step lifecycle
    STEP_START = "step_start"
    STEP_END = "step_end"

    # Tool lifecycle
    TOOL_START = "tool_start"
    TOOL_END = "tool_end"
    TOOL_ERROR = "tool_error"

    # LLM lifecycle
    LLM_START = "llm_start"
    LLM_END = "llm_end"
    LLM_ERROR = "llm_error"


@dataclass
class HookData:
    """Data passed to hooks."""

    event: HookEvent
    context: AgentContext | None = None

    # For tool events
    tool_name: str | None = None
    tool_args: dict[str, Any] | None = None
    tool_result: ToolResult | None = None

    # For step events
    step_number: int | None = None

    # For error events
    error: Exception | None = None

    # For LLM events
    messages: list[dict[str, Any]] | None = None
    response: Any | None = None

    # Custom data
    extra: dict[str, Any] = field(default_factory=dict)


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

    @property
    @abstractmethod
    def events(self) -> list[HookEvent]:
        """List of events this hook handles."""
        ...

    @abstractmethod
    def __call__(self, data: HookData) -> None:
        """Handle the event."""
        ...


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

    def __init__(
        self,
        events: list[HookEvent],
        handler: Callable[[HookData], None],
    ):
        self._events = events
        self._handler = handler

    @property
    def events(self) -> list[HookEvent]:
        return self._events

    def __call__(self, data: HookData) -> None:
        self._handler(data)


@dataclass
class HookRegistry:
    """Registry for managing hooks."""

    _hooks: dict[HookEvent, list[Hook]] = field(default_factory=dict)

    def register(self, hook: Hook) -> None:
        """Register a hook for its events."""
        for event in hook.events:
            if event not in self._hooks:
                self._hooks[event] = []
            self._hooks[event].append(hook)

    def on(
        self,
        *events: HookEvent,
    ) -> Callable[[Callable[[HookData], None]], Hook]:
        """Decorator to register a function as a hook."""

        def decorator(fn: Callable[[HookData], None]) -> Hook:
            hook = FunctionHook(list(events), fn)
            self.register(hook)
            return hook

        return decorator

    def trigger(self, data: HookData) -> None:
        """Trigger all hooks for an event."""
        hooks = self._hooks.get(data.event, [])
        for hook in hooks:
            try:
                hook(data)
            except Exception:
                # Hooks should not crash the agent
                # In production, you'd log this
                pass

    def clear(self) -> None:
        """Remove all hooks."""
        self._hooks.clear()


# Convenience functions for creating hooks
def on_tool_start(fn: Callable[[HookData], None]) -> FunctionHook:
    """Create a hook for tool start events."""
    return FunctionHook([HookEvent.TOOL_START], fn)


def on_tool_end(fn: Callable[[HookData], None]) -> FunctionHook:
    """Create a hook for tool end events."""
    return FunctionHook([HookEvent.TOOL_END], fn)


def on_step(fn: Callable[[HookData], None]) -> FunctionHook:
    """Create a hook for step events."""
    return FunctionHook([HookEvent.STEP_START, HookEvent.STEP_END], fn)


def on_error(fn: Callable[[HookData], None]) -> FunctionHook:
    """Create a hook for error events."""
    return FunctionHook(
        [HookEvent.AGENT_ERROR, HookEvent.TOOL_ERROR, HookEvent.LLM_ERROR],
        fn,
    )
