namespace Azaion.Common.Extensions;

public static class CancellationTokenExtensions
{
    public static void WaitForCancel(this CancellationToken token, TimeSpan timeout)
    {
        try
        {
            Task.Delay(timeout, token).Wait(token);
        }
        catch (OperationCanceledException)
        {
            //Don't need to catch exception, need only return from the waiting
        }
    }

    public static Task AsTask(this CancellationToken cancellationToken)
    {
        if (!cancellationToken.CanBeCanceled)
            return new TaskCompletionSource<bool>().Task;

        if (cancellationToken.IsCancellationRequested)
            return Task.FromCanceled(cancellationToken);

        var tcs = new TaskCompletionSource<bool>(TaskCreationOptions.RunContinuationsAsynchronously);
        var registration = cancellationToken.Register(() => tcs.TrySetResult(true));
        tcs.Task.ContinueWith(_ => registration.Dispose(), TaskScheduler.Default);
        return tcs.Task;
    }
}