"""DashboardEventCollector - captures agent lifecycle events for real-time dashboard.

This component hooks into the agent execution lifecycle to emit WebSocket events.
Phase 1: Events stored in in-memory buffer (max 100 events).
Phase 3: Extended to emit via WebSocket using WebSocketManager.
"""

import traceback
from collections import deque
from datetime import datetime, timezone
from typing import TYPE_CHECKING, Optional

from pydantic import PrivateAttr

from flock.components import AgentComponent
from flock.dashboard.events import (
    AgentActivatedEvent,
    AgentCompletedEvent,
    AgentErrorEvent,
    MessagePublishedEvent,
    SubscriptionInfo,
    VisibilitySpec,
)
from flock.logging.logging import get_logger
from flock.runtime import Context


logger = get_logger("dashboard.collector")

if TYPE_CHECKING:  # pragma: no cover - type hints only
    from flock.agent import Agent
    from flock.artifacts import Artifact
    from flock.dashboard.websocket import WebSocketManager


class DashboardEventCollector(AgentComponent):
    """Collects agent lifecycle events for dashboard visualization.

    Implements AgentComponent interface to hook into agent execution:
    - on_pre_consume: emits agent_activated
    - on_post_publish: emits message_published
    - on_terminate: emits agent_completed
    - on_error: emits agent_error

    Phase 1: Events stored in in-memory deque (max 100, LRU eviction).
    Phase 3: Emits events via WebSocket using WebSocketManager.
    """

    # Use PrivateAttr for non-Pydantic fields (AgentComponent extends BaseModel)
    _events: deque[
        AgentActivatedEvent | MessagePublishedEvent | AgentCompletedEvent | AgentErrorEvent
    ] = PrivateAttr(default=None)

    # Track run start times for duration calculation
    _run_start_times: dict[str, float] = PrivateAttr(default_factory=dict)

    # WebSocketManager for broadcasting events
    _websocket_manager: Optional["WebSocketManager"] = PrivateAttr(default=None)

    def __init__(self, **data):
        super().__init__(**data)
        # In-memory buffer with max 100 events (LRU eviction)
        self._events = deque(maxlen=100)
        self._run_start_times = {}
        self._websocket_manager = None

    def set_websocket_manager(self, manager: "WebSocketManager") -> None:
        """Set WebSocketManager for broadcasting events.

        Args:
            manager: WebSocketManager instance to use for broadcasting
        """
        self._websocket_manager = manager

    @property
    def events(self) -> deque:
        """Access events buffer."""
        return self._events

    async def on_pre_consume(
        self, agent: "Agent", ctx: Context, inputs: list["Artifact"]
    ) -> list["Artifact"]:
        """Emit agent_activated event when agent begins consuming.

        Args:
            agent: The agent that is consuming
            ctx: Execution context with correlation_id
            inputs: Artifacts being consumed

        Returns:
            Unmodified inputs (pass-through)
        """
        # Record start time for duration calculation
        self._run_start_times[ctx.task_id] = datetime.now(timezone.utc).timestamp()

        # Extract consumed types and artifact IDs
        consumed_types = list({artifact.type for artifact in inputs})
        consumed_artifacts = [str(artifact.id) for artifact in inputs]

        # Extract produced types from agent outputs
        produced_types = [output.spec.type_name for output in agent.outputs]

        # Build subscription info from agent's subscriptions
        subscription_info = SubscriptionInfo(from_agents=[], channels=[], mode="both")

        if agent.subscriptions:
            # Get first subscription's config (agents typically have one)
            sub = agent.subscriptions[0]
            subscription_info.from_agents = list(sub.from_agents) if sub.from_agents else []
            subscription_info.channels = list(sub.channels) if sub.channels else []
            subscription_info.mode = sub.mode

        # Create and store event
        event = AgentActivatedEvent(
            correlation_id=str(ctx.correlation_id) if ctx.correlation_id else "",
            agent_name=agent.name,
            agent_id=agent.name,
            run_id=ctx.task_id,  # Unique ID for this agent run
            consumed_types=consumed_types,
            consumed_artifacts=consumed_artifacts,
            produced_types=produced_types,
            subscription_info=subscription_info,
            labels=list(agent.labels),
            tenant_id=agent.tenant_id,
            max_concurrency=agent.max_concurrency,
        )

        self._events.append(event)
        logger.info(f"Agent activated: {agent.name} (correlation_id={event.correlation_id})")

        # Broadcast via WebSocket if manager is configured
        if self._websocket_manager:
            await self._websocket_manager.broadcast(event)
        else:
            logger.warning("WebSocket manager not configured, event not broadcast")

        return inputs

    async def on_post_publish(self, agent: "Agent", ctx: Context, artifact: "Artifact") -> None:
        """Emit message_published event when artifact is published.

        Args:
            agent: The agent that published the artifact
            ctx: Execution context with correlation_id
            artifact: The published artifact
        """
        # Convert visibility to VisibilitySpec
        visibility_spec = self._convert_visibility(artifact.visibility)

        # Create and store event
        event = MessagePublishedEvent(
            correlation_id=str(ctx.correlation_id) if ctx.correlation_id else "",
            artifact_id=str(artifact.id),
            artifact_type=artifact.type,
            produced_by=artifact.produced_by,
            payload=artifact.payload,
            visibility=visibility_spec,
            tags=list(artifact.tags) if artifact.tags else [],
            partition_key=artifact.partition_key,
            version=artifact.version,
            consumers=[],  # Phase 1: empty, Phase 3: compute from subscription matching
        )

        self._events.append(event)
        logger.info(
            f"Message published: {artifact.type} by {artifact.produced_by} (correlation_id={event.correlation_id})"
        )

        # Broadcast via WebSocket if manager is configured
        if self._websocket_manager:
            await self._websocket_manager.broadcast(event)
        else:
            logger.warning("WebSocket manager not configured, event not broadcast")

    async def on_terminate(self, agent: "Agent", ctx: Context) -> None:
        """Emit agent_completed event when agent finishes successfully.

        Args:
            agent: The agent that completed
            ctx: Execution context with final state
        """
        # Calculate duration
        start_time = self._run_start_times.get(ctx.task_id)
        if start_time:
            duration_ms = (datetime.now(timezone.utc).timestamp() - start_time) * 1000
            del self._run_start_times[ctx.task_id]
        else:
            duration_ms = 0.0

        # Extract artifacts produced from context state (if tracked)
        artifacts_produced = ctx.state.get("artifacts_produced", [])
        if not isinstance(artifacts_produced, list):
            artifacts_produced = []

        # Extract metrics from context state (if tracked)
        metrics = ctx.state.get("metrics", {})
        if not isinstance(metrics, dict):
            metrics = {}

        # Create and store event
        event = AgentCompletedEvent(
            correlation_id=str(ctx.correlation_id) if ctx.correlation_id else "",
            agent_name=agent.name,
            run_id=ctx.task_id,
            duration_ms=duration_ms,
            artifacts_produced=artifacts_produced,
            metrics=metrics,
            final_state=dict(ctx.state),
        )

        self._events.append(event)

        # Broadcast via WebSocket if manager is configured
        if self._websocket_manager:
            await self._websocket_manager.broadcast(event)

    async def on_error(self, agent: "Agent", ctx: Context, error: Exception) -> None:
        """Emit agent_error event when agent execution fails.

        Args:
            agent: The agent that failed
            ctx: Execution context
            error: The exception that was raised
        """
        # Get error details
        error_type = type(error).__name__
        error_message = str(error)
        # Use traceback.format_exception to get traceback from exception object
        error_traceback = "".join(
            traceback.format_exception(type(error), error, error.__traceback__)
        )
        failed_at = datetime.now(timezone.utc).isoformat().replace("+00:00", "Z")

        # Clean up start time tracking
        if ctx.task_id in self._run_start_times:
            del self._run_start_times[ctx.task_id]

        # Create and store event
        event = AgentErrorEvent(
            correlation_id=str(ctx.correlation_id) if ctx.correlation_id else "",
            agent_name=agent.name,
            run_id=ctx.task_id,
            error_type=error_type,
            error_message=error_message,
            traceback=error_traceback,
            failed_at=failed_at,
        )

        self._events.append(event)

        # Broadcast via WebSocket if manager is configured
        if self._websocket_manager:
            await self._websocket_manager.broadcast(event)

    def _convert_visibility(self, visibility) -> VisibilitySpec:
        """Convert flock.visibility.Visibility to VisibilitySpec.

        Args:
            visibility: Visibility object from artifact

        Returns:
            VisibilitySpec for event serialization
        """
        # Get visibility kind from class name, stripping "Visibility" suffix
        class_name = type(visibility).__name__
        kind = class_name[: -len("Visibility")] if class_name.endswith("Visibility") else class_name

        spec = VisibilitySpec(kind=kind)

        # Extract type-specific fields
        if kind == "Private":
            spec.agents = list(visibility.agents) if hasattr(visibility, "agents") else []
        elif kind == "Labelled":
            spec.required_labels = (
                list(visibility.required_labels) if hasattr(visibility, "required_labels") else []
            )
        elif kind == "Tenant":
            spec.tenant_id = visibility.tenant_id if hasattr(visibility, "tenant_id") else None

        return spec


__all__ = ["DashboardEventCollector"]
