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

using Microsoft.Testing.Platform.Extensions;
using Microsoft.Testing.Platform.Extensions.TestHost;
using Microsoft.Testing.Platform.Helpers;
using Microsoft.Testing.Platform.Logging;
using Microsoft.Testing.Platform.Messages;
using Microsoft.Testing.Platform.OutputDevice;
using Microsoft.Testing.Platform.Services;
using Microsoft.Testing.Platform.Telemetry;

namespace Microsoft.Testing.Platform.Hosts;

internal abstract partial class CommonHost
{
    private static async Task LogShutdownWarningAsync(IServiceProvider serviceProvider, string message)
    {
        try
        {
            if (serviceProvider.GetService<ILoggerFactory>()?.CreateLogger(nameof(CommonHost)) is { } logger)
            {
                await logger.LogWarningAsync(message).ConfigureAwait(false);
            }
        }
        catch (Exception)
        {
            // This runs on teardown paths that are often already unwinding, and the logging stack is itself
            // being torn down (a disposed LoggerFactoryProxy throws from CreateLogger, and providers can fail
            // from LogAsync). A failed warning must not replace the exception we were reporting, nor abort the
            // rest of the disposal.
        }
    }

    private static async Task DisplayBeforeSessionStartAsync(ProxyOutputDevice outputDevice, ITestSessionContext sessionInfo)
    {
        // Display before session start
        await outputDevice.DisplayBeforeSessionStartAsync(sessionInfo.CancellationToken).ConfigureAwait(false);

        if (outputDevice.OriginalOutputDevice is ITestSessionLifetimeHandler testSessionLifetimeHandler)
        {
            await testSessionLifetimeHandler.OnTestSessionStartingAsync(sessionInfo).ConfigureAwait(false);
        }
    }

    private static async Task DisplayAfterSessionEndRunAsync(ProxyOutputDevice outputDevice, ITestSessionContext sessionInfo)
    {
        // Display after session end even when the session cancellation token is canceled.
        // We intentionally pass a non-cancelable token so final output/cleanup notifications are not skipped.
        await outputDevice.DisplayAfterSessionEndRunAsync(CancellationToken.None).ConfigureAwait(false);

        // We want to ensure that the output service is the last one to run
        if (outputDevice.OriginalOutputDevice is ITestSessionLifetimeHandler testSessionLifetimeHandlerFinishing)
        {
            await testSessionLifetimeHandlerFinishing.OnTestSessionFinishingAsync(sessionInfo).ConfigureAwait(false);
        }
    }

    private static async Task NotifyTestSessionStartAsync(ITestSessionContext testSessionContext, BaseMessageBus baseMessageBus, ServiceProvider serviceProvider, IPlatformOpenTelemetryService? otelService)
    {
        TestSessionLifetimeHandlersContainer? testSessionLifetimeHandlersContainer = serviceProvider.GetService<TestSessionLifetimeHandlersContainer>();
        if (testSessionLifetimeHandlersContainer is null)
        {
            return;
        }

        foreach (ITestSessionLifetimeHandler testSessionLifetimeHandler in testSessionLifetimeHandlersContainer.TestSessionLifetimeHandlers)
        {
            using IPlatformActivity? activity = otelService?.StartActivity(testSessionLifetimeHandler.Uid, testSessionLifetimeHandler.ToOTelTags());
            await testSessionLifetimeHandler.OnTestSessionStartingAsync(testSessionContext).ConfigureAwait(false);
        }

        // Drain messages generated by the session start notification before to start test execution.
        await baseMessageBus.DrainDataAsync().ConfigureAwait(false);
    }

    private static async Task NotifyTestSessionEndAsync(ITestSessionContext testSessionContext, BaseMessageBus baseMessageBus, ServiceProvider serviceProvider, IPlatformOpenTelemetryService? otelService)
    {
        // Drain messages generated by the test session execution before to process the session end notification.
        await baseMessageBus.DrainDataAsync().ConfigureAwait(false);

        TestSessionLifetimeHandlersContainer? testSessionLifetimeHandlersContainer = serviceProvider.GetService<TestSessionLifetimeHandlersContainer>();
        if (testSessionLifetimeHandlersContainer is null)
        {
            // Nothing else to notify. The bus is disabled by the caller's teardown safety net
            // (EnsureMessageBusDisabledAsync), which runs for every outcome of the session.
            return;
        }

        IShutdownProgressReporter? shutdownProgressReporter = serviceProvider.GetService<IShutdownProgressReporter>();

        // First, we call OnTestSessionFinishingAsync on all non-consumers.
        bool hasNonDataConsumers = false;
        foreach (ITestSessionLifetimeHandler testSessionLifetimeHandler in testSessionLifetimeHandlersContainer.TestSessionLifetimeHandlers)
        {
            if (testSessionLifetimeHandler is IDataConsumer)
            {
                // At first, we don't call this for data consumers.
                // We want to do this for potentially publishers-only handlers.
                // The order here matters, because one handler can produce a message that is consumed by another.
                // By making the consumers last, we reduce the likelihood of issues related to ordering.
                // One case where this is important is the dependency between Code Coverage's ITestSessionLifetimeHandler and TRX's ITestSessionLifetimeHandler.
                // We must run Code Coverage ITestSessionLifetimeHandler implementation first, as it will publish SessionFileArtifact.
                // The SessionFileArtifact is expected to be consumed by TRX's implementation of ITestSessionLifetimeHandler.
                // In that case, ITestSessionLifetimeHandler of CC is a producer-only handler which we will run first.
                // Then we will drain the message bus to ensure TRX handler have consumed the SessionFileArtifact.
                // Then we run TRX OnTestSessionFinishingAsync.
                continue;
            }

            hasNonDataConsumers = true;

            using (otelService?.StartActivity(testSessionLifetimeHandler.Uid, testSessionLifetimeHandler.ToOTelTags()))
            using (shutdownProgressReporter?.Track(testSessionLifetimeHandler.Uid, testSessionLifetimeHandler.DisplayName, nameof(ITestSessionLifetimeHandler.OnTestSessionFinishingAsync)))
            {
                await testSessionLifetimeHandler.OnTestSessionFinishingAsync(testSessionContext).ConfigureAwait(false);
            }
        }

        if (hasNonDataConsumers)
        {
            // At this point, we called all non-consumer handlers.
            // Now, we want to make sure to drain the message bus before calling the consumer handlers.
            // Messages produced by non-consumer handlers could be consumed by consumer handlers.
            await baseMessageBus.DrainDataAsync().ConfigureAwait(false);
        }

        foreach (ITestSessionLifetimeHandler testSessionLifetimeHandler in testSessionLifetimeHandlersContainer.TestSessionLifetimeHandlers)
        {
            if (testSessionLifetimeHandler is not IDataConsumer)
            {
                // We already called this for non-consumers.
                continue;
            }

            using (otelService?.StartActivity(testSessionLifetimeHandler.Uid, testSessionLifetimeHandler.ToOTelTags()))
            using (shutdownProgressReporter?.Track(testSessionLifetimeHandler.Uid, testSessionLifetimeHandler.DisplayName, nameof(ITestSessionLifetimeHandler.OnTestSessionFinishingAsync)))
            {
                await testSessionLifetimeHandler.OnTestSessionFinishingAsync(testSessionContext).ConfigureAwait(false);
            }

            // OnTestSessionFinishingAsync could produce information that needs to be handled by others.
            // While in many cases we already handled this by calling all non-consumers first, it's possible that
            // two consumers might depend on each other. In that case, we solely rely on registration order.
            // And we drain in between.
            await baseMessageBus.DrainDataAsync().ConfigureAwait(false);
        }

        // We disable after the drain because it's possible that the drain will produce more messages
        await baseMessageBus.DrainDataAsync().ConfigureAwait(false);
        await baseMessageBus.DisableAsync().ConfigureAwait(false);
    }
}