"""Core observability event types and emitter.

Event types for tracing agent execution: turns, pathways, nodes, tools, LLM calls.
"""

import asyncio
import logging
from collections import defaultdict
from collections.abc import Awaitable, Callable
from datetime import UTC, datetime
from enum import StrEnum
from typing import Any

from pydantic import BaseModel, Field

from core.common import JsonObject

logger = logging.getLogger(__name__)


# =============================================================================
# EVENT TYPES
# =============================================================================


class EventType(StrEnum):
    """All event types emitted by the runtime."""

    STATE_TRANSITION = "state_transition"
    STATE_MACHINE_ERROR = "state_machine_error"
    TURN_STARTED = "turn_started"
    TURN_COMPLETED = "turn_completed"
    TURN_FAILED = "turn_failed"
    PATHWAY_CREATED = "pathway_created"
    PATHWAY_STARTED = "pathway_started"
    PATHWAY_COMPLETED = "pathway_completed"
    NODE_STARTED = "node_started"
    NODE_COMPLETED = "node_completed"
    NODE_FAILED = "node_failed"
    TOOL_CALLED = "tool_called"
    TOOL_COMPLETED = "tool_completed"
    TOOL_FAILED = "tool_failed"
    LLM_REQUEST = "llm_request"
    LLM_RESPONSE = "llm_response"


def _utc_now() -> datetime:
    """Return current UTC datetime."""
    return datetime.now(UTC)


class Event(BaseModel):
    """Base event class."""

    type: EventType
    timestamp: datetime = Field(default_factory=_utc_now)
    thread_id: str | None = None
    turn_id: str | None = None
    execution_id: str | None = None
    agent_id: str | None = None


class NodeInfo(BaseModel):
    """Info about a node in a pathway."""

    id: str
    type: str
    prompt: str | None = None


# Turn events
class TurnStartedEvent(Event):
    type: EventType = EventType.TURN_STARTED
    message: str = ""


class TurnCompletedEvent(Event):
    type: EventType = EventType.TURN_COMPLETED
    message: str = ""
    response: str | None = None
    duration_ms: float | None = None
    success: bool = True


class TurnFailedEvent(Event):
    type: EventType = EventType.TURN_FAILED
    message: str = ""
    error: str | None = None
    duration_ms: float | None = None
    success: bool = False


# Pathway events
class PathwayCreatedEvent(Event):
    type: EventType = EventType.PATHWAY_CREATED
    pathway_id: str = ""
    pathway_name: str | None = None
    node_count: int = 0
    nodes: list[NodeInfo] = Field(default_factory=list)
    connections: list[tuple[str, str]] = Field(default_factory=list)


class PathwayStartedEvent(Event):
    type: EventType = EventType.PATHWAY_STARTED
    pathway_id: str = ""
    pathway_name: str | None = None
    node_count: int = 0
    nodes: list[NodeInfo] = Field(default_factory=list)


class PathwayCompletedEvent(Event):
    type: EventType = EventType.PATHWAY_COMPLETED
    pathway_id: str = ""
    pathway_name: str | None = None
    duration_ms: float | None = None
    success: bool = True
    outputs: JsonObject = Field(default_factory=dict)


# Node events
class NodeStartedEvent(Event):
    type: EventType = EventType.NODE_STARTED
    node_id: str = ""
    node_type: str = ""
    pathway_id: str = ""
    inputs: JsonObject = Field(default_factory=dict)


class NodeCompletedEvent(Event):
    type: EventType = EventType.NODE_COMPLETED
    node_id: str = ""
    node_type: str = ""
    pathway_id: str = ""
    outputs: JsonObject = Field(default_factory=dict)
    duration_ms: float | None = None
    success: bool = True


class NodeFailedEvent(Event):
    type: EventType = EventType.NODE_FAILED
    node_id: str = ""
    node_type: str = ""
    pathway_id: str = ""
    error: str | None = None
    duration_ms: float | None = None
    success: bool = False


# Tool events
class ToolCalledEvent(Event):
    type: EventType = EventType.TOOL_CALLED
    pathway_id: str | None = None
    node_id: str | None = None
    call_id: str | None = None
    tool_name: str = ""
    tool_args: JsonObject = Field(default_factory=dict)


class ToolCompletedEvent(Event):
    type: EventType = EventType.TOOL_COMPLETED
    pathway_id: str | None = None
    node_id: str | None = None
    call_id: str | None = None
    tool_name: str = ""
    tool_args: JsonObject = Field(default_factory=dict)
    result: Any = None
    duration_ms: float | None = None
    success: bool = True


class ToolFailedEvent(Event):
    type: EventType = EventType.TOOL_FAILED
    pathway_id: str | None = None
    node_id: str | None = None
    call_id: str | None = None
    tool_name: str = ""
    tool_args: JsonObject = Field(default_factory=dict)
    result: Any = None
    error: str | None = None
    duration_ms: float | None = None
    success: bool = False


# LLM events
class LLMRequestEvent(Event):
    type: EventType = EventType.LLM_REQUEST
    pathway_id: str | None = None
    node_id: str | None = None
    call_id: str | None = None
    model: str = ""
    prompt: str = ""


class LLMResponseEvent(Event):
    type: EventType = EventType.LLM_RESPONSE
    pathway_id: str | None = None
    node_id: str | None = None
    call_id: str | None = None
    model: str = ""
    prompt: str = ""
    response: str = ""
    tokens_in: int | None = None
    tokens_out: int | None = None
    duration_ms: float | None = None


# State transition
class StateTransitionEvent(Event):
    type: EventType = EventType.STATE_TRANSITION
    from_state: str = ""
    to_state: str = ""
    trigger: str = ""
    duration_ms: float | None = None


EventHandler = Callable[[Event], None] | Callable[[Event], Awaitable[None]]


# =============================================================================
# EVENT EMITTER
# =============================================================================


class EventEmitter:
    """Simple event emitter with sync and async handler support."""

    def __init__(self) -> None:
        self._handlers: dict[str, list[EventHandler]] = defaultdict(list)
        self._all_handlers: list[EventHandler] = []
        self._background_tasks: set[asyncio.Task[None]] = set()

    def on(self, event_type: str | EventType, handler: EventHandler) -> None:
        key = event_type.value if isinstance(event_type, EventType) else event_type
        self._handlers[key].append(handler)

    def on_all(self, handler: EventHandler) -> None:
        self._all_handlers.append(handler)

    def off(self, event_type: str | EventType, handler: EventHandler) -> None:
        key = event_type.value if isinstance(event_type, EventType) else event_type
        if handler in self._handlers[key]:
            self._handlers[key].remove(handler)

    def off_all(self, handler: EventHandler) -> None:
        if handler in self._all_handlers:
            self._all_handlers.remove(handler)

    def _track_task(self, coro: object) -> None:
        """Create a tracked background task that removes itself on completion."""
        task = asyncio.create_task(coro)  # type: ignore[arg-type]
        self._background_tasks.add(task)
        task.add_done_callback(self._background_tasks.discard)

    def emit(self, event: Event) -> None:
        key = event.type.value
        for handler in self._handlers[key]:
            try:
                result = handler(event)
                if asyncio.iscoroutine(result):
                    self._track_task(result)
            except Exception as e:
                logger.warning("Event handler error for %s: %s", key, e)

        for handler in self._all_handlers:
            try:
                result = handler(event)
                if asyncio.iscoroutine(result):
                    self._track_task(result)
            except Exception as e:
                logger.warning("Event handler error (all): %s", e)


__all__ = [
    "Event",
    # Emitter
    "EventEmitter",
    "EventHandler",
    # Event types
    "EventType",
    "LLMRequestEvent",
    "LLMResponseEvent",
    "NodeCompletedEvent",
    "NodeFailedEvent",
    "NodeInfo",
    "NodeStartedEvent",
    "PathwayCompletedEvent",
    "PathwayCreatedEvent",
    "PathwayStartedEvent",
    "StateTransitionEvent",
    "ToolCalledEvent",
    "ToolCompletedEvent",
    "ToolFailedEvent",
    "TurnCompletedEvent",
    "TurnFailedEvent",
    "TurnStartedEvent",
]
