LumaTunnel

using System.Collections.Concurrent;
using System.Net.WebSockets;
using System.Threading.Channels;
using LumaTunnel.Client.Core.Abstractions;
using LumaTunnel.Shared.Models;
using LumaTunnel.Shared.Protocol;

namespace LumaTunnel.Client.Core.Tunnel;

public sealed class TunnelService : ITunnelService, IProxyConnector
{
    private readonly ConcurrentDictionary<uint, ClientTunnelState> _streams = new();
    private readonly SemaphoreSlim _lifecycleGate = new(1, 1);
    private readonly Channel<TunnelFrame> _outgoing = Channel.CreateBounded<TunnelFrame>(new BoundedChannelOptions(512)
    {
        FullMode = BoundedChannelFullMode.Wait,
        SingleReader = true,
        SingleWriter = false
    });
    private CancellationTokenSource? _sessionCancellation;
    private ClientWebSocket? _webSocket;
    private Task? _receiveTask;
    private Task? _sendTask;
    private Task? _pingTask;
    private int _nextStreamId = -1;

    public bool IsConnected => _webSocket?.State == WebSocketState.Open;
    public event EventHandler<bool>? ConnectionStateChanged;

    public async Task ConnectAsync(TunnelCredentials credentials, CancellationToken cancellationToken = default)
    {
        ArgumentNullException.ThrowIfNull(credentials);
        await _lifecycleGate.WaitAsync(cancellationToken).ConfigureAwait(false);
        try
        {
            await DisconnectCoreAsync(CancellationToken.None).ConfigureAwait(false);
            while (_outgoing.Reader.TryRead(out _))
            {
            }
            var socket = new ClientWebSocket();
            socket.Options.UseDefaultCredentials = false;
            socket.Options.Proxy = null;
            socket.Options.AddSubProtocol(TunnelProtocol.WebSocketSubProtocol);
            socket.Options.SetRequestHeader("Authorization", $"Bearer {credentials.DeviceToken}");
            socket.Options.SetRequestHeader("X-Luma-Device-Id", credentials.DeviceId);
            socket.Options.SetRequestHeader("X-Luma-Client-Version", credentials.ClientVersion);

            var sessionCancellation = new CancellationTokenSource();
            using var linked = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken, sessionCancellation.Token);
            await socket.ConnectAsync(credentials.TunnelUri, linked.Token).ConfigureAwait(false);
            // CyberComm currently exposes the requested header to the original service but may
            // not echo a selected subprotocol. Reject an explicit mismatch while remaining
            // compatible with that HTTP.sys behavior.
            if (socket.SubProtocol is not null
                && !string.Equals(socket.SubProtocol, TunnelProtocol.WebSocketSubProtocol, StringComparison.Ordinal))
            {
                socket.Abort();
                socket.Dispose();
                throw new TunnelProtocolException(TunnelErrorCode.UnsupportedVersion, "The server did not accept the LumaTunnel v1 subprotocol.");
            }

            _webSocket = socket;
            _sessionCancellation = sessionCancellation;
            _receiveTask = ReceiveLoopAsync(socket, sessionCancellation.Token);
            _sendTask = SendLoopAsync(socket, sessionCancellation.Token);
            _pingTask = PingLoopAsync(sessionCancellation.Token);
            ConnectionStateChanged?.Invoke(this, true);
        }
        finally
        {
            _lifecycleGate.Release();
        }
    }

    public async ValueTask<Stream> OpenStreamAsync(ProxyTarget target, CancellationToken cancellationToken = default)
    {
        var socket = _webSocket;
        var sessionCancellation = _sessionCancellation;
        if (socket?.State != WebSocketState.Open || sessionCancellation is null)
            throw new InvalidOperationException("The tunnel is not connected.");

        var streamId = unchecked((uint)Interlocked.Add(ref _nextStreamId, 2));
        if (streamId == 0)
            streamId = unchecked((uint)Interlocked.Add(ref _nextStreamId, 2));
        var state = new ClientTunnelState(streamId);
        if (!_streams.TryAdd(streamId, state))
            throw new InvalidOperationException("Unable to allocate a tunnel stream identifier.");

        using var timeout = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken, sessionCancellation.Token);
        timeout.CancelAfter(TimeSpan.FromSeconds(10));
        try
        {
            await QueueFrameAsync(new TunnelFrame(TunnelFrameType.Open, streamId, TunnelMessageSerializer.Serialize(new TunnelOpenRequest(target.Host, target.Port))), timeout.Token).ConfigureAwait(false);
            await state.Opened.Task.WaitAsync(timeout.Token).ConfigureAwait(false);
            return new TunnelClientStream(this, state);
        }
        catch
        {
            _streams.TryRemove(streamId, out _);
            await QueueFrameBestEffortAsync(TunnelFrame.Empty(TunnelFrameType.Reset, streamId)).ConfigureAwait(false);
            state.Fail(new IOException("The tunnel stream could not be opened."));
            throw;
        }
    }

    ValueTask<Stream> IProxyConnector.ConnectAsync(ProxyTarget target, CancellationToken cancellationToken) =>
        OpenStreamAsync(target, cancellationToken);

    public async Task DisconnectAsync(CancellationToken cancellationToken = default)
    {
        await _lifecycleGate.WaitAsync(cancellationToken).ConfigureAwait(false);
        try
        {
            await DisconnectCoreAsync(cancellationToken).ConfigureAwait(false);
        }
        finally
        {
            _lifecycleGate.Release();
        }
    }

    internal ValueTask SendDataAsync(uint streamId, ReadOnlyMemory<byte> data, CancellationToken cancellationToken)
    {
        if (data.Length > TunnelProtocol.MaxPayloadLength)
            throw new ArgumentOutOfRangeException(nameof(data));
        return QueueFrameAsync(new TunnelFrame(TunnelFrameType.Data, streamId, data.ToArray()), cancellationToken);
    }

    internal ValueTask SendHalfCloseAsync(uint streamId, CancellationToken cancellationToken) =>
        QueueFrameAsync(TunnelFrame.Empty(TunnelFrameType.HalfClose, streamId), cancellationToken);

    internal async ValueTask CloseStreamAsync(uint streamId)
    {
        if (_streams.TryRemove(streamId, out var state))
        {
            state.Complete();
            await QueueFrameBestEffortAsync(TunnelFrame.Empty(TunnelFrameType.Close, streamId)).ConfigureAwait(false);
        }
    }

    public async ValueTask DisposeAsync()
    {
        await DisconnectAsync().ConfigureAwait(false);
        _lifecycleGate.Dispose();
        GC.SuppressFinalize(this);
    }

    private async Task ReceiveLoopAsync(ClientWebSocket socket, CancellationToken cancellationToken)
    {
        var buffer = new byte[TunnelProtocol.HeaderLength + TunnelProtocol.MaxPayloadLength];
        try
        {
            while (!cancellationToken.IsCancellationRequested && socket.State == WebSocketState.Open)
            {
                var length = 0;
                ValueWebSocketReceiveResult result;
                do
                {
                    if (length == buffer.Length)
                        throw new TunnelProtocolException(TunnelErrorCode.FrameTooLarge, "The server sent an oversized WebSocket message.");
                    result = await socket.ReceiveAsync(buffer.AsMemory(length), cancellationToken).ConfigureAwait(false);
                    if (result.MessageType == WebSocketMessageType.Close)
                        return;
                    if (result.MessageType != WebSocketMessageType.Binary)
                        throw new TunnelProtocolException(TunnelErrorCode.MalformedFrame, "Tunnel messages must be binary.");
                    length += result.Count;
                }
                while (!result.EndOfMessage);

                HandleFrame(TunnelProtocol.Decode(buffer.AsSpan(0, length)));
            }
        }
        catch (OperationCanceledException) when (cancellationToken.IsCancellationRequested)
        {
        }
        catch (Exception exception)
        {
            FailAll(exception);
        }
        finally
        {
            _sessionCancellation?.Cancel();
            ConnectionStateChanged?.Invoke(this, false);
        }
    }

    private async Task SendLoopAsync(ClientWebSocket socket, CancellationToken cancellationToken)
    {
        try
        {
            await foreach (var frame in _outgoing.Reader.ReadAllAsync(cancellationToken).ConfigureAwait(false))
            {
                var bytes = TunnelProtocol.Encode(frame);
                await socket.SendAsync(bytes, WebSocketMessageType.Binary, true, cancellationToken).ConfigureAwait(false);
            }
        }
        catch (OperationCanceledException) when (cancellationToken.IsCancellationRequested)
        {
        }
        catch (Exception exception)
        {
            FailAll(exception);
            _sessionCancellation?.Cancel();
        }
    }

    private async Task PingLoopAsync(CancellationToken cancellationToken)
    {
        try
        {
            while (!cancellationToken.IsCancellationRequested)
            {
                await Task.Delay(TimeSpan.FromSeconds(30), cancellationToken).ConfigureAwait(false);
                await QueueFrameAsync(new TunnelFrame(TunnelFrameType.Ping, 0, BitConverter.GetBytes(DateTimeOffset.UtcNow.ToUnixTimeMilliseconds())), cancellationToken).ConfigureAwait(false);
            }
        }
        catch (OperationCanceledException) when (cancellationToken.IsCancellationRequested)
        {
        }
    }

    private void HandleFrame(TunnelFrame frame)
    {
        if (frame.Type == TunnelFrameType.Ping)
        {
            _ = QueueFrameBestEffortAsync(new TunnelFrame(TunnelFrameType.Pong, 0, frame.Payload)).AsTask();
            return;
        }
        if (frame.Type == TunnelFrameType.Pong)
            return;
        if (!_streams.TryGetValue(frame.StreamId, out var state))
            return;

        switch (frame.Type)
        {
            case TunnelFrameType.OpenOk:
                state.Opened.TrySetResult();
                break;
            case TunnelFrameType.OpenError:
                var error = TunnelMessageSerializer.Deserialize<TunnelOpenError>(frame.Payload);
                state.Fail(new TunnelProtocolException(error.Code, error.Message));
                _streams.TryRemove(frame.StreamId, out _);
                break;
            case TunnelFrameType.Data:
                if (!state.Incoming.Writer.TryWrite(frame.Payload.ToArray()))
                {
                    state.Fail(new IOException("The local tunnel stream receive buffer is full."));
                    _streams.TryRemove(frame.StreamId, out _);
                    _ = QueueFrameBestEffortAsync(TunnelFrame.Empty(TunnelFrameType.Reset, frame.StreamId)).AsTask();
                }
                break;
            case TunnelFrameType.HalfClose:
                state.CompleteIncoming();
                break;
            case TunnelFrameType.Close:
                state.Complete();
                _streams.TryRemove(frame.StreamId, out _);
                break;
            case TunnelFrameType.Reset:
                state.Fail(new IOException("The server reset the tunnel stream."));
                _streams.TryRemove(frame.StreamId, out _);
                break;
        }
    }

    private ValueTask QueueFrameAsync(TunnelFrame frame, CancellationToken cancellationToken) =>
        _outgoing.Writer.WriteAsync(frame, cancellationToken);

    private async ValueTask QueueFrameBestEffortAsync(TunnelFrame frame)
    {
        try
        {
            if (!_outgoing.Writer.TryWrite(frame))
            {
                using var timeout = new CancellationTokenSource(TimeSpan.FromSeconds(1));
                await _outgoing.Writer.WriteAsync(frame, timeout.Token).ConfigureAwait(false);
            }
        }
        catch (OperationCanceledException)
        {
        }
    }

    private async Task DisconnectCoreAsync(CancellationToken cancellationToken)
    {
        var socket = _webSocket;
        var sessionCancellation = _sessionCancellation;
        _webSocket = null;
        _sessionCancellation = null;
        if (socket is null)
            return;

        sessionCancellation?.Cancel();
        try
        {
            if (socket.State is WebSocketState.Open or WebSocketState.CloseReceived)
                await socket.CloseAsync(WebSocketCloseStatus.NormalClosure, "Client is disconnecting.", cancellationToken).ConfigureAwait(false);
        }
        catch (WebSocketException)
        {
            socket.Abort();
        }

        var tasks = new[] { _receiveTask, _sendTask, _pingTask }.Where(static task => task is not null).Cast<Task>().ToArray();
        if (tasks.Length > 0)
            await Task.WhenAny(Task.WhenAll(tasks), Task.Delay(TimeSpan.FromSeconds(2), cancellationToken)).ConfigureAwait(false);
        socket.Dispose();
        sessionCancellation?.Dispose();
        FailAll(new IOException("The tunnel session was closed."));
        ConnectionStateChanged?.Invoke(this, false);
    }

    private void FailAll(Exception exception)
    {
        foreach (var pair in _streams.ToArray())
        {
            if (_streams.TryRemove(pair.Key, out var state))
                state.Fail(exception);
        }
    }
}