//------------------------------------------------------------------------------
// <copyright file="WithinCancellableCallbackTaskAwaitable.cs" company="Microsoft">
//     Copyright (c) Microsoft Corporation.  All rights reserved.
// </copyright>
//------------------------------------------------------------------------------

namespace System.Web.Util {
    using System;
    using System.Runtime.CompilerServices;
    using System.Security.Permissions;
    using System.Threading;
    using System.Threading.Tasks;

    // An awaitable type that invokes a continuation callback under a call to HttpContext.InvokeCancellableCallback.
    // This type is for compiler use; the ASP.NET runtime is not expected to call into these APIs directly.
    internal struct WithinCancellableCallbackTaskAwaitable {

        internal static readonly WithinCancellableCallbackTaskAwaitable Completed = new WithinCancellableCallbackTaskAwaitable(null, ((Task)Task.FromResult((object)null)).GetAwaiter());

        private readonly WithinCancellableCallbackTaskAwaiter _awaiter;

        public WithinCancellableCallbackTaskAwaitable(HttpContext context, TaskAwaiter innerAwaiter) {
            _awaiter = new WithinCancellableCallbackTaskAwaiter(context, innerAwaiter);
        }

        public WithinCancellableCallbackTaskAwaiter GetAwaiter() {
            return _awaiter;
        }

        // The awaiter type that backs WithinCancellableCallbackTaskAwaitable.
        internal struct WithinCancellableCallbackTaskAwaiter : ICriticalNotifyCompletion {

            private static readonly WaitCallback _shunt = state => ((Action)state)();

            private readonly HttpContext _context;
            private readonly TaskAwaiter _innerAwaiter;

            internal WithinCancellableCallbackTaskAwaiter(HttpContext context, TaskAwaiter innerAwaiter) {
                _context = context;
                _innerAwaiter = innerAwaiter;
            }

            public bool IsCompleted {
                get { return _innerAwaiter.IsCompleted; }
            }

            public void GetResult() {
                _innerAwaiter.GetResult();

                // If Response.End was called, need to observe it here.
                HttpContext context = _context;
                if (context != null) {
                    context.Response.ObserveResponseEndCalled();
                }
            }

            public void OnCompleted(Action continuation) {
                Action wrappedContinuation = WrapContinuation(continuation);
                _innerAwaiter.OnCompleted(wrappedContinuation);
            }

            [SecurityPermission(SecurityAction.LinkDemand, Unrestricted = true)] // equivalent of [SecurityCritical]
            public void UnsafeOnCompleted(Action continuation) {
                Action wrappedContinuation = WrapContinuation(continuation);
                _innerAwaiter.UnsafeOnCompleted(wrappedContinuation);
            }

            private Action WrapContinuation(Action continuation) {
                HttpContext context = _context;
                return (context != null)
                    ? () => context.InvokeCancellableCallback(_shunt, continuation)
                    : continuation;
            }

        }

    }
}