mcdiscordbot/MinecraftDiscordBot/ChunkWaiter.cs

namespace MinecraftDiscordBot;

public class ChunkWaiter<T> : IChunkWaiter {
    public int ID { get; }
    private readonly CancellationToken _ct;
    public ChunkWaiter(int id, Func<string, T> resultParser, CancellationToken ct) {
        ID = id;
        this.resultParser = resultParser;
        _ct = ct;
    }
    private readonly TaskCompletionSource<T> tcs = new();
    private readonly Func<string, T> resultParser;
    public Task<T> Task => tcs.Task.WaitAsync(_ct);
    public bool Finished { get; private set; } = false;
    public bool IsCancellationRequested => _ct.IsCancellationRequested;
    private string?[]? _chunks = null;
    private int _receivedChunks = 0;
    private bool _success = true;
    private readonly object _syncRoot = new();
    public void AddChunk(int chunkId, int totalChunks, string value) {
        lock (_syncRoot) {
            if (_chunks is null) _chunks = new string[totalChunks];
            else if (_chunks.Length != totalChunks) {
                Program.LogErrorAsync(Program.WebSocketSource, new InvalidOperationException("Different numbers of chunks in same message ID!"));
                return;
            }
            ref string? chunk = ref _chunks[chunkId - 1]; // Lua 1-indexed
            if (chunk is not null) {
                Program.LogErrorAsync(Program.WebSocketSource, new InvalidOperationException($"Chunk with ID {chunkId} was already received!"));
                return;
            }
            chunk = value;
        }
        if (++_receivedChunks == totalChunks) FinalizeResult(_chunks);
    }
    private void FinalizeResult(string?[] _chunks) {
        var resultString = string.Concat(_chunks);
        if (_success) tcs.SetResult(resultParser(resultString));
        else tcs.SetException(new ReplyException(resultString));
        Finished = true;
    }
    public void SetUnsuccessful() => _success = false;
}