import type { ServerWebSocket } from 'bun'

const connections = new Map<string, Set<ServerWebSocket<{ userId: string }>>>()
const adminSubs = new Set<ServerWebSocket<{ userId: string }>>()
const notifSubs = new Set<ServerWebSocket<{ userId: string }>>()

export function getOnlineUserIds(): string[] {
  return Array.from(connections.keys())
}

function broadcast() {
  const online = getOnlineUserIds()
  const msg = JSON.stringify({ type: 'presence', online })
  for (const ws of adminSubs) ws.send(msg)
}

export function addConnection(
  ws: ServerWebSocket<{ userId: string }>,
  userId: string,
  isAdmin: boolean,
  canReceiveNotifs: boolean,
) {
  let set = connections.get(userId)
  if (!set) {
    set = new Set()
    connections.set(userId, set)
  }
  set.add(ws)
  if (isAdmin) {
    adminSubs.add(ws)
    ws.send(JSON.stringify({ type: 'presence', online: getOnlineUserIds() }))
  }
  if (canReceiveNotifs) notifSubs.add(ws)
  broadcast()
}

export function broadcastToAdmins(message: object) {
  const msg = JSON.stringify(message)
  for (const ws of adminSubs) ws.send(msg)
}

export function broadcastNotification(message: object) {
  const msg = JSON.stringify(message)
  for (const ws of notifSubs) ws.send(msg)
}

export function removeConnection(ws: ServerWebSocket<{ userId: string }>) {
  const userId = ws.data.userId
  const set = connections.get(userId)
  if (set) {
    set.delete(ws)
    if (set.size === 0) connections.delete(userId)
  }
  adminSubs.delete(ws)
  notifSubs.delete(ws)
  broadcast()
}