from fastapi import APIRouter, WebSocket, WebSocketDisconnect, Query
from typing import Dict
import json

router = APIRouter(tags=["websocket"])

# Store active connections
active_connections: Dict[int, WebSocket] = {}


@router.websocket("/api/v1/ws")
async def websocket_endpoint(websocket: WebSocket, token: str = Query(...)):
    await websocket.accept()

    # In a real implementation, you would validate the token here
    # For MVP, we'll accept all connections
    user_id = 1  # Placeholder

    active_connections[user_id] = websocket

    try:
        while True:
            data = await websocket.receive_text()
            # Handle incoming messages if needed
    except WebSocketDisconnect:
        if user_id in active_connections:
            del active_connections[user_id]


async def broadcast_event(event_type: str, data: dict, user_ids: list[int] = None):
    """Broadcast an event to specific users or all connected users"""
    message = json.dumps({
        "type": event_type,
        "data": data
    })

    if user_ids:
        for user_id in user_ids:
            if user_id in active_connections:
                await active_connections[user_id].send_text(message)
    else:
        for connection in active_connections.values():
            await connection.send_text(message)
