namespace Tnb.Common.Extension
{
    public static class TaskEx
    {
        public static Task<T> Catch<T, TError>(this Task<T> task, Func<TError, T> onError) where TError : Exception
        {
            var tcs = new TaskCompletionSource<T>();    // #A
            task.ContinueWith(innerTask =>
            {
                if (innerTask.IsFaulted && innerTask?.Exception?.InnerException is TError)
                    tcs.SetResult(onError((TError)innerTask.Exception.InnerException)); // #B
                else if (innerTask.IsCanceled)
                    tcs.SetCanceled();      // #B
                else if (innerTask.IsFaulted)
                    tcs.SetException(innerTask?.Exception?.InnerException ?? throw new InvalidOperationException()); // #B
                else
                    tcs.SetResult(innerTask.Result);  // #B
            });
            return tcs.Task;
        }

        public static async Task<T> Retry<T>(Func<Task<T>> task, int retries,
                                TimeSpan delay, CancellationToken cts = default) =>
            await task().ContinueWith(async innerTask =>
            {
                cts.ThrowIfCancellationRequested();
                if (innerTask.Status != TaskStatus.Faulted)
                    return innerTask.Result;
                if (retries == 0)
                    throw innerTask.Exception ?? throw new Exception();
                await Task.Delay(delay, cts);
                return await Retry(task, retries - 1, delay, cts);
            }).Unwrap();



        public static async Task Catch(this Task task, Func<Exception, Task> exceptionHandler)
        {
            try
            {
                await task;
            }
            catch (Exception ex)
            {
                await exceptionHandler(ex);
            }
        }

    }
}