// Licensed to the .NET Foundation under one or more agreements.
// The .NET Foundation licenses this file to you under the MIT license.

using System.Diagnostics;
using System.Runtime.ExceptionServices;
using System.Runtime.InteropServices;

namespace System.Threading.Tasks.Sources
{
    /// <summary>Provides the core logic for implementing a manual-reset <see cref="IValueTaskSource"/> or <see cref="IValueTaskSource{TResult}"/>.</summary>
    /// <typeparam name="TResult">Specifies the type of results of the operation represented by this instance.</typeparam>
    [StructLayout(LayoutKind.Auto)]
    public struct ManualResetValueTaskSourceCore<TResult>
    {
        /// <summary>
        /// The callback to invoke when the operation completes if <see cref="OnCompleted"/> was called before the operation completed,
        /// or <see cref="ManualResetValueTaskSourceCoreShared.s_sentinel"/> if the operation completed before a callback was supplied,
        /// or null if a callback hasn't yet been provided and the operation hasn't yet completed.
        /// </summary>
        private Action<object?>? _continuation;
        /// <summary>State to pass to <see cref="_continuation"/>.</summary>
        private object? _continuationState;
        /// <summary>
        /// Null if no special context was found.
        /// ExecutionContext if one was captured due to needing to be flowed.
        /// A scheduler (TaskScheduler or SynchronizationContext) if one was captured and needs to be used for callback scheduling.
        /// Or a CapturedContext if there's both an ExecutionContext and a scheduler.
        /// The most common and the fast path case to optimize for is null.
        /// </summary>
        private object? _capturedContext;
        /// <summary>The exception with which the operation failed, or null if it hasn't yet completed or completed successfully.</summary>
        private ExceptionDispatchInfo? _error;
        /// <summary>The result with which the operation succeeded, or the default value if it hasn't yet completed or failed.</summary>
        private TResult? _result;
        /// <summary>The current version of this value, used to help prevent misuse.</summary>
        private short _version;
        /// <summary>Whether to force continuations to run asynchronously.</summary>
        private bool _runContinuationsAsynchronously;

        /// <summary>Gets or sets whether to force continuations to run asynchronously.</summary>
        /// <remarks>Continuations may run asynchronously if this is false, but they'll never run synchronously if this is true.</remarks>
        public bool RunContinuationsAsynchronously
        {
            get => _runContinuationsAsynchronously;
            set => _runContinuationsAsynchronously = value;
        }

        /// <summary>Resets to prepare for the next operation.</summary>
        public void Reset()
        {
            // Reset/update state for the next use/await of this instance.
            // Order of assignments is unimportant here.
            // The outer user always ensures that the state is not accessed across
            // the reset point when implementing Rent/Return operations.
            _version++;
            _continuation = null;
            _continuationState = null;
            _capturedContext = null;
            _error = null;
            _result = default;
        }

        /// <summary>Completes with a successful result.</summary>
        /// <param name="result">The result.</param>
        public void SetResult(TResult result)
        {
            _result = result;
            SignalCompletion();
        }

        /// <summary>Completes with an error.</summary>
        /// <param name="error">The exception.</param>
        public void SetException(Exception error)
        {
            _error = ExceptionDispatchInfo.Capture(error);
            SignalCompletion();
        }

        /// <summary>Gets the operation version.</summary>
        public short Version => _version;

        /// <summary>Gets whether the operation has completed.</summary>
        internal bool IsCompleted => ReferenceEquals(Volatile.Read(ref _continuation), ManualResetValueTaskSourceCoreShared.s_sentinel);

        /// <summary>Gets the continuation object for diagnostic purposes only.</summary>
        internal object? ContinuationForDiagnostics => _continuationState;

        /// <summary>Gets the status of the operation.</summary>
        /// <param name="token">Opaque value that was provided to the <see cref="ValueTask"/>'s constructor.</param>
        public ValueTaskSourceStatus GetStatus(short token)
        {
            ValidateToken(token);
            return
                !IsCompleted ? ValueTaskSourceStatus.Pending :
                _error is null ? ValueTaskSourceStatus.Succeeded :
                _error.SourceException is OperationCanceledException ? ValueTaskSourceStatus.Canceled :
                ValueTaskSourceStatus.Faulted;
        }

        /// <summary>Gets the result of the operation.</summary>
        /// <param name="token">Opaque value that was provided to the <see cref="ValueTask"/>'s constructor.</param>
        [StackTraceHidden]
        public TResult GetResult(short token)
        {
            if (token != _version || !IsCompleted || _error is not null)
            {
                ThrowForFailedGetResult();
            }

            return _result!;
        }

        /// <summary>Throws an exception in response to a failed <see cref="GetResult"/>.</summary>
        [StackTraceHidden]
        private void ThrowForFailedGetResult()
        {
            _error?.Throw();
            throw new InvalidOperationException(); // not using ThrowHelper.ThrowInvalidOperationException so that the JIT sees ThrowForFailedGetResult as always throwing
        }

        /// <summary>Schedules the continuation action for this operation.</summary>
        /// <param name="continuation">The continuation to invoke when the operation has completed.</param>
        /// <param name="state">The state object to pass to <paramref name="continuation"/> when it's invoked.</param>
        /// <param name="token">Opaque value that was provided to the <see cref="ValueTask"/>'s constructor.</param>
        /// <param name="flags">The flags describing the behavior of the continuation.</param>
        public void OnCompleted(Action<object?> continuation, object? state, short token, ValueTaskSourceOnCompletedFlags flags)
        {
            if (continuation is null)
            {
                ThrowHelper.ThrowArgumentNullException(ExceptionArgument.continuation);
            }
            ValidateToken(token);

            // We need to store the state before the CompareExchange, so that if it completes immediately
            // after the CompareExchange, it'll find the state already stored.  If someone misuses this
            // and schedules multiple continuations erroneously, we could end up using the wrong state.
            // Make a best-effort attempt to catch such misuse.
            if (_continuationState is not null)
            {
                ThrowHelper.ThrowInvalidOperationException();
            }
            _continuationState = state;

            Debug.Assert(_capturedContext is null);
            if ((flags & ValueTaskSourceOnCompletedFlags.FlowExecutionContext) != 0)
            {
                _capturedContext = ExecutionContext.Capture();
            }

            if ((flags & ValueTaskSourceOnCompletedFlags.UseSchedulingContext) != 0)
            {
                if (SynchronizationContext.Current is SynchronizationContext sc &&
                    sc.GetType() != typeof(SynchronizationContext))
                {
                    _capturedContext = _capturedContext is null ?
                        sc :
                        new CapturedSchedulerAndExecutionContext(sc, (ExecutionContext)_capturedContext);
                }
                else
                {
                    TaskScheduler ts = TaskScheduler.Current;
                    if (ts != TaskScheduler.Default)
                    {
                        _capturedContext = _capturedContext is null ?
                            ts :
                            new CapturedSchedulerAndExecutionContext(ts, (ExecutionContext)_capturedContext);
                    }
                }
            }

            // Try to set the provided continuation into _continuation.  If this succeeds, that means the operation
            // has not yet completed, and the completer will be responsible for invoking the callback.  If this fails,
            // that means the operation has already completed, and we must invoke the callback, but because we're still
            // inside the awaiter's OnCompleted method and we want to avoid possible stack dives, we must invoke
            // the continuation asynchronously rather than synchronously.
            Action<object?>? prevContinuation = Interlocked.CompareExchange(ref _continuation, continuation, null);
            if (prevContinuation is null)
            {
                // Operation hadn't already completed, so we're done. The continuation will be
                // invoked when SetResult/Exception is called at some later point.
                return;
            }

            // Queue the continuation.  We always queue here, even if !RunContinuationsAsynchronously, in order
            // to avoid stack diving; this path happens in the rare race when we're setting up to await and the
            // object is completed after the awaiter.IsCompleted but before the awaiter.OnCompleted.

            // We no longer need the stored values as we will be passing the state when queuing directly
            _continuationState = null;
            object? capturedContext = _capturedContext;
            _capturedContext = null;

            // If the set failed because there's already a delegate in _continuation, but that delegate is
            // something other than the completion sentinel, something went wrong, which should only happen if
            // the instance was erroneously used, likely to hook up multiple continuations.
            if (!ReferenceEquals(prevContinuation, ManualResetValueTaskSourceCoreShared.s_sentinel))
            {
                ThrowHelper.ThrowInvalidOperationException();
            }

            switch (capturedContext)
            {
                case null:
                    ThreadPool.UnsafeQueueUserWorkItem(continuation, state, preferLocal: true);
                    break;

                case ExecutionContext:
                    ThreadPool.QueueUserWorkItem(continuation, state, preferLocal: true);
                    break;

                default:
                    ManualResetValueTaskSourceCoreShared.ScheduleCapturedContext(capturedContext, continuation, state);
                    break;
            }
        }

        /// <summary>Ensures that the specified token matches the current version.</summary>
        /// <param name="token">The token supplied by <see cref="ValueTask"/>.</param>
        private void ValidateToken(short token)
        {
            if (token != _version)
            {
                ThrowHelper.ThrowInvalidOperationException();
            }
        }

        /// <summary>Signals that the operation has completed.  Invoked after the result or error has been set.</summary>
        private void SignalCompletion()
        {
            if (IsCompleted)
            {
                ThrowHelper.ThrowInvalidOperationException();
            }

            Action<object?>? continuation = Interlocked.Exchange(ref _continuation, ManualResetValueTaskSourceCoreShared.s_sentinel);

            if (continuation is not null)
            {
                object? state = _continuationState;
                _continuationState = null;
                object? context = _capturedContext;
                _capturedContext = null;

                if (context is null)
                {
                    if (_runContinuationsAsynchronously)
                    {
                        ThreadPool.UnsafeQueueUserWorkItem(continuation, state, preferLocal: true);
                    }
                    else
                    {
                        continuation(state);
                    }
                }
                else if (context is ExecutionContext or CapturedSchedulerAndExecutionContext)
                {
                    ManualResetValueTaskSourceCoreShared.InvokeContinuationWithContext(context, continuation, state, _runContinuationsAsynchronously);
                }
                else
                {
                    Debug.Assert(context is TaskScheduler or SynchronizationContext, $"context is {context}");
                    ManualResetValueTaskSourceCoreShared.ScheduleCapturedContext(context, continuation, state);
                }
            }
        }
    }

    /// <summary>A tuple of both a non-null scheduler and a non-null ExecutionContext.</summary>
    internal sealed class CapturedSchedulerAndExecutionContext
    {
        internal readonly object _scheduler;
        internal readonly ExecutionContext _executionContext;

        public CapturedSchedulerAndExecutionContext(object scheduler, ExecutionContext executionContext)
        {
            Debug.Assert(scheduler is SynchronizationContext or TaskScheduler, $"{nameof(scheduler)} is {scheduler}");
            Debug.Assert(executionContext is not null, $"{nameof(executionContext)} is null");

            _scheduler = scheduler;
            _executionContext = executionContext;
        }
    }

    internal static class ManualResetValueTaskSourceCoreShared // separated out of generic to avoid unnecessary duplication
    {
        internal static readonly Action<object?> s_sentinel = CompletionSentinel;

        private static void CompletionSentinel(object? _) // named method to aid debugging
        {
            Debug.Fail("The sentinel delegate should never be invoked.");
            ThrowHelper.ThrowInvalidOperationException();
        }

        internal static void ScheduleCapturedContext(object context, Action<object?> continuation, object? state)
        {
            Debug.Assert(
                context is SynchronizationContext or TaskScheduler or CapturedSchedulerAndExecutionContext,
                $"{nameof(context)} is {context}");

            switch (context)
            {
                case SynchronizationContext sc:
                    ScheduleSynchronizationContext(sc, continuation, state);
                    break;

                case TaskScheduler ts:
                    ScheduleTaskScheduler(ts, continuation, state);
                    break;

                default:
                    CapturedSchedulerAndExecutionContext cc = (CapturedSchedulerAndExecutionContext)context;
                    if (cc._scheduler is SynchronizationContext ccsc)
                    {
                        ScheduleSynchronizationContext(ccsc, continuation, state);
                    }
                    else
                    {
                        Debug.Assert(cc._scheduler is TaskScheduler, $"{nameof(cc._scheduler)} is {cc._scheduler}");
                        ScheduleTaskScheduler((TaskScheduler)cc._scheduler, continuation, state);
                    }
                    break;
            }

            static void ScheduleSynchronizationContext(SynchronizationContext sc, Action<object?> continuation, object? state) =>
                sc.Post(continuation.Invoke, state);

            static void ScheduleTaskScheduler(TaskScheduler scheduler, Action<object?> continuation, object? state) =>
                Task.Factory.StartNew(continuation, state, CancellationToken.None, TaskCreationOptions.DenyChildAttach, scheduler);
        }

        internal static void InvokeContinuationWithContext(object capturedContext, Action<object?> continuation, object? continuationState, bool runContinuationsAsynchronously)
        {
            // This is in a helper as the error handling causes the generated asm
            // for the surrounding code to become less efficient (stack spills etc)
            // and it is an uncommon path.
            Debug.Assert(continuation is not null, $"{nameof(continuation)} is null");
            Debug.Assert(capturedContext is ExecutionContext or CapturedSchedulerAndExecutionContext, $"{nameof(capturedContext)} is {capturedContext}");

            // Capture the current EC.  We'll switch over to the target EC and then restore back to this one.
            ExecutionContext? currentContext = ExecutionContext.CaptureForRestore();

            if (capturedContext is ExecutionContext ec)
            {
                ExecutionContext.RestoreInternal(ec); // Restore the captured ExecutionContext before executing anything.
                if (runContinuationsAsynchronously)
                {
                    try
                    {
                        ThreadPool.QueueUserWorkItem(continuation, continuationState, preferLocal: true);
                    }
                    finally
                    {
                        ExecutionContext.RestoreInternal(currentContext); // Restore the current ExecutionContext.
                    }
                }
                else
                {
                    // Running inline may throw; capture the edi if it does as we changed the ExecutionContext,
                    // so need to restore it back before propagating the throw.
                    ExceptionDispatchInfo? edi = null;
                    SynchronizationContext? syncContext = SynchronizationContext.Current;
                    try
                    {
                        continuation(continuationState);
                    }
                    catch (Exception ex)
                    {
                        // Note: we have a "catch" rather than a "finally" because we want
                        // to stop the first pass of EH here.  That way we can restore the previous
                        // context before any of our callers' EH filters run.
                        edi = ExceptionDispatchInfo.Capture(ex);
                    }
                    finally
                    {
                        // Set sync context back to what it was prior to coming in.
                        // Then restore the current ExecutionContext.
                        SynchronizationContext.SetSynchronizationContext(syncContext);
                        ExecutionContext.RestoreInternal(currentContext);
                    }

                    // Now rethrow the exception; if there is one.
                    edi?.Throw();
                }
            }
            else
            {
                CapturedSchedulerAndExecutionContext cc = (CapturedSchedulerAndExecutionContext)capturedContext;
                ExecutionContext.Restore(cc._executionContext); // Restore the captured ExecutionContext before executing anything.
                try
                {
                    ScheduleCapturedContext(capturedContext, continuation, continuationState);
                }
                finally
                {
                    ExecutionContext.RestoreInternal(currentContext); // Restore the current ExecutionContext.
                }
            }
        }
    }
}