server/livekit_stream.py

"""
LiveKit Stream Manager for real-time video streaming.
Uses LiveKit Python SDK for WebRTC video publishing.
"""
import asyncio
import logging
import os
from typing import Dict, Optional
import numpy as np
import cv2

from livekit import rtc, api

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

# LiveKit server configuration
LIVEKIT_URL = os.getenv("LIVEKIT_URL", "ws://localhost:7880")
LIVEKIT_API_KEY = os.getenv("LIVEKIT_API_KEY", "devkey")
LIVEKIT_API_SECRET = os.getenv("LIVEKIT_API_SECRET", "secret")


class LiveKitVideoPublisher:
    """Publishes video frames to a LiveKit room."""

    def __init__(self, room_name: str, participant_identity: str, fps: int = 30):
        self.room_name = room_name
        self.participant_identity = participant_identity
        self.fps = fps
        self.room: Optional[rtc.Room] = None
        self.video_source: Optional[rtc.VideoSource] = None
        self.video_track: Optional[rtc.LocalVideoTrack] = None
        self._connected = False
        self._width = 256
        self._height = 256

    async def connect(self) -> bool:
        """Connect to LiveKit room and publish video track."""
        try:
            # Generate access token
            token = api.AccessToken(LIVEKIT_API_KEY, LIVEKIT_API_SECRET)
            token.with_identity(self.participant_identity)
            token.with_name("Avatar Bot")
            token.with_grants(api.VideoGrants(
                room_join=True,
                room=self.room_name,
                can_publish=True,
                can_subscribe=False,
            ))
            jwt_token = token.to_jwt()

            # Create and connect to room
            self.room = rtc.Room()

            @self.room.on("connected")
            def on_connected():
                logger.info(f"[LiveKit] Connected to room: {self.room_name}")
                self._connected = True

            @self.room.on("disconnected")
            def on_disconnected():
                logger.info(f"[LiveKit] Disconnected from room: {self.room_name}")
                self._connected = False

            await self.room.connect(LIVEKIT_URL, jwt_token)

            # Create video source and track
            self.video_source = rtc.VideoSource(self._width, self._height)
            self.video_track = rtc.LocalVideoTrack.create_video_track(
                "avatar-video",
                self.video_source
            )

            # Publish the track
            options = rtc.TrackPublishOptions()
            options.source = rtc.TrackSource.SOURCE_CAMERA
            await self.room.local_participant.publish_track(self.video_track, options)

            logger.info(f"[LiveKit] Published video track to room: {self.room_name}")
            self._connected = True
            return True

        except Exception as e:
            logger.error(f"[LiveKit] Connection error: {e}")
            self._connected = False
            return False

    async def send_frame(self, frame: np.ndarray):
        """Send a video frame to the room."""
        if not self._connected or self.video_source is None:
            return

        try:
            # Resize frame if needed
            h, w = frame.shape[:2]
            if w != self._width or h != self._height:
                frame = cv2.resize(frame, (self._width, self._height))

            # Convert BGR to RGBA
            if frame.shape[2] == 3:
                frame_rgba = cv2.cvtColor(frame, cv2.COLOR_BGR2RGBA)
            else:
                frame_rgba = frame

            # Create VideoFrame and capture
            video_frame = rtc.VideoFrame(
                self._width,
                self._height,
                rtc.VideoBufferType.RGBA,
                frame_rgba.tobytes()
            )
            self.video_source.capture_frame(video_frame)

        except Exception as e:
            logger.error(f"[LiveKit] Error sending frame: {e}")

    def is_connected(self) -> bool:
        """Check if connected to room."""
        return self._connected

    async def disconnect(self):
        """Disconnect from the room."""
        if self.room:
            await self.room.disconnect()
        self._connected = False
        logger.info(f"[LiveKit] Disconnected from room: {self.room_name}")


class LiveKitManager:
    """Manages LiveKit rooms and video publishers."""

    def __init__(self):
        self.publishers: Dict[str, LiveKitVideoPublisher] = {}
        self._api: Optional[api.LiveKitAPI] = None

    def _get_api(self) -> api.LiveKitAPI:
        """Get or create LiveKit API client."""
        if self._api is None:
            self._api = api.LiveKitAPI(
                LIVEKIT_URL.replace("ws://", "http://").replace("wss://", "https://"),
                LIVEKIT_API_KEY,
                LIVEKIT_API_SECRET
            )
        return self._api

    def generate_viewer_token(self, room_name: str, viewer_identity: str) -> str:
        """Generate a token for a viewer to join the room."""
        token = api.AccessToken(LIVEKIT_API_KEY, LIVEKIT_API_SECRET)
        token.with_identity(viewer_identity)
        token.with_name(f"Viewer {viewer_identity}")
        token.with_grants(api.VideoGrants(
            room_join=True,
            room=room_name,
            can_publish=False,
            can_subscribe=True,
        ))
        return token.to_jwt()

    async def create_room(self, room_name: str) -> bool:
        """Create a LiveKit room."""
        try:
            lk_api = self._get_api()
            await lk_api.room.create_room(
                api.CreateRoomRequest(name=room_name)
            )
            logger.info(f"[LiveKit] Created room: {room_name}")
            return True
        except Exception as e:
            logger.error(f"[LiveKit] Error creating room: {e}")
            return False

    async def get_or_create_publisher(
        self,
        session_id: str,
        fps: int = 30
    ) -> LiveKitVideoPublisher:
        """Get existing publisher or create a new one."""
        if session_id in self.publishers:
            return self.publishers[session_id]

        room_name = f"avatar-{session_id}"
        publisher = LiveKitVideoPublisher(
            room_name=room_name,
            participant_identity=f"avatar-bot-{session_id}",
            fps=fps
        )

        # Create room and connect publisher
        await self.create_room(room_name)
        await publisher.connect()

        self.publishers[session_id] = publisher
        return publisher

    async def handle_offer(
        self,
        session_id: str,
        sdp: str,
        type: str,
        fps: int = 30
    ) -> dict:
        """
        Handle WebRTC offer - with LiveKit, we don't use SDP directly.
        Instead, return room info for the client to connect.
        """
        room_name = f"avatar-{session_id}"

        # Create publisher if not exists
        publisher = await self.get_or_create_publisher(session_id, fps)

        # Generate viewer token for client
        viewer_token = self.generate_viewer_token(room_name, f"viewer-{session_id}")

        return {
            "type": "livekit",
            "room_name": room_name,
            "token": viewer_token,
            "url": LIVEKIT_URL,
            "sdp": ""  # Not used with LiveKit
        }

    def is_connected(self, session_id: str) -> bool:
        """Check if a session has an active publisher."""
        publisher = self.publishers.get(session_id)
        return publisher.is_connected() if publisher else False

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

        start = asyncio.get_event_loop().time()
        while not publisher.is_connected():
            if asyncio.get_event_loop().time() - start > timeout:
                return False
            await asyncio.sleep(0.1)
        return True

    async def send_frame(self, session_id: str, frame: np.ndarray):
        """Send a video frame to a session."""
        publisher = self.publishers.get(session_id)
        if publisher and publisher.is_connected():
            await publisher.send_frame(frame)

    async def close_connection(self, session_id: str):
        """Close a session's connection."""
        publisher = self.publishers.pop(session_id, None)
        if publisher:
            await publisher.disconnect()
            logger.info(f"[LiveKit] Closed session: {session_id}")


# Global manager instance
livekit_manager = LiveKitManager()

# Alias for compatibility with existing code
webrtc_manager = livekit_manager