using System.Net.WebSockets;
using System.Text;

namespace MeetingAssistant.Transcription;

public sealed class ClientFunAsrWebSocketConnectionFactory : IFunAsrWebSocketConnectionFactory
{
    public async Task<IFunAsrWebSocketConnection> ConnectAsync(Uri endpoint, CancellationToken cancellationToken)
    {
        var socket = new ClientWebSocket();
        await socket.ConnectAsync(endpoint, cancellationToken);
        return new ClientFunAsrWebSocketConnection(socket);
    }
}

public sealed class ClientFunAsrWebSocketConnection : IFunAsrWebSocketConnection
{
    private readonly ClientWebSocket socket;

    public ClientFunAsrWebSocketConnection(ClientWebSocket socket)
    {
        this.socket = socket;
    }

    public Task SendTextAsync(string message, CancellationToken cancellationToken)
    {
        return socket.SendAsync(
            Encoding.UTF8.GetBytes(message),
            WebSocketMessageType.Text,
            endOfMessage: true,
            cancellationToken);
    }

    public async Task SendBinaryAsync(ReadOnlyMemory<byte> message, CancellationToken cancellationToken)
    {
        await socket.SendAsync(message, WebSocketMessageType.Binary, endOfMessage: true, cancellationToken);
    }

    public async Task<string?> ReceiveTextAsync(CancellationToken cancellationToken)
    {
        var buffer = new byte[8192];
        using var message = new MemoryStream();

        while (true)
        {
            var result = await socket.ReceiveAsync(buffer, cancellationToken);
            if (result.MessageType == WebSocketMessageType.Close)
            {
                return null;
            }

            if (result.MessageType != WebSocketMessageType.Text)
            {
                continue;
            }

            message.Write(buffer, 0, result.Count);
            if (result.EndOfMessage)
            {
                return Encoding.UTF8.GetString(message.ToArray());
            }
        }
    }

    public async Task CloseOutputAsync(CancellationToken cancellationToken)
    {
        if (socket.State is WebSocketState.Open or WebSocketState.CloseReceived)
        {
            await socket.CloseOutputAsync(WebSocketCloseStatus.NormalClosure, "meeting assistant finished sending audio", cancellationToken);
        }
    }

    public ValueTask DisposeAsync()
    {
        socket.Dispose();
        return ValueTask.CompletedTask;
    }
}