// 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.Helpers;
using Microsoft.Testing.Platform.Hosts;
using Microsoft.Testing.Platform.IPC;
using Microsoft.Testing.Platform.Logging;
using Microsoft.Testing.Platform.Services;

using Moq;

namespace Microsoft.Testing.Platform.UnitTests;

[TestClass]
[UnsupportedOSPlatform("browser")]
public sealed class TestHostControllerCancellationTests
{
    // Production closes the stalled control channel after one second. Allow loaded CI agents extra time
    // to schedule the blocking Dispose call and observe the resulting pipe disconnect.
    private static readonly TimeSpan DisposalTestTimeout = TimeSpan.FromSeconds(30);

    public TestContext TestContext { get; set; }

    [TestMethod]
    public async Task RequestCancellation_CancelsChildApplication()
    {
        Mock<ILoggerFactory> loggerFactory = CreateLoggerFactory();
        SystemEnvironment environment = new();
        using var server = new TestHostControllerCancellationServer(
            authorizedSecurityIdentities: null,
            environment,
            loggerFactory.Object,
            new SystemTask());
        server.Start();
        using CTRLPlusCCancellationTokenSource applicationCancellationTokenSource = new();
        using var listener = new TestHostControllerCancellationListener(
            server.PipeName,
            applicationCancellationTokenSource,
            environment,
            new NopLogger());

        await server.WaitForRequestAsync().TimeoutAfterAsync(TimeoutHelper.DefaultHangTimeSpanTimeout);
        server.RequestCancellation();

        Assert.IsTrue(
            applicationCancellationTokenSource.CancellationToken.WaitHandle.WaitOne(TimeoutHelper.DefaultHangTimeSpanTimeout),
            "The controller cancellation request was not propagated to the child application.");
        Assert.IsTrue(listener.WasCancellationRequestedByController);
    }

    [TestMethod]
    public async Task RequestCancellation_BeforeChildConnects_CancelsChildAfterConnection()
    {
        Mock<ILoggerFactory> loggerFactory = CreateLoggerFactory();
        SystemEnvironment environment = new();
        UnexpectedListenerFailureLogger listenerLogger = new();
        using var server = new TestHostControllerCancellationServer(
            authorizedSecurityIdentities: null,
            environment,
            loggerFactory.Object,
            new SystemTask());
        server.Start();
        server.RequestCancellation();
        using CTRLPlusCCancellationTokenSource applicationCancellationTokenSource = new();
        using var listener = new TestHostControllerCancellationListener(
            server.PipeName,
            applicationCancellationTokenSource,
            environment,
            listenerLogger);

        await WaitForOperationOrListenerFailureAsync(server.WaitForRequestAsync(), listenerLogger);

        TaskCompletionSource<bool> cancellationObserved = new(TaskCreationOptions.RunContinuationsAsynchronously);
        using CancellationTokenRegistration registration =
            applicationCancellationTokenSource.CancellationToken.Register(() => cancellationObserved.TrySetResult(true));
        await WaitForOperationOrListenerFailureAsync(cancellationObserved.Task, listenerLogger);

        Assert.IsTrue(listener.WasCancellationRequestedByController);
    }

    [TestMethod]
    public async Task ServerDisposal_WhenClientConnectsWithoutSendingRequest_DoesNotHang()
    {
        Mock<ILoggerFactory> loggerFactory = CreateLoggerFactory();
        SystemEnvironment environment = new();
        var server = new TestHostControllerCancellationServer(
            authorizedSecurityIdentities: null,
            environment,
            loggerFactory.Object,
            new SystemTask());
        server.Start();
        using var client = new System.IO.Pipes.NamedPipeClientStream(
            ".",
            server.PipeName,
            System.IO.Pipes.PipeDirection.InOut,
            System.IO.Pipes.PipeOptions.Asynchronous);
        await client.ConnectAsync(TestContext.CancellationToken);

        await Task.Run(server.Dispose, TestContext.CancellationToken).TimeoutAfterAsync(DisposalTestTimeout);

        byte[] responseBuffer = new byte[1];
        try
        {
            Task<int> readTask = client.ReadAsync(responseBuffer, 0, responseBuffer.Length, TestContext.CancellationToken);
            Task completedTask = await Task.WhenAny(readTask, Task.Delay(DisposalTestTimeout, TestContext.CancellationToken));
            Assert.AreSame(readTask, completedTask, "The client remained connected after server disposal.");
            int bytesRead = await readTask;
            Assert.AreEqual(0, bytesRead);
        }
        catch (IOException ex) when (ex.InnerException is System.Net.Sockets.SocketException
        {
            SocketErrorCode: System.Net.Sockets.SocketError.ConnectionReset,
        })
        {
            // Unix can report a forced close of an idle pipe as a connection reset instead of EOF.
        }
    }

    [TestMethod]
    public async Task ServerDisposal_WhenNoClientConnects_DoesNotHang()
    {
        Mock<ILoggerFactory> loggerFactory = CreateLoggerFactory();
        SystemEnvironment environment = new();
        var server = new TestHostControllerCancellationServer(
            authorizedSecurityIdentities: null,
            environment,
            loggerFactory.Object,
            new SystemTask());
        server.Start();

        await Task.Run(server.Dispose, TestContext.CancellationToken).TimeoutAfterAsync(DisposalTestTimeout);

        using var lateClient = new System.IO.Pipes.NamedPipeClientStream(
            ".",
            server.PipeName,
            System.IO.Pipes.PipeDirection.InOut,
            System.IO.Pipes.PipeOptions.Asynchronous);
        using var timeout = CancellationTokenSource.CreateLinkedTokenSource(TestContext.CancellationToken);
        timeout.CancelAfter(TimeSpan.FromMilliseconds(100));
        await Assert.ThrowsAsync<Exception>(() => lateClient.ConnectAsync(timeout.Token));
    }

    [TestMethod]
    public async Task ServerDisposal_WhenClientKeepsConnectionOpenAfterRequest_DoesNotHang()
    {
        Mock<ILoggerFactory> loggerFactory = CreateLoggerFactory();
        SystemEnvironment environment = new();
        var server = new TestHostControllerCancellationServer(
            authorizedSecurityIdentities: null,
            environment,
            loggerFactory.Object,
            new SystemTask());
        server.Start();
        using var client = new System.IO.Pipes.NamedPipeClientStream(
            ".",
            server.PipeName,
            System.IO.Pipes.PipeDirection.InOut,
            System.IO.Pipes.PipeOptions.Asynchronous);
        await client.ConnectAsync(TestContext.CancellationToken);
        byte[] requestFrame = [.. BitConverter.GetBytes(sizeof(int)), .. BitConverter.GetBytes(13)];
        await client.WriteAsync(requestFrame, 0, requestFrame.Length, TestContext.CancellationToken);
        await client.FlushAsync(TestContext.CancellationToken);
        await server.WaitForRequestAsync().TimeoutAfterAsync(TimeoutHelper.DefaultHangTimeSpanTimeout);

        await Task.Run(server.Dispose, TestContext.CancellationToken).TimeoutAfterAsync(DisposalTestTimeout);

        byte[] responseBuffer = new byte[64];
        int bytesRead;
        do
        {
            bytesRead = await client.ReadAsync(responseBuffer, 0, responseBuffer.Length, TestContext.CancellationToken);
        }
        while (bytesRead > 0);

        Assert.AreEqual(0, bytesRead);
    }

    [TestMethod]
    public async Task NormalServerDisposal_DoesNotCancelChildApplication()
    {
        for (int i = 0; i < 20; i++)
        {
            Mock<ILoggerFactory> loggerFactory = CreateLoggerFactory();
            SystemEnvironment environment = new();
            var server = new TestHostControllerCancellationServer(
                authorizedSecurityIdentities: null,
                environment,
                loggerFactory.Object,
                new SystemTask());
            server.Start();
            using CTRLPlusCCancellationTokenSource applicationCancellationTokenSource = new();
            using var listener = new TestHostControllerCancellationListener(
                server.PipeName,
                applicationCancellationTokenSource,
                environment,
                new NopLogger());

            await server.WaitForRequestAsync().TimeoutAfterAsync(TimeoutHelper.DefaultHangTimeSpanTimeout);
            server.Dispose();

            Assert.IsFalse(applicationCancellationTokenSource.CancellationToken.IsCancellationRequested);
            Assert.IsFalse(listener.WasCancellationRequestedByController);
        }
    }

    private static Mock<ILoggerFactory> CreateLoggerFactory()
    {
        Mock<ILoggerFactory> loggerFactory = new();
        loggerFactory.Setup(x => x.CreateLogger(It.IsAny<string>())).Returns(new NopLogger());
        return loggerFactory;
    }

    private async Task WaitForOperationOrListenerFailureAsync(Task operation, UnexpectedListenerFailureLogger listenerLogger)
    {
        Task<Task> completionTask = Task.WhenAny(operation, listenerLogger.Failure);
        await completionTask.TimeoutAfterAsync(TimeoutHelper.DefaultHangTimeSpanTimeout, TestContext.CancellationToken);

        if (await completionTask == listenerLogger.Failure)
        {
            Assert.Fail(await listenerLogger.Failure);
        }

        await operation;
    }

    private sealed class UnexpectedListenerFailureLogger : ILogger
    {
        private readonly TaskCompletionSource<string> _failure = new(TaskCreationOptions.RunContinuationsAsynchronously);

        public Task<string> Failure => _failure.Task;

        public bool IsEnabled(LogLevel logLevel) => true;

        public void Log<TState>(LogLevel logLevel, TState state, Exception? exception, Func<TState, Exception?, string> formatter)
            => _failure.TrySetResult(formatter(state, exception));

        public Task LogAsync<TState>(LogLevel logLevel, TState state, Exception? exception, Func<TState, Exception?, string> formatter)
        {
            _failure.TrySetResult(formatter(state, exception));
            return Task.CompletedTask;
        }
    }
}