socket_api.py

from fastapi import APIRouter, WebSocket, Depends, BackgroundTasks, WebSocketDisconnect
from uuid import UUID
from typing import Optional, Dict, Any
import json

from channel.socket_service import SocketService

router = APIRouter(prefix="/ws", tags=["websocket"])

# Singleton instance of SocketService
_socket_service: Optional[SocketService] = None

def get_socket_service() -> SocketService:
    global _socket_service
    if _socket_service is None:
        _socket_service = SocketService()
    return _socket_service

@router.websocket("/{user_id}/{device_id}")
async def websocket_endpoint(
    websocket: WebSocket,
    user_id: UUID,
    device_id: str,
    background_tasks: BackgroundTasks,
    socket_service: SocketService = Depends(get_socket_service)
):
    """
    WebSocket endpoint for real-time communication
    
    Args:
        websocket: The WebSocket connection
        user_id: The ID of the user
        device_id: The ID of the device
    """
    try:
        print(f"New connection request - User: {user_id}, Device: {device_id}")
        # Store background tasks reference
        socket_service.background_tasks = background_tasks
        
        # Start cleanup task if not already running
        background_tasks.add_task(socket_service.cleanup_expired_connections)
        
        # Get initial metadata from client
        metadata = None
        try:
            initial_message = await websocket.receive_json()
            print(f"Initial message received: {initial_message}")
            if isinstance(initial_message, dict):
                metadata = initial_message.get("metadata", {})
        except Exception:
            # If no initial message or invalid, proceed without metadata
            pass
            
        # Connect the websocket
        await socket_service.connect(websocket, user_id, device_id, metadata)
        
        # Handle incoming messages
        while True:
            try:
                message = await websocket.receive_json()
                
                # Handle heartbeat responses
                if isinstance(message, dict) and message.get("type") == "heartbeat_response":
                    await socket_service.update_heartbeat(user_id, device_id)
                    continue
                    
                # For other messages, broadcast to all other devices of this user
                if isinstance(message, dict):
                    # Extract the content from the message
                    if message.get("type") == "message" and "content" in message:
                        # Create a new message with the content
                        broadcast_message = {
                            "type": "message",
                            "content": message["content"]
                        }
                        # Broadcast the content to all other devices of this user
                        broadcast_count = await socket_service.broadcast_message(
                            user_id=user_id,
                            message=broadcast_message,
                            exclude_device_id=device_id  # Exclude the sender's device
                        )
                        
                        # Send acknowledgment to the sender
                        await websocket.send_json({
                            "type": "message_status",
                            "status": "delivered",
                            "broadcast_count": broadcast_count
                        })
                
            except WebSocketDisconnect:
                break
            except Exception as e:
                # Send error message to client
                try:
                    await websocket.send_json({
                        "type": "error",
                        "message": str(e)
                    })
                except Exception:
                    break
                    
    finally:
        # Clean up the connection
        await socket_service.disconnect(user_id, device_id)

@router.post("/{user_id}/broadcast")
async def broadcast_to_user(
    user_id: UUID,
    message: Dict[str, Any],
    sender_device_id: Optional[str] = None,
    socket_service: SocketService = Depends(get_socket_service)
):
    """
    Broadcast a message to all active connections of a user, optionally excluding a sender device
    
    Args:
        user_id: The ID of the user to broadcast to
        message: The message to broadcast
        sender_device_id: Optional device ID to exclude from broadcast
        
    Returns:
        Dict containing broadcast status and count of successful broadcasts
    """
    broadcast_count = await socket_service.broadcast_message(user_id, message, exclude_device_id=sender_device_id)
    return {
        "status": "message_broadcast",
        "broadcast_count": broadcast_count
    }