suitelinkclient/src/SuiteLink.Client/SuiteLinkClient.cs

using SuiteLink.Client.Internal;
using SuiteLink.Client.Protocol;
using SuiteLink.Client.Transport;

namespace SuiteLink.Client;

public sealed class SuiteLinkClient : IAsyncDisposable
{
    private readonly ISuiteLinkTransport _transport;
    private readonly bool _ownsTransport;
    private readonly SemaphoreSlim _connectGate = new(1, 1);
    private readonly SemaphoreSlim _operationGate = new(1, 1);
    private readonly SuiteLinkSession _session = new();
    private byte[] _receiveBuffer = new byte[1024];
    private int _receiveCount;
    private int _nextSubscriptionTagId;
    private bool _disposed;

    public SuiteLinkClient()
        : this(new SuiteLinkTcpTransport(), ownsTransport: true)
    {
    }

    public SuiteLinkClient(ISuiteLinkTransport transport, bool ownsTransport = false)
    {
        _transport = transport ?? throw new ArgumentNullException(nameof(transport));
        _ownsTransport = ownsTransport;
    }

    public bool IsConnected =>
        !_disposed &&
        _session.State is SuiteLinkSessionState.SessionConnected or SuiteLinkSessionState.Subscribed;

    public async Task ConnectAsync(SuiteLinkConnectionOptions options, CancellationToken cancellationToken = default)
    {
        ArgumentNullException.ThrowIfNull(options);
        ThrowIfDisposed();

        await _connectGate.WaitAsync(cancellationToken).ConfigureAwait(false);
        try
        {
            ThrowIfDisposed();

            if (IsConnected || _session.State == SuiteLinkSessionState.ConnectSent)
            {
                return;
            }

            if (_session.State == SuiteLinkSessionState.Faulted)
            {
                throw new InvalidOperationException("Client is faulted and cannot be reused.");
            }

            await _transport.ConnectAsync(options.Host, options.Port, cancellationToken).ConfigureAwait(false);
            _session.SetState(SuiteLinkSessionState.TcpConnected);

            var handshakeBytes = SuiteLinkHandshakeCodec.EncodeNormalQueryHandshake(
                options.Application,
                options.ClientNode,
                options.UserName);
            await _transport.SendAsync(handshakeBytes, cancellationToken).ConfigureAwait(false);

            var handshakeAckBytes = await ReceiveSingleFrameAsync(cancellationToken).ConfigureAwait(false);
            _ = SuiteLinkHandshakeCodec.ParseNormalHandshakeAck(handshakeAckBytes);
            _session.SetState(SuiteLinkSessionState.HandshakeComplete);

            var connectBytes = SuiteLinkConnectCodec.Encode(options);
            await _transport.SendAsync(connectBytes, cancellationToken).ConfigureAwait(false);
            // At this stage we've only submitted CONNECT. Do not report ready yet.
            _session.SetState(SuiteLinkSessionState.ConnectSent);
        }
        catch
        {
            try
            {
                _session.SetState(SuiteLinkSessionState.Faulted);
            }
            catch
            {
                // Preserve original exception.
            }

            throw;
        }
        finally
        {
            _connectGate.Release();
        }
    }

    public async Task DisconnectAsync(CancellationToken cancellationToken = default)
    {
        await DisposeCoreAsync(cancellationToken).ConfigureAwait(false);
    }

    public async Task<SubscriptionHandle> SubscribeAsync(
        string itemName,
        Action<SuiteLinkTagUpdate> onUpdate,
        CancellationToken cancellationToken = default)
    {
        ThrowIfDisposed();

        SubscriptionRegistration registration;
        await _operationGate.WaitAsync(cancellationToken).ConfigureAwait(false);
        try
        {
            registration = await SubscribeCoreAsync(itemName, onUpdate, cancellationToken).ConfigureAwait(false);
        }
        finally
        {
            _operationGate.Release();
        }

        DispatchDecodedUpdates(registration.DeferredUpdates);
        return registration.Handle;
    }

    public async Task<SuiteLinkTagUpdate> ReadAsync(
        string itemName,
        TimeSpan timeout,
        CancellationToken cancellationToken = default)
    {
        if (timeout <= TimeSpan.Zero && timeout != Timeout.InfiniteTimeSpan)
        {
            throw new ArgumentOutOfRangeException(nameof(timeout), timeout, "Timeout must be positive or infinite.");
        }

        ThrowIfDisposed();

        using var timeoutCts = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);
        if (timeout != Timeout.InfiniteTimeSpan)
        {
            timeoutCts.CancelAfter(timeout);
        }

        var updateCompletion = new TaskCompletionSource<SuiteLinkTagUpdate>(
            TaskCreationOptions.RunContinuationsAsynchronously);
        SubscriptionHandle? temporaryHandle = null;
        Exception? primaryFailure = null;

        await _operationGate.WaitAsync(timeoutCts.Token).ConfigureAwait(false);
        try
        {
            var registration = await SubscribeCoreAsync(
                itemName,
                update => updateCompletion.TrySetResult(update),
                timeoutCts.Token).ConfigureAwait(false);
            temporaryHandle = registration.Handle;
            DispatchDecodedUpdates(registration.DeferredUpdates);
        }
        finally
        {
            _operationGate.Release();
        }

        try
        {
            while (!updateCompletion.Task.IsCompleted)
            {
                IReadOnlyList<DecodedUpdate> decodedUpdates;
                await _operationGate.WaitAsync(timeoutCts.Token).ConfigureAwait(false);
                try
                {
                    decodedUpdates = await ProcessSingleIncomingFrameAsync(timeoutCts.Token).ConfigureAwait(false);
                }
                finally
                {
                    _operationGate.Release();
                }

                DispatchDecodedUpdates(decodedUpdates);
            }

            return await updateCompletion.Task.ConfigureAwait(false);
        }
        catch (OperationCanceledException) when (!cancellationToken.IsCancellationRequested)
        {
            primaryFailure = new TimeoutException($"No update for '{itemName}' was received within {timeout}.");
            throw primaryFailure;
        }
        catch (Exception ex)
        {
            primaryFailure = ex;
            throw;
        }
        finally
        {
            if (temporaryHandle is not null)
            {
                try
                {
                    await _operationGate.WaitAsync(CancellationToken.None).ConfigureAwait(false);
                    try
                    {
                        await UnsubscribeCoreAsync(temporaryHandle.TagId, CancellationToken.None).ConfigureAwait(false);
                    }
                    finally
                    {
                        _operationGate.Release();
                    }
                }
                catch when (primaryFailure is not null)
                {
                    // Preserve the original read failure.
                }
            }
        }
    }

    public async Task ProcessIncomingAsync(CancellationToken cancellationToken = default)
    {
        ThrowIfDisposed();
        EnsureTagOperationsAllowed();

        IReadOnlyList<DecodedUpdate> decodedUpdates;
        await _operationGate.WaitAsync(cancellationToken).ConfigureAwait(false);
        try
        {
            decodedUpdates = await ProcessSingleIncomingFrameAsync(cancellationToken).ConfigureAwait(false);
        }
        finally
        {
            _operationGate.Release();
        }

        DispatchDecodedUpdates(decodedUpdates);
    }

    public async Task WriteAsync(
        string itemName,
        SuiteLinkValue value,
        CancellationToken cancellationToken = default)
    {
        ArgumentException.ThrowIfNullOrWhiteSpace(itemName);
        ThrowIfDisposed();

        await _operationGate.WaitAsync(cancellationToken).ConfigureAwait(false);
        try
        {
            ThrowIfDisposed();
            EnsureTagOperationsAllowed();

            if (!_session.TryGetTagId(itemName, out var tagId))
            {
                throw new InvalidOperationException(
                    $"Tag '{itemName}' is not subscribed. Subscribe before writing.");
            }

            var pokeBytes = SuiteLinkWriteCodec.Encode(tagId, value);
            await _transport.SendAsync(pokeBytes, cancellationToken).ConfigureAwait(false);
        }
        finally
        {
            _operationGate.Release();
        }
    }

    public async ValueTask DisposeAsync()
    {
        await DisposeCoreAsync(CancellationToken.None).ConfigureAwait(false);
    }

    private async ValueTask<byte[]> ReceiveSingleFrameAsync(CancellationToken cancellationToken)
    {
        while (true)
        {
            if (SuiteLinkFrameReader.TryParseFrame(
                    _receiveBuffer.AsSpan(0, _receiveCount),
                    out _,
                    out var consumed))
            {
                var frameBytes = _receiveBuffer.AsSpan(0, consumed).ToArray();
                var remaining = _receiveCount - consumed;
                if (remaining > 0)
                {
                    _receiveBuffer.AsSpan(consumed, remaining).CopyTo(_receiveBuffer);
                }

                _receiveCount = remaining;
                return frameBytes;
            }

            EnsureReceiveCapacity();
            var bytesRead = await _transport.ReceiveAsync(
                _receiveBuffer.AsMemory(_receiveCount),
                cancellationToken).ConfigureAwait(false);

            if (bytesRead == 0)
            {
                throw new IOException("Remote endpoint closed while waiting for a full frame.");
            }

            _receiveCount += bytesRead;
        }
    }

    private void EnsureReceiveCapacity()
    {
        if (_receiveCount < _receiveBuffer.Length)
        {
            return;
        }

        if (_receiveBuffer.Length >= 1024 * 1024)
        {
            throw new FormatException("Incoming frame exceeds maximum supported size.");
        }

        Array.Resize(ref _receiveBuffer, _receiveBuffer.Length * 2);
    }

    private void ThrowIfDisposed()
    {
        ObjectDisposedException.ThrowIf(_disposed, this);
    }

    private async ValueTask DisposeCoreAsync(CancellationToken cancellationToken)
    {
        var connectGateHeld = false;
        var operationGateHeld = false;

        try
        {
            await _connectGate.WaitAsync(cancellationToken).ConfigureAwait(false);
            connectGateHeld = true;

            await _operationGate.WaitAsync(cancellationToken).ConfigureAwait(false);
            operationGateHeld = true;

            if (_disposed)
            {
                return;
            }

            _disposed = true;
            _session.SetState(SuiteLinkSessionState.Disconnected);
            _receiveCount = 0;
            _receiveBuffer = new byte[1024];

            if (_ownsTransport)
            {
                await _transport.DisposeAsync().ConfigureAwait(false);
            }
        }
        finally
        {
            if (operationGateHeld)
            {
                _operationGate.Release();
            }

            if (connectGateHeld)
            {
                _connectGate.Release();
            }
        }
    }

    private async Task<SubscriptionRegistration> SubscribeCoreAsync(
        string itemName,
        Action<SuiteLinkTagUpdate> onUpdate,
        CancellationToken cancellationToken)
    {
        ArgumentException.ThrowIfNullOrWhiteSpace(itemName);
        ArgumentNullException.ThrowIfNull(onUpdate);

        ThrowIfDisposed();
        EnsureTagOperationsAllowed();

        var requestedTagId = unchecked((uint)Interlocked.Increment(ref _nextSubscriptionTagId));
        var adviseBytes = SuiteLinkSubscriptionCodec.EncodeAdvise(requestedTagId, itemName);
        await _transport.SendAsync(adviseBytes, cancellationToken).ConfigureAwait(false);

        var adviseAckResult = await ReceiveAndCollectUpdatesUntilAsync(
            messageType => messageType == SuiteLinkSubscriptionCodec.AdviseAckMessageType,
            cancellationToken).ConfigureAwait(false);
        var adviseAckBytes = adviseAckResult.FrameBytes;

        var ackItems = SuiteLinkSubscriptionCodec.DecodeAdviseAckMany(adviseAckBytes);
        if (ackItems.Count != 1)
        {
            throw new FormatException(
                $"Expected exactly one advise ACK item for a single subscribe request, but decoded {ackItems.Count}.");
        }

        var acknowledgedTagId = ackItems[0].TagId;
        if (acknowledgedTagId != requestedTagId)
        {
            throw new FormatException(
                $"Advise ACK tag id 0x{acknowledgedTagId:x8} did not match requested tag id 0x{requestedTagId:x8}.");
        }

        _session.RegisterSubscription(itemName, acknowledgedTagId, onUpdate);
        if (_session.State == SuiteLinkSessionState.ConnectSent)
        {
            _session.SetState(SuiteLinkSessionState.SessionConnected);
        }

        if (_session.State == SuiteLinkSessionState.SessionConnected)
        {
            _session.SetState(SuiteLinkSessionState.Subscribed);
        }

        var handle = new SubscriptionHandle(
            itemName,
            acknowledgedTagId,
            () => UnsubscribeAsync(acknowledgedTagId, CancellationToken.None));
        return new SubscriptionRegistration(handle, adviseAckResult.DeferredUpdates);
    }

    private async ValueTask UnsubscribeAsync(uint tagId, CancellationToken cancellationToken)
    {
        if (_disposed)
        {
            return;
        }

        await _operationGate.WaitAsync(cancellationToken).ConfigureAwait(false);
        try
        {
            await UnsubscribeCoreAsync(tagId, cancellationToken).ConfigureAwait(false);
        }
        finally
        {
            _operationGate.Release();
        }
    }

    private async ValueTask UnsubscribeCoreAsync(uint tagId, CancellationToken cancellationToken)
    {
        if (!_session.TryUnregisterByTagId(tagId, out _))
        {
            return;
        }

        var unadviseBytes = SuiteLinkSubscriptionCodec.EncodeUnadvise(tagId);
        await _transport.SendAsync(unadviseBytes, cancellationToken).ConfigureAwait(false);

        if (_session.State == SuiteLinkSessionState.Subscribed && _session.SubscriptionCount == 0)
        {
            _session.SetState(SuiteLinkSessionState.SessionConnected);
        }
    }

    private async Task<FrameReadResult> ReceiveAndCollectUpdatesUntilAsync(
        Func<ushort, bool> messageTypePredicate,
        CancellationToken cancellationToken)
    {
        var deferredUpdates = new List<DecodedUpdate>();
        while (true)
        {
            var frameBytes = await ReceiveSingleFrameAsync(cancellationToken).ConfigureAwait(false);
            var frame = SuiteLinkFrameReader.ParseFrame(frameBytes);

            if (frame.MessageType == SuiteLinkUpdateCodec.UpdateMessageType)
            {
                deferredUpdates.AddRange(SuiteLinkUpdateCodec.DecodeMany(frameBytes));
            }

            if (messageTypePredicate(frame.MessageType))
            {
                return new FrameReadResult(frameBytes, deferredUpdates);
            }
        }
    }

    private async Task<IReadOnlyList<DecodedUpdate>> ProcessSingleIncomingFrameAsync(CancellationToken cancellationToken)
    {
        var frameBytes = await ReceiveSingleFrameAsync(cancellationToken).ConfigureAwait(false);
        var frame = SuiteLinkFrameReader.ParseFrame(frameBytes);
        if (frame.MessageType == SuiteLinkUpdateCodec.UpdateMessageType)
        {
            return SuiteLinkUpdateCodec.DecodeMany(frameBytes);
        }

        return [];
    }

    private void DispatchDecodedUpdates(IReadOnlyList<DecodedUpdate> decodedUpdates)
    {
        if (decodedUpdates.Count == 0)
        {
            return;
        }

        var receivedAtUtc = DateTimeOffset.UtcNow;
        foreach (var decodedUpdate in decodedUpdates)
        {
            _ = _session.TryDispatchUpdate(decodedUpdate, receivedAtUtc, out _, out _);
        }
    }

    private void EnsureTagOperationsAllowed()
    {
        if (_session.State is
            not SuiteLinkSessionState.ConnectSent and
            not SuiteLinkSessionState.SessionConnected and
            not SuiteLinkSessionState.Subscribed)
        {
            throw new InvalidOperationException("Client is not ready for tag operations.");
        }
    }

    private readonly record struct SubscriptionRegistration(
        SubscriptionHandle Handle,
        IReadOnlyList<DecodedUpdate> DeferredUpdates);

    private readonly record struct FrameReadResult(
        byte[] FrameBytes,
        IReadOnlyList<DecodedUpdate> DeferredUpdates);
}