using System.Threading.Tasks;

namespace SquadDash;

internal sealed class PostedUiActionTracker {
    private readonly object _gate = new();
    private long _postedCount;
    private long _completedCount;
    private TaskCompletionSource<bool>? _drainedTcs;

    public long RegisterPostedAction() {
        lock (_gate) {
            _postedCount++;
            _drainedTcs ??= CreateCompletionSource();
            return _postedCount;
        }
    }

    public void MarkCompleted(long _) {
        TaskCompletionSource<bool>? completion = null;

        lock (_gate) {
            _completedCount++;

            if (_completedCount >= _postedCount && _drainedTcs is not null) {
                completion = _drainedTcs;
                _drainedTcs = null;
            }
        }

        completion?.TrySetResult(true);
    }

    public long PendingCount {
        get {
            lock (_gate)
                return _postedCount - _completedCount;
        }
    }

    public Task WaitForDrainAsync() {
        lock (_gate) {
            if (_completedCount >= _postedCount)
                return Task.CompletedTask;

            _drainedTcs ??= CreateCompletionSource();
            return _drainedTcs.Task;
        }
    }

    private static TaskCompletionSource<bool> CreateCompletionSource() =>
        new(TaskCreationOptions.RunContinuationsAsynchronously);
}
