"""Streaming callback for real-time progress events.

Emits phase_started, phase_completed, phase_failed, and tool_called
events. Compatible with the WebSocket console broadcast if running.
"""

from __future__ import annotations

import logging
from datetime import datetime
from typing import Any, Callable, Dict, List, Optional

logger = logging.getLogger(__name__)


class StreamingProgressHandler:
    """Collects and broadcasts progress events for LangGraph execution.

    Events are stored in-memory and optionally forwarded to:
    - A callback function (for custom integrations)
    - The DataBridge console WebSocket server (if running)
    """

    def __init__(
        self,
        run_id: str = "",
        callback: Optional[Callable[[Dict[str, Any]], None]] = None,
        enable_console: bool = False,
    ):
        self.run_id = run_id
        self._callback = callback
        self._enable_console = enable_console
        self._events: List[Dict[str, Any]] = []

    # ------------------------------------------------------------------
    # Event emission
    # ------------------------------------------------------------------

    def emit(self, event_type: str, data: Optional[Dict[str, Any]] = None) -> Dict[str, Any]:
        """Emit a progress event.

        Args:
            event_type: One of phase_started, phase_completed, phase_failed,
                        tool_called, workflow_started, workflow_completed.
            data: Event-specific payload.

        Returns:
            The full event dict.
        """
        event = {
            "event_type": event_type,
            "run_id": self.run_id,
            "timestamp": datetime.utcnow().isoformat() + "Z",
            "data": data or {},
        }
        self._events.append(event)

        # Forward to callback
        if self._callback:
            try:
                self._callback(event)
            except Exception as e:
                logger.debug("Progress callback error: %s", e)

        # Forward to console WebSocket
        if self._enable_console:
            self._broadcast_to_console(event)

        return event

    def phase_started(self, phase: str, agent: str = "") -> Dict[str, Any]:
        """Emit a phase_started event."""
        return self.emit("phase_started", {"phase": phase, "agent": agent})

    def phase_completed(self, phase: str, agent: str = "", duration: float = 0.0) -> Dict[str, Any]:
        """Emit a phase_completed event."""
        return self.emit("phase_completed", {
            "phase": phase, "agent": agent, "duration_seconds": duration,
        })

    def phase_failed(self, phase: str, agent: str = "", error: str = "") -> Dict[str, Any]:
        """Emit a phase_failed event."""
        return self.emit("phase_failed", {
            "phase": phase, "agent": agent, "error": error,
        })

    def tool_called(self, tool_name: str, agent: str = "", args: Optional[Dict] = None) -> Dict[str, Any]:
        """Emit a tool_called event."""
        return self.emit("tool_called", {
            "tool": tool_name, "agent": agent, "args": args or {},
        })

    def workflow_started(self) -> Dict[str, Any]:
        """Emit a workflow_started event."""
        return self.emit("workflow_started", {"run_id": self.run_id})

    def workflow_completed(self, status: str = "completed", duration: float = 0.0) -> Dict[str, Any]:
        """Emit a workflow_completed event."""
        return self.emit("workflow_completed", {
            "status": status, "duration_seconds": duration,
        })

    # ------------------------------------------------------------------
    # History
    # ------------------------------------------------------------------

    def get_history(self, limit: int = 100) -> List[Dict[str, Any]]:
        """Get recent events for this run."""
        return list(self._events[-limit:])

    def get_by_type(self, event_type: str, limit: int = 50) -> List[Dict[str, Any]]:
        """Get events filtered by type."""
        matches = [e for e in self._events if e["event_type"] == event_type]
        return matches[-limit:]

    @property
    def event_count(self) -> int:
        return len(self._events)

    # ------------------------------------------------------------------
    # Console integration
    # ------------------------------------------------------------------

    def _broadcast_to_console(self, event: Dict[str, Any]) -> None:
        """Try to broadcast event to console WebSocket server."""
        try:
            from src.console_ws.broadcaster import ConsoleBroadcaster
            from src.console_ws.types import WebSocketMessage, WebSocketMessageType

            broadcaster = ConsoleBroadcaster.get_instance()
            payload = {
                "run_id": self.run_id,
                "event_type": event.get("event_type", "info"),
                "timestamp": event.get("timestamp"),
                "data": event.get("data", {}),
            }
            broadcaster.publish(
                "agents",
                WebSocketMessage(
                    type=WebSocketMessageType.AGENT_MESSAGE,
                    payload=payload,
                ),
            )

            data = event.get("data", {}) or {}
            agent_name = str(data.get("agent", "")).strip()
            if agent_name:
                broadcaster.agent_status(
                    agent_id=agent_name,
                    agent_name=agent_name,
                    status=self._status_from_event_type(str(event.get("event_type", ""))),
                    current_task=str(data.get("phase") or event.get("event_type") or ""),
                )
        except ImportError:
            pass  # Console not available
        except Exception as e:
            logger.debug("Console broadcast failed: %s", e)

    @staticmethod
    def _status_from_event_type(event_type: str) -> str:
        et = (event_type or "").lower()
        if et.endswith("failed"):
            return "error"
        if et.endswith("completed"):
            return "active"
        if et.endswith("started") or et == "tool_called":
            return "busy"
        return "active"

    # ------------------------------------------------------------------
    # As callback
    # ------------------------------------------------------------------

    def as_callback(self) -> Callable[[Dict[str, Any]], None]:
        """Return a callback function suitable for agent progress_callback."""
        def callback(data: Dict[str, Any]) -> None:
            event_type = data.get("status", "info")
            self.emit(event_type, data)
        return callback
