using System.Collections.Concurrent;
using System.Net.WebSockets;
using System.Text;

namespace MyApp.WebSockets;

public class WebSocketManager
{
    private readonly ConcurrentDictionary<string, WebSocket> _sockets = new();

    public IReadOnlyDictionary<string, WebSocket> Sockets => _sockets;

    public void AddSocket(string id, WebSocket socket)
    {
        _sockets[id] = socket;
    }

    public bool RemoveSocket(string id)
    {
        return _sockets.TryRemove(id, out _);
    }

    public async Task SendToClientAsync(string id, string message, CancellationToken ct = default)
    {
        if (_sockets.TryGetValue(id, out var socket) && socket.State == WebSocketState.Open)
        {
            var buffer = Encoding.UTF8.GetBytes(message);
            await socket.SendAsync(buffer, WebSocketMessageType.Text, true, ct);
        }
    }

    public async Task BroadcastAsync(string message, CancellationToken ct = default)
    {
        var buffer = Encoding.UTF8.GetBytes(message);

        foreach (var (id, socket) in _sockets.ToArray())
        {
            if (socket.State == WebSocketState.Open)
            {
                await socket.SendAsync(buffer, WebSocketMessageType.Text, true, ct);
            }
            else
            {
                RemoveSocket(id);
            }
        }
    }
}