api/ws.py

"""
WebSocket handlers for real-time camera feeds and interactive VQA chat.

Two WebSocket endpoints:
- /ws/feed/{cam_id} — Streams detection results from a camera
- /ws/chat — Interactive VQA chat (send question + camera_id, get answer)
"""

from __future__ import annotations

import json
import time
import asyncio
import logging
from typing import Optional

from fastapi import WebSocket, WebSocketDisconnect

logger = logging.getLogger(__name__)


class ConnectionManager:
    """Manages active WebSocket connections grouped by channel (camera_id or chat)."""

    def __init__(self) -> None:
        self._connections: dict[str, list[WebSocket]] = {}

    async def connect(self, ws: WebSocket, channel: str) -> None:
        await ws.accept()
        self._connections.setdefault(channel, []).append(ws)
        logger.info("WebSocket connected: channel=%s, total=%d", channel, len(self._connections[channel]))

    def disconnect(self, ws: WebSocket, channel: str) -> None:
        conns = self._connections.get(channel, [])
        if ws in conns:
            conns.remove(ws)
        if not conns:
            self._connections.pop(channel, None)
        logger.info("WebSocket disconnected: channel=%s", channel)

    async def broadcast(self, channel: str, data: dict) -> None:
        """Send data to all connections in a channel."""
        conns = self._connections.get(channel, [])
        dead = []
        for ws in conns:
            try:
                await ws.send_json(data)
            except Exception:
                dead.append(ws)
        for ws in dead:
            conns.remove(ws)

    @property
    def active_channels(self) -> list[str]:
        return list(self._connections.keys())

    def connection_count(self, channel: str) -> int:
        return len(self._connections.get(channel, []))


# Global connection manager
ws_manager = ConnectionManager()


async def feed_websocket(ws: WebSocket, camera_id: str) -> None:
    """
    Stream detection results from a camera feed.

    Sends JSON messages with detection/alert results as they occur.
    Client receives periodic updates while the camera is active.
    """
    await ws_manager.connect(ws, f"feed:{camera_id}")

    try:
        while True:
            # Check for client messages (ping/config)
            try:
                data = await asyncio.wait_for(ws.receive_text(), timeout=5.0)
                msg = json.loads(data)
                # Client can send config updates
                if msg.get("type") == "ping":
                    await ws.send_json({"type": "pong", "timestamp": time.time()})
            except asyncio.TimeoutError:
                pass
            except json.JSONDecodeError:
                pass

            # Send latest results for this camera
            from api.deps import get_model_manager
            manager = get_model_manager()

            status = {
                "type": "status",
                "camera_id": camera_id,
                "timestamp": time.time(),
                "model_loaded": manager.model is not None,
            }
            await ws.send_json(status)

    except WebSocketDisconnect:
        ws_manager.disconnect(ws, f"feed:{camera_id}")
    except Exception as e:
        logger.error("Feed WebSocket error for %s: %s", camera_id, e)
        ws_manager.disconnect(ws, f"feed:{camera_id}")


async def chat_websocket(ws: WebSocket) -> None:
    """
    Interactive VQA chat over WebSocket.

    Client sends: {"question": "...", "camera_id": "...", "task_type": "vqa"}
    Server responds: {"answer": "...", "confidence": 0.85, "expert_used": "vqa", ...}
    """
    await ws_manager.connect(ws, "chat")

    try:
        while True:
            data = await ws.receive_text()
            msg = json.loads(data)

            question = msg.get("question", "")
            camera_id = msg.get("camera_id")
            task_type = msg.get("task_type", "vqa")

            if not question:
                await ws.send_json({"error": "Missing 'question' field"})
                continue

            # Process through the agent framework
            from api.deps import get_model_manager
            from agents.base import Task

            manager = get_model_manager()
            start = time.time()

            # Resolve image from camera if provided
            image_ref = msg.get("image_path", "")
            if camera_id and manager.camera_manager:
                cam_result = manager.camera_manager.get_frame(camera_id)
                if cam_result is not None:
                    frame, _ts = cam_result
                    import cv2
                    import tempfile
                    tmp_path = tempfile.mktemp(suffix=".jpg")
                    cv2.imwrite(tmp_path, frame)
                    image_ref = tmp_path

            task = Task(
                type=task_type,
                payload={
                    "image_ref": image_ref,
                    "query": question,
                    "context": {"camera_id": camera_id},
                },
                source="ws_chat",
                expert_hint=task_type if task_type in {
                    "vqa", "detect", "alert", "caption", "count", "ocr", "reason"
                } else "vqa",
            )

            result = manager.submit_task(task)

            response = {
                "type": "answer",
                "answer": result.answer,
                "confidence": result.confidence,
                "expert_used": result.expert_used,
                "processing_time_ms": (time.time() - start) * 1000,
                "task_id": result.task_id,
                "error": result.error,
            }
            await ws.send_json(response)

    except WebSocketDisconnect:
        ws_manager.disconnect(ws, "chat")
    except Exception as e:
        logger.error("Chat WebSocket error: %s", e)
        ws_manager.disconnect(ws, "chat")