src/core/events.py

"""
Event and hook system for monitoring.

Provides:
- Typed events for all operations
- Event subscription via handlers
- Event hooks for external monitoring
- Structured change logging
"""

from abc import ABC, abstractmethod
from collections.abc import Callable
from datetime import UTC, datetime
from enum import Enum
from typing import Any

from pydantic import BaseModel, Field

from config.logging import logger

__all__ = [
    # Budget events
    "BudgetEvent",
    "BudgetExceededEvent",
    "BudgetWarningEvent",
    "EdgeAddedEvent",
    "EdgeRemovedEvent",
    "EdgeUpdatedEvent",
    # Base event
    "Event",
    "EventBus",
    # Event handler
    "EventHandler",
    "EventPriority",
    # Event types
    "EventType",
    # Execution events
    "ExecutionEvent",
    "GlobalEventBus",
    # Graph events
    "GraphEvent",
    # Logging handler
    "LoggingEventHandler",
    # Memory events
    "MemoryEvent",
    "MemoryExpiredEvent",
    "MemoryReadEvent",
    "MemoryWriteEvent",
    "MetricsEventHandler",
    "NodeAddedEvent",
    "NodeRemovedEvent",
    "NodeReplacedEvent",
    "RunCompletedEvent",
    "RunStartedEvent",
    "StepCompletedEvent",
    "StepFailedEvent",
    "StepRetriedEvent",
    "StepStartedEvent",
    "global_event_bus",
]


class EventType(str, Enum):
    """Event types."""

    # Graph events
    NODE_ADDED = "node_added"
    NODE_REMOVED = "node_removed"
    NODE_REPLACED = "node_replaced"
    EDGE_ADDED = "edge_added"
    EDGE_REMOVED = "edge_removed"
    EDGE_UPDATED = "edge_updated"

    # Execution events
    RUN_STARTED = "run_started"
    RUN_COMPLETED = "run_completed"
    STEP_STARTED = "step_started"
    STEP_COMPLETED = "step_completed"
    STEP_FAILED = "step_failed"
    STEP_RETRIED = "step_retried"
    PLAN_CREATED = "plan_created"
    PLAN_TOPOLOGY_CHANGED = "plan_topology_changed"

    # Memory events
    MEMORY_WRITE = "memory_write"
    MEMORY_READ = "memory_read"
    MEMORY_EXPIRED = "memory_expired"
    MEMORY_COMPRESSED = "memory_compressed"
    MEMORY_PROMOTED = "memory_promoted"

    # Budget events
    BUDGET_WARNING = "budget_warning"
    BUDGET_EXCEEDED = "budget_exceeded"

    # Metrics events
    METRICS_UPDATED = "metrics_updated"


class EventPriority(int, Enum):
    """Event processing priority."""

    LOW = 0
    NORMAL = 1
    HIGH = 2
    CRITICAL = 3


class Event(BaseModel):
    event_type: EventType
    timestamp: datetime = Field(default_factory=lambda: datetime.now(UTC))
    source: str | None = None
    priority: EventPriority = EventPriority.NORMAL
    metadata: dict[str, Any] = Field(default_factory=dict)

    def to_dict(self) -> dict[str, Any]:
        """Serialize the event to a dict for logging/transport."""
        return {
            "event_type": self.event_type.value,
            "timestamp": self.timestamp.isoformat(),
            "source": self.source,
            "priority": self.priority.value,
            "metadata": self.metadata,
        }


class GraphEvent(Event):
    graph_id: str | None = None


class NodeAddedEvent(GraphEvent):
    event_type: EventType = EventType.NODE_ADDED
    node_id: str = ""
    node_data: dict[str, Any] = Field(default_factory=dict)
    connected_to: list[str] = Field(default_factory=list)


class NodeRemovedEvent(GraphEvent):
    event_type: EventType = EventType.NODE_REMOVED
    node_id: str = ""
    migration_policy: str = "discard"
    state_archived: bool = False


class NodeReplacedEvent(GraphEvent):
    event_type: EventType = EventType.NODE_REPLACED
    old_node_id: str = ""
    new_node_id: str = ""
    state_migrated: bool = False


class EdgeAddedEvent(GraphEvent):
    event_type: EventType = EventType.EDGE_ADDED
    source_id: str = ""
    target_id: str = ""
    weight: float = 1.0
    edge_data: dict[str, Any] = Field(default_factory=dict)


class EdgeRemovedEvent(GraphEvent):
    event_type: EventType = EventType.EDGE_REMOVED
    source_id: str = ""
    target_id: str = ""


class EdgeUpdatedEvent(GraphEvent):
    event_type: EventType = EventType.EDGE_UPDATED
    source_id: str = ""
    target_id: str = ""
    old_weight: float = 0.0
    new_weight: float = 0.0
    changes: dict[str, Any] = Field(default_factory=dict)


class ExecutionEvent(Event):
    run_id: str | None = None
    graph_id: str | None = None


class RunStartedEvent(ExecutionEvent):
    event_type: EventType = EventType.RUN_STARTED
    query: str = ""
    config: dict[str, Any] = Field(default_factory=dict)
    num_agents: int = 0


class RunCompletedEvent(ExecutionEvent):
    event_type: EventType = EventType.RUN_COMPLETED
    success: bool = True
    answer: str = ""
    total_steps: int = 0
    total_tokens: int = 0
    duration_ms: float = 0.0
    errors: list[str] = Field(default_factory=list)


class StepStartedEvent(ExecutionEvent):
    event_type: EventType = EventType.STEP_STARTED
    agent_id: str = ""
    step_index: int = 0
    predecessors: list[str] = Field(default_factory=list)


class StepCompletedEvent(ExecutionEvent):
    event_type: EventType = EventType.STEP_COMPLETED
    agent_id: str = ""
    step_index: int = 0
    success: bool = True
    response_length: int = 0
    tokens_used: int = 0
    duration_ms: float = 0.0


class StepFailedEvent(ExecutionEvent):
    event_type: EventType = EventType.STEP_FAILED
    priority: EventPriority = EventPriority.HIGH
    agent_id: str = ""
    step_index: int = 0
    error_type: str = ""
    error_message: str = ""
    will_retry: bool = False


class StepRetriedEvent(ExecutionEvent):
    event_type: EventType = EventType.STEP_RETRIED
    agent_id: str = ""
    attempt: int = 0
    max_attempts: int = 0
    delay_ms: float = 0.0


class MemoryEvent(Event):
    agent_id: str | None = None


class MemoryWriteEvent(MemoryEvent):
    event_type: EventType = EventType.MEMORY_WRITE
    key: str = ""
    value_size: int = 0
    memory_level: str = "working"


class MemoryReadEvent(MemoryEvent):
    event_type: EventType = EventType.MEMORY_READ
    key: str = ""
    found: bool = False
    memory_level: str = "working"


class MemoryExpiredEvent(MemoryEvent):
    event_type: EventType = EventType.MEMORY_EXPIRED
    key: str = ""
    ttl_seconds: float = 0.0


class BudgetEvent(Event):
    run_id: str | None = None


class BudgetWarningEvent(BudgetEvent):
    event_type: EventType = EventType.BUDGET_WARNING
    budget_type: str = ""
    current_value: float = 0.0
    limit: float = 0.0
    ratio: float = 0.0


class BudgetExceededEvent(BudgetEvent):
    event_type: EventType = EventType.BUDGET_EXCEEDED
    priority: EventPriority = EventPriority.CRITICAL
    budget_type: str = ""
    current_value: float = 0.0
    limit: float = 0.0
    action_taken: str = ""


class EventHandler[T: Event](ABC):
    """
    Base event handler.

    Attributes:
        raise_error: If True — exceptions from handle() propagate outward.
                     If False (default) — errors are logged and the bus
                     continues with the next handlers.

    """

    raise_error: bool = False

    @abstractmethod
    def handle(self, event: T) -> None:
        """Handle the event."""

    def can_handle(self, event: Event) -> bool:
        """Return True if the handler is ready to accept the event."""
        return True


class EventBus:
    """Simple event bus with type-based subscribers and global listeners."""

    def __init__(self):
        self._handlers: dict[EventType | None, list[EventHandler]] = {}
        self._global_handlers: list[EventHandler] = []
        self._enabled: bool = True

    def subscribe(
        self,
        event_type: EventType | None,
        handler: EventHandler | Callable[[Event], None],
    ) -> None:
        """Subscribe to events of a specific type or all events (when None)."""
        event_handler: EventHandler
        if isinstance(handler, EventHandler):
            event_handler = handler
        elif callable(handler):
            event_handler = CallableHandler(handler)
        else:
            msg = "Handler must be EventHandler or callable"
            raise TypeError(msg)

        if event_type is None:
            self._global_handlers.append(event_handler)
        else:
            if event_type not in self._handlers:
                self._handlers[event_type] = []
            self._handlers[event_type].append(event_handler)

    def unsubscribe(
        self,
        event_type: EventType | None,
        handler: EventHandler,
    ) -> None:
        """Unsubscribe from events of the specified type or from global events."""
        if event_type is None:
            if handler in self._global_handlers:
                self._global_handlers.remove(handler)
        elif event_type in self._handlers and handler in self._handlers[event_type]:
            self._handlers[event_type].remove(handler)

    def publish(self, event: Event) -> None:
        """Dispatch the event to all matching subscribers."""
        if not self._enabled:
            return

        all_handlers = []
        all_handlers.extend(self._global_handlers)
        if event.event_type in self._handlers:
            all_handlers.extend(self._handlers[event.event_type])
        for handler in all_handlers:
            try:
                if handler.can_handle(event):
                    handler.handle(event)
            except Exception as e:
                if handler.raise_error:
                    raise
                logger.debug(
                    "Event handler '{}' failed (set raise_error=True to propagate): {}",
                    type(handler).__name__,
                    e,
                )

    def enable(self) -> None:
        """Enable event processing."""
        self._enabled = True

    def disable(self) -> None:
        """Disable event processing."""
        self._enabled = False

    def clear(self) -> None:
        """Clear all registered handlers."""
        self._handlers.clear()
        self._global_handlers.clear()


class CallableHandler(EventHandler):
    """Wrapper that allows subscribing plain functions as event handlers."""

    def __init__(self, fn: Callable[[Event], None]):
        self._fn = fn

    def handle(self, event: Event) -> None:
        self._fn(event)


_global_bus: EventBus | None = None


def global_event_bus() -> EventBus:
    """Return the global event bus instance (create on first access)."""
    global _global_bus  # noqa: PLW0603
    if _global_bus is None:
        _global_bus = EventBus()
    return _global_bus


GlobalEventBus = global_event_bus


class LoggingEventHandler(EventHandler):
    """Handler that formats events and sends them to the log."""

    def __init__(
        self,
        log_level: str = "INFO",
        include_metadata: bool = True,
        format_func: Callable[[Event], str] | None = None,
    ):
        self.log_level = log_level
        self.include_metadata = include_metadata
        self.format_func = format_func or self._default_format
        self._logger = logger

    def handle(self, event: Event) -> None:
        """Write the event to the log according to priority and formatter."""
        message = self.format_func(event)

        level = self.log_level
        if event.priority == EventPriority.CRITICAL:
            level = "ERROR"
        elif event.priority == EventPriority.HIGH:
            level = "WARNING"

        if hasattr(self._logger, "log"):
            import logging

            level_int = getattr(logging, level.upper(), logging.INFO) if isinstance(level, str) else level
            self._logger.log(level_int, message)
        else:
            getattr(self._logger, level.lower(), self._logger.info)(message)

    def _default_format(self, event: Event) -> str:
        """Build a default string for different event types."""
        parts = [
            f"[{event.event_type.value}]",
            f"source={event.source}" if event.source else "",
        ]

        if isinstance(event, NodeAddedEvent):
            parts.append(f"node={event.node_id}")
        elif isinstance(event, NodeRemovedEvent):
            parts.append(f"node={event.node_id} policy={event.migration_policy}")
        elif isinstance(event, EdgeAddedEvent):
            parts.append(f"edge={event.source_id}->{event.target_id} weight={event.weight}")
        elif isinstance(event, StepCompletedEvent):
            parts.append(f"agent={event.agent_id} success={event.success} tokens={event.tokens_used}")
        elif isinstance(event, StepFailedEvent):
            parts.append(f"agent={event.agent_id} error={event.error_message}")
        elif isinstance(event, BudgetWarningEvent):
            parts.append(f"{event.budget_type}: {event.current_value}/{event.limit} ({event.ratio:.1%})")
        elif isinstance(event, RunCompletedEvent):
            parts.append(f"success={event.success} steps={event.total_steps} tokens={event.total_tokens}")

        if self.include_metadata and event.metadata:
            parts.append(f"metadata={event.metadata}")

        return " ".join(filter(None, parts))


class MetricsEventHandler(EventHandler):
    """Metrics aggregator based on incoming events."""

    def __init__(self):
        self._event_counts: dict[str, int] = {}
        self._total_tokens: int = 0
        self._total_duration_ms: float = 0.0
        self._errors: list[dict[str, Any]] = []
        self._step_durations: list[float] = []
        self._budget_warnings: int = 0
        self._runs_completed: int = 0
        self._runs_failed: int = 0

    def handle(self, event: Event) -> None:
        """Accumulate event statistics (tokens, errors, durations)."""
        event_type = event.event_type.value
        self._event_counts[event_type] = self._event_counts.get(event_type, 0) + 1

        if isinstance(event, StepCompletedEvent):
            self._total_tokens += event.tokens_used
            self._step_durations.append(event.duration_ms)
            self._total_duration_ms += event.duration_ms

        elif isinstance(event, StepFailedEvent):
            self._errors.append(
                {
                    "agent_id": event.agent_id,
                    "error_type": event.error_type,
                    "error_message": event.error_message,
                    "timestamp": event.timestamp.isoformat(),
                }
            )

        elif isinstance(event, BudgetWarningEvent):
            self._budget_warnings += 1

        elif isinstance(event, RunCompletedEvent):
            if event.success:
                self._runs_completed += 1
            else:
                self._runs_failed += 1

    def get_metrics(self) -> dict[str, Any]:
        """Return aggregated metrics from collected events."""
        return {
            "event_counts": dict(self._event_counts),
            "total_tokens": self._total_tokens,
            "total_duration_ms": self._total_duration_ms,
            "avg_step_duration_ms": (
                sum(self._step_durations) / len(self._step_durations) if self._step_durations else 0.0
            ),
            "errors_count": len(self._errors),
            "errors": self._errors[-10:],
            "budget_warnings": self._budget_warnings,
            "runs_completed": self._runs_completed,
            "runs_failed": self._runs_failed,
        }

    def reset(self) -> None:
        """Reset accumulated metrics."""
        self._event_counts.clear()
        self._total_tokens = 0
        self._total_duration_ms = 0.0
        self._errors.clear()
        self._step_durations.clear()
        self._budget_warnings = 0
        self._runs_completed = 0
        self._runs_failed = 0


def emit_event(event: Event) -> None:
    """Publish an event through the global bus."""
    global_event_bus().publish(event)


def on_event(event_type: EventType | None = None):
    """Decorator for subscribing a function to events of the given type."""

    def decorator(fn: Callable[[Event], None]):
        global_event_bus().subscribe(event_type, fn)
        return fn

    return decorator