// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT license. See LICENSE file in the project root for full license information.

#if NETFRAMEWORK
using System.Runtime.Remoting.Messaging;
#endif

using Microsoft.VisualStudio.TestPlatform.MSTest.TestAdapter.Extensions;
using Microsoft.VisualStudio.TestPlatform.MSTest.TestAdapter.Helpers;
using Microsoft.VisualStudio.TestPlatform.MSTest.TestAdapter.ObjectModel;
using Microsoft.VisualStudio.TestPlatform.MSTestAdapter.PlatformServices.Execution;
using Microsoft.VisualStudio.TestPlatform.MSTestAdapter.PlatformServices.Extensions;
using Microsoft.VisualStudio.TestTools.UnitTesting;

namespace Microsoft.VisualStudio.TestPlatform.MSTest.TestAdapter.Execution;

#pragma warning disable CA1852 // Seal internal types - This class is inherited in tests.
internal partial class TestMethodInfo
{
    /// <summary>
    /// Runs TestCleanup methods of parent TestClass and base classes.
    /// </summary>
    /// <param name="result">Instance of TestResult.</param>
    /// <param name="timeoutTokenSource">The timeout token source.</param>
    /// <param name="invocationState">The activation and cleanup state for this invocation.</param>
    [SuppressMessage("Microsoft.Design", "CA1031:DoNotCatchGeneralExceptionTypes", Justification = "Requirement is to handle all kinds of user exceptions and message appropriately.")]
    private async SynchronizationContextPreservingTask RunTestCleanupMethodAsync(
        TestResult result,
        CancellationTokenSource? timeoutTokenSource,
        TestInvocationState invocationState)
    {
        _ = result ?? throw new ArgumentNullException(nameof(result));

        if (!invocationState.TryGetCleanupSnapshot(out TestInvocationCleanupSnapshot snapshot))
        {
            return;
        }

        bool hasCleanupsToInvoke = HasCleanupsToInvoke(snapshot.ClassInstanceLease);
        bool shouldResetCancellationTokenSource = snapshot.IsTestContextSet
            && TestContext.Context.CancellationTokenSource.IsCancellationRequested;
        if ((!hasCleanupsToInvoke && !shouldResetCancellationTokenSource)
            || !invocationState.TryBeginCleanup())
        {
            return;
        }

        if (shouldResetCancellationTokenSource)
        {
            TestContext.Context.CancellationTokenSource = new CancellationTokenSource();
        }

        if (!hasCleanupsToInvoke)
        {
            return;
        }

        if (!snapshot.IsTestContextSet)
        {
            try
            {
                if (snapshot.ClassInstanceLease is { RequiresCleanup: true } classInstanceLease)
                {
                    await classInstanceLease.DisposeAsync().ConfigureAwait(false);
                }
            }
            catch (Exception exception)
            {
                result.Outcome = result.Outcome.GetMoreImportantOutcome(UnitTestOutcome.Failed);
                result.TestFailureException = exception;
            }

            return;
        }

        MethodInfo? testCleanupMethod = Parent.TestCleanupMethod;
        Exception? testCleanupException;
        TestFailedException? globalTestCleanupException = null;
        MethodInfo? globalTestCleanupFailingMethod = null;
        MethodInfo? currentGlobalTestCleanupMethod = null;
        try
        {
            try
            {
                // If we are running with a method timeout, we need to cancel the cleanup when the overall timeout expires. If it already expired, nothing to do.
                if (timeoutTokenSource is { IsCancellationRequested: false })
                {
                    timeoutTokenSource?.Token.Register(TestContext.Context.CancellationTokenSource.Cancel);
                }

                // Test cleanups are called in the order of discovery
                // Current TestClass -> Parent -> Grandparent
                testCleanupException = testCleanupMethod is not null
                    ? await InvokeCleanupMethodAsync(testCleanupMethod, snapshot.ClassInstance, timeoutTokenSource).ConfigureAwait(false)
                    : null;
                if (testCleanupException is null)
                {
                    foreach (MethodInfo baseCleanupMethod in Parent.BaseTestCleanupMethodsQueue)
                    {
                        testCleanupMethod = baseCleanupMethod;
                        testCleanupException = await InvokeCleanupMethodAsync(baseCleanupMethod, snapshot.ClassInstance, timeoutTokenSource).ConfigureAwait(false);
                        if (testCleanupException is not null)
                        {
                            break;
                        }
                    }
                }
            }
            finally
            {
                if (snapshot.ClassInstanceLease is { RequiresCleanup: true } classInstanceLease)
                {
                    await classInstanceLease.DisposeAsync().ConfigureAwait(false);
                }

                foreach ((MethodInfo method, TimeoutInfo? timeoutInfo) in Parent.Parent.GlobalTestCleanups)
                {
                    currentGlobalTestCleanupMethod = method;
                    globalTestCleanupException = await InvokeGlobalCleanupMethodAsync(method, timeoutInfo, timeoutTokenSource).ConfigureAwait(false);

                    if (globalTestCleanupException is not null)
                    {
                        // Stop on first global cleanup failure, mirroring the local test cleanup loop above.
                        globalTestCleanupFailingMethod = method;
                        break;
                    }

                    currentGlobalTestCleanupMethod = null;
                }
            }
        }
        catch (Exception ex)
        {
            testCleanupException = ex;
            if (globalTestCleanupFailingMethod is null && currentGlobalTestCleanupMethod is not null)
            {
                globalTestCleanupFailingMethod = currentGlobalTestCleanupMethod;
            }

            if (globalTestCleanupFailingMethod is not null)
            {
                testCleanupMethod = globalTestCleanupFailingMethod;
            }
        }

        // If no earlier cleanup failure was recorded but a global test cleanup failed, surface it
        // so the wrapping logic below produces a TestFailedException with the per-method message.
        if (testCleanupException is null && globalTestCleanupException is not null)
        {
            testCleanupException = globalTestCleanupException;
            testCleanupMethod = globalTestCleanupFailingMethod;
        }

        // If testCleanup was successful, then don't do anything
        if (testCleanupException == null)
        {
            return;
        }

        Exception realException = testCleanupException.GetRealException();

        // Check `realException` (not `testCleanupException`) so that a `TestFailedException` wrapped in
        // a `TargetInvocationException` / `TypeInitializationException` still preserves its Outcome
        // (for example a Timeout outcome must not be silently downgraded to Failed).
        UnitTestOutcome outcomeFromRealException = realException is TestFailedException testFailedException
            ? testFailedException.Outcome
            : realException is AssertInconclusiveException
                ? UnitTestOutcome.Inconclusive
                : UnitTestOutcome.Failed;
        result.Outcome = result.Outcome.GetMoreImportantOutcome(outcomeFromRealException);

        realException = testCleanupMethod != null
            ? new TestFailedException(
                outcomeFromRealException,
                string.Format(CultureInfo.CurrentCulture, Resource.UTA_CleanupMethodThrows, TestClassName, testCleanupMethod.Name, realException.GetFormattedExceptionMessage()),
                realException.TryGetStackTraceInformation(),
                realException)
            : new TestFailedException(
                outcomeFromRealException,
                string.Format(CultureInfo.CurrentCulture, Resource.UTA_CleanupMethodThrowsGeneralError, TestClassName, realException.GetFormattedExceptionMessage()),
                realException.TryGetStackTraceInformation(),
                realException);

        result.TestFailureException = realException;
    }

    private bool HasCleanupsToInvoke(ITestClassInstanceLease? classInstanceLease) =>
        Parent.TestCleanupMethod is not null ||
        Parent.BaseTestCleanupMethodsQueue is { Count: > 0 } ||
        classInstanceLease is { RequiresCleanup: true } ||
        Parent.Parent.GlobalTestCleanups is { Count: > 0 };

    /// <summary>
    /// Runs TestInitialize methods of parent TestClass and the base classes.
    /// </summary>
    /// <param name="classInstance">Instance of TestClass.</param>
    /// <param name="result">Instance of TestResult.</param>
    /// <param name="timeoutTokenSource">The timeout token source.</param>
    /// <returns>True if the TestInitialize method(s) did not throw an exception.</returns>
    [SuppressMessage("Microsoft.Design", "CA1031:DoNotCatchGeneralExceptionTypes", Justification = "Requirement is to handle all kinds of user exceptions and message appropriately.")]
    private async SynchronizationContextPreservingTask<bool> RunTestInitializeMethodAsync(object? classInstance, TestResult result, CancellationTokenSource? timeoutTokenSource)
    {
        DebugEx.Assert(classInstance is not null || Parent.IsFSharpModule, "classInstance is null for a non-module test");
        DebugEx.Assert(result != null, "result != null");

        MethodInfo? testInitializeMethod = null;
        Exception? testInitializeException = null;

        try
        {
            // TestInitialize methods for base classes are called in reverse order of discovery
            // Grandparent -> Parent -> Child TestClass
            List<MethodInfo> baseTestInitializeList = Parent.BaseTestInitializeMethodsQueue;
            for (int i = baseTestInitializeList.Count - 1; i >= 0; i--)
            {
                testInitializeMethod = baseTestInitializeList[i];
                testInitializeException = testInitializeMethod is not null
                    ? await InvokeInitializeMethodAsync(testInitializeMethod, classInstance, timeoutTokenSource).ConfigureAwait(false)
                    : null;
                if (testInitializeException is not null)
                {
                    break;
                }
            }

            if (testInitializeException == null)
            {
                testInitializeMethod = Parent.TestInitializeMethod;
                testInitializeException = testInitializeMethod is not null
                    ? await InvokeInitializeMethodAsync(testInitializeMethod, classInstance, timeoutTokenSource).ConfigureAwait(false)
                    : null;
            }
        }
        catch (Exception ex)
        {
            testInitializeException = ex;
        }

        // If testInitialization was successful, then don't do anything
        if (testInitializeException == null)
        {
            return true;
        }

        // If the exception is already a `TestFailedException` we throw it as-is
        if (testInitializeException is TestFailedException tfe)
        {
            result.Outcome = tfe.Outcome;
            result.TestFailureException = testInitializeException;
            return false;
        }

        Exception realException = testInitializeException.GetRealException();

        // Prefix the exception message with the exception type name as prefix when exception is not assert exception.
        string exceptionMessage = realException is UnitTestAssertException
            ? realException.TryGetMessage()
            : realException.GetFormattedExceptionMessage();
        string errorMessage = string.Format(
            CultureInfo.CurrentCulture,
            Resource.UTA_InitMethodThrows,
            TestClassName,
            testInitializeMethod?.Name,
            exceptionMessage);
        StackTraceInformation? stackTrace = realException.GetStackTraceInformation();

        result.Outcome = realException is AssertInconclusiveException
            ? UnitTestOutcome.Inconclusive
            : UnitTestOutcome.Failed;
        result.TestFailureException = new TestFailedException(
            result.Outcome,
            errorMessage,
            stackTrace,
            realException);

        return false;
    }

    private async SynchronizationContextPreservingTask<TestFailedException?> InvokeInitializeMethodAsync(MethodInfo methodInfo, object? classInstance, CancellationTokenSource? timeoutTokenSource)
    {
        TimeoutInfo? timeout = null;
        if (Parent.TestInitializeMethodTimeoutMilliseconds.TryGetValue(methodInfo, out TimeoutInfo localTimeout))
        {
            timeout = localTimeout;
        }

        using IMSTestActivity? activity = MSTestInstrumentation.IsEnabled
            ? MSTestInstrumentation.StartFixtureActivity(
                MSTestInstrumentation.ActivityNames.TestInitialize,
                "test_initialize",
                methodInfo.DeclaringType?.FullName,
                methodInfo.DeclaringType?.Assembly.GetName().Name)
            : null;

        TestFailedException? result = await InvokeFixtureMethodAsync(
            methodInfo,
            classInstance,
            arguments: null,
            timeout,
            timeoutTokenSource,
            Resource.TestInitializeWasCancelled,
            Resource.TestInitializeTimedOut).ConfigureAwait(false);

        if (result is not null)
        {
            activity?.RecordException(result);
        }

        return result;
    }

    private async SynchronizationContextPreservingTask<TestFailedException?> InvokeGlobalInitializeMethodAsync(MethodInfo methodInfo, TimeoutInfo? timeoutInfo, CancellationTokenSource? timeoutTokenSource)
        => await InvokeFixtureMethodAsync(
            methodInfo,
            classInstance: null,
            arguments: [TestContext],
            timeoutInfo,
            timeoutTokenSource,
            Resource.GlobalTestInitializeWasCancelled,
            Resource.GlobalTestInitializeTimedOut).ConfigureAwait(false);

    private async SynchronizationContextPreservingTask<TestFailedException?> InvokeCleanupMethodAsync(MethodInfo methodInfo, object? classInstance, CancellationTokenSource? timeoutTokenSource)
    {
        TimeoutInfo? timeout = null;
        if (Parent.TestCleanupMethodTimeoutMilliseconds.TryGetValue(methodInfo, out TimeoutInfo localTimeout))
        {
            timeout = localTimeout;
        }

        using IMSTestActivity? activity = MSTestInstrumentation.IsEnabled
            ? MSTestInstrumentation.StartFixtureActivity(
                MSTestInstrumentation.ActivityNames.TestCleanup,
                "test_cleanup",
                methodInfo.DeclaringType?.FullName,
                methodInfo.DeclaringType?.Assembly.GetName().Name)
            : null;

        TestFailedException? result = await InvokeFixtureMethodAsync(
            methodInfo,
            classInstance,
            arguments: null,
            timeout,
            timeoutTokenSource,
            Resource.TestCleanupWasCancelled,
            Resource.TestCleanupTimedOut).ConfigureAwait(false);

        if (result is not null)
        {
            activity?.RecordException(result);
        }

        return result;
    }

    private async SynchronizationContextPreservingTask<TestFailedException?> InvokeGlobalCleanupMethodAsync(MethodInfo methodInfo, TimeoutInfo? timeoutInfo, CancellationTokenSource? timeoutTokenSource)
        => await InvokeFixtureMethodAsync(
            methodInfo,
            classInstance: null,
            arguments: [TestContext],
            timeoutInfo,
            timeoutTokenSource,
            Resource.GlobalTestCleanupWasCancelled,
            Resource.GlobalTestCleanupTimedOut).ConfigureAwait(false);

    private async SynchronizationContextPreservingTask<TestFailedException?> InvokeFixtureMethodAsync(
        MethodInfo methodInfo,
        object? classInstance,
        object?[]? arguments,
        TimeoutInfo? timeout,
        CancellationTokenSource? timeoutTokenSource,
        string cancelledMessage,
        string timedOutMessage)
    {
        TestFailedException? result = await FixtureMethodRunner.RunWithTimeoutAndCancellationAsync(
            async () =>
            {
#if NETFRAMEWORK
                CallContext.HostContext = _hostContext;
#endif

                Task? task = methodInfo.GetInvokeResultAsync(classInstance, arguments);
                if (task is not null)
                {
                    await task.ConfigureAwait(false);
                }

                CaptureExecutionContextAfterFixtureIfNeeded(timeout);

#if NETFRAMEWORK
                _hostContext = CallContext.HostContext;
#endif
            },
            TestContext.Context.CancellationTokenSource,
            timeout,
            methodInfo,
            _executionContext,
            cancelledMessage,
            timedOutMessage,
            timeoutTokenSource is null
                ? null
                : (timeoutTokenSource, TimeoutInfo.Timeout)).ConfigureAwait(false);

        return result;
    }

    private void CaptureExecutionContextAfterFixtureIfNeeded(TimeoutInfo? timeoutInfo)
    {
        // After we execute a global test initialize, test initialize, test cleanup, or global test cleanup, we
        // might need to capture the execution context.
        // Generally, we do so only if the method has a timeout and that timeout is non-cooperative.
        // For all other cases, we already use a custom task that preserves synchronization and execution contexts.
        // NOTE: it seems that in .NET Framework, the synchronization context is part of the execution context.
        // However, this doesn't appear to be the case in .NET (Core) where the synchronization context is strictly tied to current thread.
        // In addition, if execution context was captured before (due to use of non-cooperative timeout in a previously run fixture), we still capture here again.
        if (timeoutInfo?.CooperativeCancellation == false || _executionContext is not null)
        {
            _executionContext = ExecutionContext.Capture() ?? _executionContext;
        }
    }
}