"""WebSocket API for real-time agent events.

This module implements a WebSocket endpoint at /ws/agents that broadcasts
real-time agent activity events to connected browser clients.

The WebSocket:
- Accepts connections from browser clients
- Sends initial state snapshot on connect
- Background task polls SessionEvent table every 2 seconds
- Broadcasts parsed events to all connected clients
- Cleans up on disconnect
"""

from __future__ import annotations

import asyncio
import json
from datetime import datetime, timedelta, timezone
from typing import Any

from fastapi import APIRouter, WebSocket, WebSocketDisconnect
from sqlmodel import select

from app.core.logging import get_logger
from app.db.session import async_session_maker
from app.models.monitoring import SessionEvent
from app.services.monitoring.event_parser import parse_session_event

logger = get_logger(__name__)

router = APIRouter()


class WebSocketConnectionManager:
    """Manages WebSocket connections and event broadcasting."""

    def __init__(self) -> None:
        self.active_connections: set[WebSocket] = set()
        self._broadcast_task: asyncio.Task[None] | None = None

    async def connect(self, websocket: WebSocket) -> None:
        """Accept a new WebSocket connection and send initial state."""
        await websocket.accept()
        self.active_connections.add(websocket)
        logger.info(
            "app.websocket.client_connected total=%d",
            len(self.active_connections),
        )
        await self._send_initial_state(websocket)

    async def disconnect(self, websocket: WebSocket) -> None:
        """Remove a WebSocket connection."""
        self.active_connections.discard(websocket)
        logger.info(
            "app.websocket.client_disconnected total=%d",
            len(self.active_connections),
        )

    async def broadcast_to_all(self, event: dict[str, Any]) -> None:
        """Send an event to all connected clients."""
        data = json.dumps(event)
        disconnected: set[WebSocket] = set()
        for connection in list(self.active_connections):
            try:
                await connection.send_text(data)
            except Exception:
                disconnected.add(connection)
        # Clean up broken connections
        for conn in disconnected:
            self.active_connections.discard(conn)

    async def _send_initial_state(self, websocket: WebSocket) -> None:
        """Send recent session events as initial state snapshot."""
        try:
            async with async_session_maker() as session:
                stmt = (
                    select(SessionEvent)
                    .order_by(SessionEvent.created_at.desc())
                    .limit(50)
                )
                result = await session.execute(stmt)
                events = result.scalars().all()

            # Send parsed events oldest first
            for event in reversed(events):
                try:
                    dashboard_events = parse_session_event(event)
                    for de in dashboard_events:
                        await websocket.send_json(de)
                except Exception as exc:
                    logger.warning(
                        "app.websocket.parse_error event_id=%s error=%s",
                        event.id,
                        exc,
                    )

            await websocket.send_json({
                "type": "initialState",
                "timestamp": datetime.now(timezone.utc).isoformat(),
            })
        except Exception as exc:
            logger.error(
                "app.websocket.initial_state_error error=%s", exc, exc_info=True
            )

    async def _broadcast_loop(self) -> None:
        """Background task: poll for new session events and broadcast them."""
        last_poll_time: datetime | None = None
        poll_interval = 2  # seconds

        while True:
            try:
                await asyncio.sleep(poll_interval)

                async with async_session_maker() as session:
                    if last_poll_time:
                        stmt = (
                            select(SessionEvent)
                            .where(SessionEvent.created_at > last_poll_time)
                            .order_by(SessionEvent.created_at.asc())
                        )
                    else:
                        # First poll: get events from last hour
                        cutoff = datetime.now(timezone.utc).replace(tzinfo=None) - timedelta(hours=1)
                        stmt = (
                            select(SessionEvent)
                            .where(SessionEvent.created_at > cutoff)
                            .order_by(SessionEvent.created_at.asc())
                        )

                    result = await session.execute(stmt)
                    events = result.scalars().all()

                if events:
                    last_poll_time = max(e.created_at for e in events)

                    for event in events:
                        try:
                            dashboard_events = parse_session_event(event)
                            for de in dashboard_events:
                                await self.broadcast_to_all(de)
                        except Exception as exc:
                            logger.warning(
                                "app.websocket.parse_error event_id=%s error=%s",
                                event.id,
                                exc,
                            )

            except asyncio.CancelledError:
                raise
            except Exception as exc:
                logger.error(
                    "app.websocket.broadcast_loop_error error=%s",
                    exc,
                    exc_info=True,
                )

    def start(self) -> None:
        """Start the background broadcast task."""
        loop = asyncio.get_running_loop()
        if self._broadcast_task is None or self._broadcast_task.done():
            self._broadcast_task = loop.create_task(self._broadcast_loop())
            logger.info("app.websocket.broadcast_started")

    async def stop(self) -> None:
        """Stop the background broadcast task."""
        if self._broadcast_task:
            self._broadcast_task.cancel()
            try:
                await self._broadcast_task
            except asyncio.CancelledError:
                pass
            self._broadcast_task = None
            logger.info("app.websocket.broadcast_stopped")


# Global connection manager instance
ws_manager = WebSocketConnectionManager()


@router.websocket("/ws/agents")
async def websocket_endpoint(websocket: WebSocket) -> None:
    """WebSocket endpoint for real-time agent events.

    Connect at: ws://localhost:8080/ws/agents

    Event types sent:
        - agentStatus: agent became active/idle/stalled
        - agentToolStart: agent started using a tool
        - agentToolDone: agent finished using a tool
        - agentChat: agent sent a text message
        - agentTask: new user task arrived
        - sessionStart: new session started
        - subagentSpawned: sub-agent spawned
        - subagentDespawned: sub-agent completed
        - initialState: initial state snapshot on connect
    """
    await ws_manager.connect(websocket)
    try:
        while True:
            try:
                data = await asyncio.wait_for(
                    websocket.receive_text(),
                    timeout=30,
                )
                # Client can send pings or commands — currently just keep-alive
                if data:
                    logger.debug("app.websocket.recv data=%s", data[:100])
            except asyncio.TimeoutError:
                # Send ping to keep connection alive
                try:
                    await websocket.send_text("ping")
                except Exception:
                    break
    except WebSocketDisconnect:
        pass
    except Exception as exc:
        logger.error("app.websocket.error error=%s", exc, exc_info=True)
    finally:
        await ws_manager.disconnect(websocket)


def broadcast_event(event: dict[str, Any]) -> None:
    """Queue an event for broadcasting to all connected clients.

    This is a synchronous function that can be called from any context
    (including background tasks) to push events to WebSocket clients.

    Args:
        event: Dashboard event dict to broadcast
    """
    try:
        loop = asyncio.get_running_loop()
        loop.create_task(ws_manager.broadcast_to_all(event))
    except RuntimeError:
        logger.debug("app.websocket.no_loop skipping broadcast")


def start_websocket_broadcast() -> None:
    """Start the WebSocket broadcast background task."""
    ws_manager.start()


async def stop_websocket_broadcast() -> None:
    """Stop the WebSocket broadcast background task."""
    await ws_manager.stop()