type ConnectionPool* = ref object
    connections* : seq[WebSocket]
    lock*        : Lock

proc newConnectionPool*(): ConnectionPool =
    result = ConnectionPool(connections: @[])
    result.lock.initLock()

var socket_pool = initTable[string, var ConnectionPool]()

proc removeConnection(session_id: string, closed_socket: WebSocket) {.gcsafe.} =
    icb "Removing socket from " & session_id & " connection pool"
    var
        num_connections_before : int
        num_connections_after  : int
    
    {.cast(gcsafe).}: 
        num_connections_before = socket_pool[session_id].connections.len
        icb "Num connections before removal: " & $num_connections_before

        withLock socket_pool[session_id].lock:
            socket_pool[session_id].connections = 
                socket_pool[session_id].connections
                    .filterit( it != closed_socket )

        num_connections_after = socket_pool[session_id].connections.len
        icb "Num connections after removal: " & $num_connections_after

        if num_connections_after != num_connections_before - 1:
            echo "Error removing socket from connection pool"
            echo "Num connections before removal: " & $num_connections_before
            echo "Num connections after removal: " & $num_connections_after

proc broadcastMessage(session_id: string, message: string) {.async.} =
    var connectionsToSend, failedConnections: seq[WebSocket] = @[]

    # Extract the relevant WebSocket connections while holding the lock
    {.cast(gcsafe).}:
        withLock socket_pool[session_id].lock:
            for ws in socket_pool[session_id].connections:
                if ws.readyState == Open:
                    connectionsToSend.add(ws)

    icb "Num sockets in connection pool for broadcast: " & $connectionsToSend.len

    # Send messages and handle exceptions
    for ws in connectionsToSend:
        try:
            await ws.send(message)
        except:
            icr "Error sending message to websocket; removing from pool."
            failedConnections.add(ws)

    # Remove failed connections
    for ws in failedConnections:
        removeConnection(session_id, ws)
        try:
            ws.close()
        except:
            discard