using System;
using System.Runtime.CompilerServices;
using System.Threading.Tasks;

namespace System.Threading.Tasks
{
    public static class NamedTaskExtensions
    {
        public static Task<T> WithName<T>(this Task<T> task, string name)
        {
            return task;
        }

        public static Task WithName<T>(this Task task, string name)
        {
            return task;
        }
    }
}

namespace JonSkeet.Diagnostics
{
    public static class DiagnosticTaskExtensions
    {
        /// <summary>
        /// Associates a task with a user-specified name before GetAwaiter is called
        /// </summary>
        public static NamedTask<T> WithName<T>(this Task<T> task, string name)
        {
            return new NamedTask<T>(task, name);
        }

        /// <summary>
        /// Gets a diagnostic awaiter for a task, based only on its ID.
        /// </summary>
        public static NamedAwaiter<T> GetAwaiter<T>(this Task<T> task)
        {
            return new NamedTask<T>(task, "[" + task.Id + "]").GetAwaiter();
        }

        public struct NamedTask<T>
        {
            private readonly Task<T> task;
            private readonly string name;

            public NamedTask(Task<T> task, string name)
            {
                this.task = task;
                this.name = name;
            }

            public NamedAwaiter<T> GetAwaiter()
            {
                Console.WriteLine("GetAwaiter called for task \"{0}\"", name);
                return new NamedAwaiter<T>(AsyncCtpThreadingExtensions.GetAwaiter(task), name);
            }
        }

        public struct NamedAwaiter<T>
        {
            private readonly TaskAwaiter<T> awaiter;
            private readonly string name;

            public NamedAwaiter(TaskAwaiter<T> awaiter, string name)
            {
                this.awaiter = awaiter;
                this.name = name;
            }

            public bool BeginAwait(Action continuation)
            {
                Console.WriteLine("BeginAwait called for task \"{0}\"...", name);
                bool ret = awaiter.BeginAwait(continuation);
                Console.WriteLine("... BeginAwait for task \"{0}\" returning {1}", name, ret);
                return ret;
            }

            public T EndAwait()
            {
                Console.WriteLine("EndAwait called for task \"{0}\"", name);
                // We could potentially report the result here
                return awaiter.EndAwait();
            }
        }

        /// <summary>
        /// Associates a task with a user-specified name before GetAwaiter is called
        /// </summary>
        public static NamedTask WithName(this Task task, string name)
        {
            return new NamedTask(task, name);
        }

        /// <summary>
        /// Gets a diagnostic awaiter for a task, based only on its ID.
        /// </summary>
        public static NamedAwaiter GetAwaiter(this Task task)
        {
            return new NamedTask(task, "[" + task.Id + "]").GetAwaiter();
        }

        public struct NamedTask
        {
            private readonly Task task;
            private readonly string name;

            public NamedTask(Task task, string name)
            {
                this.task = task;
                this.name = name;
            }

            public NamedAwaiter GetAwaiter()
            {
                Console.WriteLine("GetAwaiter called for task \"{0}\"", name);
                return new NamedAwaiter(AsyncCtpThreadingExtensions.GetAwaiter(task), name);
            }
        }

        public struct NamedAwaiter
        {
            private readonly TaskAwaiter awaiter;
            private readonly string name;

            public NamedAwaiter(TaskAwaiter awaiter, string name)
            {
                this.awaiter = awaiter;
                this.name = name;
            }

            public bool BeginAwait(Action continuation)
            {
                Console.WriteLine("BeginAwait called for task \"{0}\"...", name);
                bool ret = awaiter.BeginAwait(continuation);
                Console.WriteLine("... BeginAwait for task \"{0}\" returning {1}", name, ret);
                return ret;
            }

            public void EndAwait()
            {
                Console.WriteLine("EndAwait called for task \"{0}\"", name);
                // We could potentially report the result here
                awaiter.EndAwait();
            }
        }
    }
}