using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using System.Threading.Tasks;

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 Catch(this Task task, Func<Exception,Task> exceptionHandler)
        {
            try
            {
                await task;
            }
            catch (Exception ex)
            {
                await exceptionHandler(ex);
            }
        }

    }
}