COPILOT/copilot-coordinator/Services/WsBroadcaster.cs

using System.Collections.Concurrent;
using System.Net.WebSockets;
using System.Text;
using System.Text.Json;
using CopilotCoordinator.Models;

namespace CopilotCoordinator.Services;

/// <summary>
/// WebSocket broadcaster. Manages connected keyboard clients and sends events.
/// </summary>
public class WsBroadcaster
{
    private readonly ConcurrentDictionary<string, (WebSocket Socket, string KeyboardId)> _clients = new();
    private readonly ILogger<WsBroadcaster> _logger;

    public WsBroadcaster(ILogger<WsBroadcaster> logger)
    {
        _logger = logger;
    }

    public IReadOnlyList<KeyboardInfo> GetConnectedKeyboards()
    {
        return _clients.Values
            .Select(c => new KeyboardInfo(c.KeyboardId, null, DateTime.UtcNow))
            .ToList();
    }

    public async Task HandleConnection(WebSocket ws, string clientId, string keyboardId)
    {
        _clients[clientId] = (ws, keyboardId);
        _logger.LogInformation("Keyboard {KeyboardId} connected (client {ClientId})", keyboardId, clientId);

        await Broadcast("keyboard_online", new { keyboardId });

        try
        {
            var buffer = new byte[4096];
            while (ws.State == WebSocketState.Open)
            {
                var result = await ws.ReceiveAsync(buffer, CancellationToken.None);
                if (result.MessageType == WebSocketMessageType.Close)
                    break;

                // Parse incoming messages (keyboard can send commands via WebSocket too)
                if (result.MessageType == WebSocketMessageType.Text)
                {
                    var message = Encoding.UTF8.GetString(buffer, 0, result.Count);
                    _logger.LogDebug("Received from {KeyboardId}: {Message}", keyboardId, message);
                }
            }
        }
        catch (WebSocketException ex)
        {
            _logger.LogWarning("WebSocket error for {KeyboardId}: {Message}", keyboardId, ex.Message);
        }
        finally
        {
            _clients.TryRemove(clientId, out _);
            _logger.LogInformation("Keyboard {KeyboardId} disconnected", keyboardId);
            await Broadcast("keyboard_offline", new { keyboardId });
        }
    }

    /// <summary>
    /// Broadcast a message to all connected keyboards.
    /// </summary>
    public async Task Broadcast(string type, object? data = null)
    {
        var json = WsMessage.Serialize(type, data);
        var bytes = Encoding.UTF8.GetBytes(json);

        var deadClients = new List<string>();

        foreach (var (clientId, (socket, _)) in _clients)
        {
            try
            {
                if (socket.State == WebSocketState.Open)
                {
                    await socket.SendAsync(bytes, WebSocketMessageType.Text, true, CancellationToken.None);
                }
                else
                {
                    deadClients.Add(clientId);
                }
            }
            catch
            {
                deadClients.Add(clientId);
            }
        }

        foreach (var id in deadClients)
            _clients.TryRemove(id, out _);
    }

    /// <summary>
    /// Send a message to a specific keyboard.
    /// </summary>
    public async Task SendTo(string keyboardId, string type, object? data = null)
    {
        var json = WsMessage.Serialize(type, data);
        var bytes = Encoding.UTF8.GetBytes(json);

        foreach (var (clientId, (socket, kbId)) in _clients)
        {
            if (kbId.Equals(keyboardId, StringComparison.OrdinalIgnoreCase) &&
                socket.State == WebSocketState.Open)
            {
                try
                {
                    await socket.SendAsync(bytes, WebSocketMessageType.Text, true, CancellationToken.None);
                }
                catch (Exception ex)
                {
                    _logger.LogWarning("Failed to send to {KeyboardId}: {Error}", keyboardId, ex.Message);
                }
            }
        }
    }
}