using System;
using System.Threading.Tasks;
using System.Threading;
using System.Collections;


namespace ConsoleApplication1
{
    public static class TaskUtils
    {
        private class TaskBuilder<T> : IDisposable
        {
            private readonly TaskCompletionSource<T> completionSource;
            private readonly CancellationToken cancellation;
            private volatile IEnumerator stateMachine;
            private volatile int status;
            private volatile Task currentTask;

            public Task<T> Task { get { return this.completionSource.Task; } }
            public TaskStatus Status { get { return (TaskStatus)this.status; } }

            public TaskBuilder(IEnumerator enumerator, object state, TaskCreationOptions creationFlags = TaskCreationOptions.None, CancellationToken cancellationToken = default(CancellationToken))
            {
                if (enumerator == null)
                    throw new ArgumentNullException("enumerator");

                this.stateMachine = enumerator;
                this.completionSource = new TaskCompletionSource<T>(state, creationFlags);
                this.cancellation = cancellationToken;
                this.status = (int)TaskStatus.WaitingForActivation;
            }

            private void Next()
            {
                try
                {
                    cancellation.ThrowIfCancellationRequested();

                    var hasNext = stateMachine.MoveNext();
                    var nextValue = hasNext ? this.stateMachine.Current : null;
                    if (!hasNext)
                    {
                        // yielded break
                        Interlocked.CompareExchange(ref this.status, (int)TaskStatus.RanToCompletion, (int)TaskStatus.Running);
                        completionSource.TrySetResult(default(T));
                    }
                    else if (nextValue is Task)
                    {
                        // yielded task
                        var nextTask = (Task)nextValue;
                        nextTask.ContinueWith((t) =>
                        {
                            Interlocked.CompareExchange(ref this.currentTask, null, t);

                            if (t.IsFaulted)
                            {
                                Interlocked.CompareExchange(ref this.status, (int)TaskStatus.Faulted, (int)TaskStatus.Running);
                                this.completionSource.TrySetException(t.Exception);
                                this.Dispose();
                            }
                            else if (t.IsCanceled)
                            {
                                Interlocked.CompareExchange(ref this.status, (int)TaskStatus.Canceled, (int)TaskStatus.Running);
                                this.completionSource.TrySetCanceled();
                                this.Dispose();
                            }
                            else
                                this.Next();
                        }, CancellationToken.None, TaskContinuationOptions.LongRunning, TaskScheduler.Current);

                        // save next task in root's state, so GC would't collect it
                        this.currentTask = nextTask;
                    }
                    else
                    {
                        // yeilded value
                        Interlocked.CompareExchange(ref this.status, (int)TaskStatus.Faulted, (int)TaskStatus.Running);
                        this.completionSource.TrySetResult((T)nextValue);
                    }
                }
                catch (ThreadAbortException) { throw; }
                catch (OutOfMemoryException) { throw; }
                catch (OperationCanceledException)
                {
                    this.completionSource.TrySetCanceled();
                }
                catch (Exception e)
                {
                    Interlocked.CompareExchange(ref this.status, (int)TaskStatus.Faulted, (int)TaskStatus.Running);
                    this.completionSource.TrySetException(e);
                }

                if (this.status == (int)TaskStatus.RanToCompletion || this.status == (int)TaskStatus.Faulted || this.status == (int)TaskStatus.Canceled)
                    this.Dispose();
            }

            public void Start()
            {
                if (Interlocked.CompareExchange(ref this.status, (int)TaskStatus.Running, (int)TaskStatus.WaitingForActivation) != (int)TaskStatus.WaitingForActivation)
                    return;

                this.currentTask = System.Threading.Tasks.Task.Factory.StartNew(this.Next, this.cancellation);
            }
            public void StartNow()
            {
                if (Interlocked.CompareExchange(ref this.status, (int)TaskStatus.Running, (int)TaskStatus.WaitingForActivation) != (int)TaskStatus.WaitingForActivation)
                    return;

                this.Next();
            }

            #region IDisposable Members

            public void Dispose()
            {
                Interlocked.CompareExchange(ref this.status, (int)TaskStatus.Canceled, (int)TaskStatus.WaitingForActivation);
                Interlocked.CompareExchange(ref this.status, (int)TaskStatus.Canceled, (int)TaskStatus.Running);
                try
                {
                    var s = Interlocked.Exchange(ref this.stateMachine, null);
                    if (s is IDisposable)
                        (s as IDisposable).Dispose();
                }
                catch (ThreadAbortException) { throw; }
                catch (OutOfMemoryException) { throw; }
                catch (Exception e)
                {
                    // Log.Error("An error while disposing builder's state machine", e);
                }

                this.completionSource.TrySetCanceled();
            }

            #endregion

            public override string ToString()
            {
                return string.Format("TaskBuilder, Status={0}, Generator={1}", this.Status, this.stateMachine);
            }
        }

        public static Task FromIterator(Func<object, IEnumerable> iteratorFn, object state, TaskCreationOptions creationFlags = TaskCreationOptions.None, CancellationToken cancellationToken = default(CancellationToken))
        {
            return FromIterator<object>(iteratorFn, state, creationFlags, cancellationToken);
        }
        public static Task FromIterator(IEnumerable iterator, TaskCreationOptions creationFlags = TaskCreationOptions.None, CancellationToken cancellationToken = default(CancellationToken))
        {
            return FromIterator<object>(iterator, null, creationFlags, cancellationToken);
        }
        public static Task<T> FromIterator<T>(Func<object, IEnumerable> iteratorFn, object state, TaskCreationOptions creationFlags = TaskCreationOptions.None, CancellationToken cancellationToken = default(CancellationToken))
        {
            if (iteratorFn == null)
                throw new ArgumentNullException("iteratorFn");

            var iterator = iteratorFn(state);

            return FromIterator<T>(iterator, state, creationFlags, cancellationToken);
        }
        public static Task<T> FromIterator<T>(IEnumerable iterator, object state, TaskCreationOptions creationFlags = TaskCreationOptions.None, CancellationToken cancellationToken = default(CancellationToken))
        {
            if (iterator == null)
                throw new ArgumentNullException("iterator");

            var isLongRunning = (creationFlags & TaskCreationOptions.LongRunning) == TaskCreationOptions.LongRunning;
            if (isLongRunning)
                creationFlags = creationFlags ^ TaskCreationOptions.LongRunning; // remove LongRunning flag


            var enumerator = iterator.GetEnumerator();
            var builder = new TaskBuilder<T>(enumerator, state, creationFlags, cancellationToken);

            // start state machine
            if (isLongRunning)
                builder.Start();
            else
                builder.StartNow();

            // return root task
            return builder.Task;
        }

    }
}