import type { WebSocket } from "ws";

// Map<lobbyId, Map<userId, WebSocket>>
const connections = new Map<string, Map<string, WebSocket>>();

export const addConnection = (
  lobbyId: string,
  userId: string,
  ws: WebSocket,
): void => {
  if (!connections.has(lobbyId)) {
    connections.set(lobbyId, new Map());
  }
  connections.get(lobbyId)!.set(userId, ws);
};

export const removeConnection = (lobbyId: string, userId: string): void => {
  const lobby = connections.get(lobbyId);
  if (!lobby) return;
  lobby.delete(userId);
  if (lobby.size === 0) {
    connections.delete(lobbyId);
  }
};

export const getConnections = (lobbyId: string): Map<string, WebSocket> => {
  return connections.get(lobbyId) ?? new Map();
};

export const broadcastToLobby = (
  lobbyId: string,
  message: unknown,
  excludeUserId?: string,
): void => {
  const lobby = connections.get(lobbyId);
  if (!lobby) return;
  const payload = JSON.stringify(message);
  for (const [userId, ws] of lobby) {
    if (excludeUserId && userId === excludeUserId) continue;
    if (ws.readyState === ws.OPEN) {
      ws.send(payload);
    }
  }
};