using System.Net.WebSockets;
using MyApp.WebSockets;

namespace MyApp.WebSockets;

public class WebSocketMiddleware
{
    private readonly RequestDelegate _next;
    private readonly WebSocketManager _manager;

    public WebSocketMiddleware(RequestDelegate next, WebSocketManager manager)
    {
        _next = next;
        _manager = manager;
    }

    public async Task InvokeAsync(HttpContext context)
    {
        if (context.Request.Path == "/ws")
        {
            if (!context.WebSockets.IsWebSocketRequest)
            {
                context.Response.StatusCode = StatusCodes.Status400BadRequest;
                return;
            }

            // Let the client pass a stable id:  ws://host/ws?clientId=abc123
            var clientId = context.Request.Query["clientId"].ToString();
            if (string.IsNullOrWhiteSpace(clientId))
            {
                clientId = Guid.NewGuid().ToString("N");
            }

            using var socket = await context.WebSockets.AcceptWebSocketAsync();
            _manager.AddSocket(clientId, socket);

            await ReceiveLoop(clientId, socket);
            _manager.RemoveSocket(clientId);
            return;
        }

        await _next(context);
    }

    private static async Task ReceiveLoop(string id, WebSocket socket)
    {
        var buffer = new byte[4 * 1024];

        while (socket.State == WebSocketState.Open)
        {
            var result = await socket.ReceiveAsync(buffer, CancellationToken.None);

            if (result.MessageType == WebSocketMessageType.Close)
            {
                await socket.CloseAsync(WebSocketCloseStatus.NormalClosure, "Closing", CancellationToken.None);
                break;
            }

            // (Optional) handle incoming messages here if you need two-way comms.
        }
    }
}