server/webrtc_stream.py

"""
WebRTC Stream Manager for real-time video streaming using aiortc.
"""
import asyncio
import logging
from typing import Dict, Optional
import numpy as np
import cv2

from aiortc import RTCPeerConnection, RTCSessionDescription, VideoStreamTrack
from aiortc.contrib.media import MediaRelay
from av import VideoFrame

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


class VideoFrameTrack(VideoStreamTrack):
    """
    Custom video track that streams frames pushed from the server.
    """
    kind = "video"

    def __init__(self, fps: int = 30):
        super().__init__()
        self.fps = fps
        self.frame_queue: asyncio.Queue = asyncio.Queue(maxsize=60)
        self._timestamp = 0
        self._time_base = 1 / fps
        self._running = True

    async def recv(self) -> VideoFrame:
        """Receive the next frame to send."""
        try:
            # Wait for a frame with timeout
            frame_data = await asyncio.wait_for(self.frame_queue.get(), timeout=1.0)

            # Convert numpy array (BGR) to VideoFrame
            if isinstance(frame_data, np.ndarray):
                # Convert BGR to RGB
                frame_rgb = cv2.cvtColor(frame_data, cv2.COLOR_BGR2RGB)
                frame = VideoFrame.from_ndarray(frame_rgb, format="rgb24")
            else:
                frame = frame_data

            # Set timing
            frame.pts = self._timestamp
            frame.time_base = self._time_base
            self._timestamp += 1

            return frame

        except asyncio.TimeoutError:
            # Return a black frame if no data available
            black_frame = np.zeros((256, 256, 3), dtype=np.uint8)
            frame = VideoFrame.from_ndarray(black_frame, format="rgb24")
            frame.pts = self._timestamp
            frame.time_base = self._time_base
            self._timestamp += 1
            return frame

    async def push_frame(self, frame: np.ndarray):
        """Push a frame to be sent via WebRTC."""
        try:
            # Drop oldest frame if queue is full
            if self.frame_queue.full():
                try:
                    self.frame_queue.get_nowait()
                except asyncio.QueueEmpty:
                    pass
            await self.frame_queue.put(frame)
        except Exception as e:
            logger.error(f"Error pushing frame: {e}")

    def stop(self):
        """Stop the track."""
        self._running = False
        super().stop()


class WebRTCConnection:
    """Manages a single WebRTC peer connection."""

    def __init__(self, session_id: str, fps: int = 30):
        self.session_id = session_id
        self.pc = RTCPeerConnection()
        self.video_track = VideoFrameTrack(fps=fps)
        self.connected = False
        self._connection_event = asyncio.Event()

        # Add video track to peer connection
        self.pc.addTrack(self.video_track)

        # Handle connection state changes
        @self.pc.on("connectionstatechange")
        async def on_connectionstatechange():
            state = self.pc.connectionState
            logger.info(f"[WebRTC:{session_id}] Connection state: {state}")
            if state == "connected":
                self.connected = True
                self._connection_event.set()
            elif state in ("failed", "closed", "disconnected"):
                self.connected = False
                self._connection_event.clear()

        @self.pc.on("iceconnectionstatechange")
        async def on_iceconnectionstatechange():
            logger.info(f"[WebRTC:{session_id}] ICE state: {self.pc.iceConnectionState}")

    async def handle_offer(self, sdp: str, type: str) -> dict:
        """Handle an SDP offer and return an answer."""
        offer = RTCSessionDescription(sdp=sdp, type=type)
        await self.pc.setRemoteDescription(offer)

        answer = await self.pc.createAnswer()
        await self.pc.setLocalDescription(answer)

        return {
            "sdp": self.pc.localDescription.sdp,
            "type": self.pc.localDescription.type
        }

    async def send_frame(self, frame: np.ndarray):
        """Send a video frame."""
        if self.connected:
            await self.video_track.push_frame(frame)

    async def wait_for_connection(self, timeout: float = 5.0) -> bool:
        """Wait for the connection to be established."""
        try:
            await asyncio.wait_for(self._connection_event.wait(), timeout=timeout)
            return True
        except asyncio.TimeoutError:
            return False

    async def close(self):
        """Close the connection."""
        self.video_track.stop()
        await self.pc.close()
        self.connected = False


class WebRTCManager:
    """Manages multiple WebRTC connections."""

    def __init__(self):
        self.connections: Dict[str, WebRTCConnection] = {}
        self.relay = MediaRelay()

    async def handle_offer(self, session_id: str, sdp: str, type: str, fps: int = 30) -> dict:
        """Handle a WebRTC offer for a session."""
        # Close existing connection if any
        if session_id in self.connections:
            await self.close_connection(session_id)

        # Create new connection
        connection = WebRTCConnection(session_id, fps=fps)
        self.connections[session_id] = connection

        # Handle the offer
        answer = await connection.handle_offer(sdp, type)
        logger.info(f"[WebRTC] Created connection for session {session_id}")

        return answer

    def is_connected(self, session_id: str) -> bool:
        """Check if a session is connected."""
        connection = self.connections.get(session_id)
        return connection.connected if connection else False

    async def wait_for_connection(self, session_id: str, timeout: float = 5.0) -> bool:
        """Wait for a session to be connected."""
        connection = self.connections.get(session_id)
        if not connection:
            return False
        return await connection.wait_for_connection(timeout)

    async def send_frame(self, session_id: str, frame: np.ndarray):
        """Send a frame to a session."""
        connection = self.connections.get(session_id)
        if connection and connection.connected:
            await connection.send_frame(frame)

    async def close_connection(self, session_id: str):
        """Close a session's connection."""
        connection = self.connections.pop(session_id, None)
        if connection:
            await connection.close()
            logger.info(f"[WebRTC] Closed connection for session {session_id}")


# Global manager instance
webrtc_manager = WebRTCManager()