using Sony.Filtr.Functional;
using System;
using System.Collections.Generic;
using System.Threading.Tasks;
using System.Threading.Tasks.Dataflow;

namespace Sony.Filtr.Utility.Extensions
{
    public static class TPLDataflowExtensions
    {
        public static BatchBlock<TOutput> Batch<TOutput>(this ISourceBlock<TOutput> sourceBlock, int batchSize, DataflowLinkOptions linkOptions)
        {
            var batchBlock = new BatchBlock<TOutput>(batchSize);

            sourceBlock.LinkTo(batchBlock, linkOptions);

            return batchBlock;
        }

        public static TransformManyBlock<TInput, TOutput> AsTransformManyBlock<TInput, TOutput>(this Func<TInput, Task<IEnumerable<TOutput>>> func, ExecutionDataflowBlockOptions blockOptions)
        {
            return new TransformManyBlock<TInput, TOutput>(func, blockOptions);
        }

        public static TransformManyBlock<TInput, Result<TOutput>> AsTransformManyBlock<TInput, TOutput>(this Func<TInput, Task<IEnumerable<Result<TOutput>>>> func, ExecutionDataflowBlockOptions blockOptions)
        {
            return new TransformManyBlock<TInput, Result<TOutput>>(func, blockOptions);
        }

        public static TransformBlock<TInput, TOutput> AsTransformBlock<TInput, TOutput>(this Func<TInput, Task<TOutput>> func, ExecutionDataflowBlockOptions blockOptions)
        {
            return new TransformBlock<TInput, TOutput>(func, blockOptions);
        }

        public static ActionBlock<TInput> AsActionBlock<TInput>(this Func<TInput, Task<Unit>> func, ExecutionDataflowBlockOptions blockOptions)
        {
            return new ActionBlock<TInput>(func, blockOptions);
        }

        public static ActionBlock<TInput> AsActionBlock<TInput>(this Func<TInput, Task<Result<Unit>>> func, ExecutionDataflowBlockOptions blockOptions)
        {
            return new ActionBlock<TInput>(func, blockOptions);
        }

        public static TransformBlock<TSourceBlockOutput, TOutput> ContinueWith<TSourceBlockOutput, TOutput>(
            this ISourceBlock<TSourceBlockOutput> sourceBlock,
            Func<TSourceBlockOutput, Task<TOutput>> func,
            ExecutionDataflowBlockOptions blockOptions,
            DataflowLinkOptions linkOptions)
        {
            var transformBlock = new TransformBlock<TSourceBlockOutput, TOutput>(func, blockOptions);

            sourceBlock.LinkTo(transformBlock, linkOptions);

            return transformBlock;
        }

        public static ActionBlock<TSourceBlockOutput> ContinueWith<TSourceBlockOutput>(
            this ISourceBlock<TSourceBlockOutput> sourceBlock,
            Func<TSourceBlockOutput, Task> func,
            ExecutionDataflowBlockOptions blockOptions,
            DataflowLinkOptions linkOptions)
        {
            var actionBlock = new ActionBlock<TSourceBlockOutput>(func, blockOptions);

            sourceBlock.LinkTo(actionBlock, linkOptions);

            return actionBlock;
        }
    }
}