// Licensed to the .NET Foundation under one or more agreements.
// The .NET Foundation licenses this file to you under the MIT license.
 
using System.Text.Json;
using Aspire.Hosting.RemoteHost.Ats;
using Aspire.Hosting.RemoteHost.Language;
using Aspire.TypeSystem;
using Microsoft.Extensions.Configuration;
using Microsoft.Extensions.DependencyInjection;
using Microsoft.Extensions.Hosting.Internal;
using Microsoft.Extensions.Logging.Abstractions;
using Xunit;
 
namespace Aspire.Hosting.RemoteHost.Tests;
 
public class IntegrationHostLauncherTests
{
    [Fact]
    public async Task StartAsync_BootstrapSkipsUnavailableIntegrationHosts()
    {
        using var services = new ServiceCollection().BuildServiceProvider();
        var configuration = new ConfigurationBuilder().AddInMemoryCollection(new Dictionary<string, string?>
        {
            ["ASPIRE_INTEGRATION_HOSTS_ENABLED"] = "true",
            ["ASPIRE_INTEGRATION_HOST_BOOTSTRAP"] = "true",
            ["IntegrationHosts:0:PackageName"] = "needs-generated-sdk"
        }).Build();
        using var registry = new ExternalCapabilityRegistry(NullLogger<ExternalCapabilityRegistry>.Instance);
        await using var launcher = IntegrationHostTestServices.CreateLauncher(
            new LanguageSupportResolver(services, () => [], NullLogger<LanguageSupportResolver>.Instance),
            registry,
            configuration,
            new ApplicationLifetime(NullLogger<ApplicationLifetime>.Instance),
            NullLogger<IntegrationHostLauncher>.Instance);
 
        await launcher.StartAsync(TestContext.Current.CancellationToken);
        await launcher.ReadyAsync(TestContext.Current.CancellationToken);
        await launcher.StopAsync(TestContext.Current.CancellationToken);
    }
 
    [Fact]
    public async Task StartAsync_InvalidIntegrationFaultsReadiness()
    {
        using var services = new ServiceCollection().BuildServiceProvider();
        var configuration = new ConfigurationBuilder().AddInMemoryCollection(new Dictionary<string, string?>
        {
            ["ASPIRE_INTEGRATION_HOSTS_ENABLED"] = "true",
            ["IntegrationHosts:0:PackageName"] = "missing-language-and-entrypoint"
        }).Build();
        using var registry = new ExternalCapabilityRegistry(NullLogger<ExternalCapabilityRegistry>.Instance);
        await using var launcher = IntegrationHostTestServices.CreateLauncher(
            new LanguageSupportResolver(services, () => [], NullLogger<LanguageSupportResolver>.Instance),
            registry,
            configuration,
            new ApplicationLifetime(NullLogger<ApplicationLifetime>.Instance),
            NullLogger<IntegrationHostLauncher>.Instance);
 
        var startupException = await Assert.ThrowsAsync<InvalidOperationException>(() =>
            launcher.StartAsync(TestContext.Current.CancellationToken));
        var readinessException = await Assert.ThrowsAsync<InvalidOperationException>(() =>
            launcher.ReadyAsync(TestContext.Current.CancellationToken));
 
        Assert.Same(startupException, readinessException);
    }
 
    [Theory]
    [InlineData(false)]
    [InlineData(true)]
    public async Task StartAsync_DisabledIntegrationHostsFaultReadinessBeforeBootstrap(bool bootstrap)
    {
        using var services = new ServiceCollection().BuildServiceProvider();
        var configuration = new ConfigurationBuilder().AddInMemoryCollection(new Dictionary<string, string?>
        {
            ["ASPIRE_INTEGRATION_HOST_BOOTSTRAP"] = bootstrap.ToString(),
            ["IntegrationHosts:0:PackageName"] = "disabled-integration"
        }).Build();
        using var registry = new ExternalCapabilityRegistry(NullLogger<ExternalCapabilityRegistry>.Instance);
        await using var launcher = IntegrationHostTestServices.CreateLauncher(
            new LanguageSupportResolver(services, () => [], NullLogger<LanguageSupportResolver>.Instance),
            registry, configuration,
            new ApplicationLifetime(NullLogger<ApplicationLifetime>.Instance),
            NullLogger<IntegrationHostLauncher>.Instance);
 
        var exception = await Assert.ThrowsAsync<InvalidOperationException>(() =>
            launcher.StartAsync(TestContext.Current.CancellationToken));
        var readinessException = await Assert.ThrowsAsync<InvalidOperationException>(() =>
            launcher.ReadyAsync(TestContext.Current.CancellationToken));
 
        Assert.Equal("Integration hosts require ASPIRE_INTEGRATION_HOSTS_ENABLED=true.", exception.Message);
        Assert.Same(exception, readinessException);
    }
 
    [Fact]
    public async Task InitializeHostsAsync_WaitsForEveryRegistrationBeforeDiscovery()
    {
        using var services = new ServiceCollection().BuildServiceProvider();
        using var registry = new ExternalCapabilityRegistry(NullLogger<ExternalCapabilityRegistry>.Instance);
        await using var launcher = IntegrationHostTestServices.CreateLauncher(
            new LanguageSupportResolver(services, () => [], NullLogger<LanguageSupportResolver>.Instance),
            registry, new ConfigurationBuilder().Build(),
            new ApplicationLifetime(NullLogger<ApplicationLifetime>.Instance), NullLogger<IntegrationHostLauncher>.Instance);
        var discoveryStarted = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);
        using var firstConnection = new IntegrationHostTestConnection(_ =>
        {
            discoveryStarted.TrySetResult();
            return Task.FromResult(JsonSerializer.SerializeToElement(new[]
            {
                new { id = "test/first", method = "first", returnType = new { typeId = "void", category = "Primitive" } }
            }));
        });
        using var secondConnection = new IntegrationHostTestConnection(
            JsonSerializer.SerializeToElement(new[]
            {
                new { id = "test/second", method = "second", returnType = new { typeId = "void", category = "Primitive" } }
            }));
        registry.AddIntegrationHost(firstConnection.ServerRpc);
 
        var initialization = launcher.InitializeHostsAsync(
            2, Timeout.InfiniteTimeSpan, TimeSpan.FromSeconds(10), TestContext.Current.CancellationToken);
 
        Assert.False(initialization.IsCompleted);
        Assert.False(discoveryStarted.Task.IsCompleted);
        registry.AddIntegrationHost(secondConnection.ServerRpc);
        await initialization.WaitAsync(TimeSpan.FromSeconds(10), TestContext.Current.CancellationToken);
 
        Assert.True(registry.IsRegistered("test/first"));
        Assert.True(registry.IsRegistered("test/second"));
    }
 
    [Fact]
    public async Task InitializeHostsAsync_MissingRegistrationStopsDiscovery()
    {
        using var services = new ServiceCollection().BuildServiceProvider();
        using var registry = new ExternalCapabilityRegistry(NullLogger<ExternalCapabilityRegistry>.Instance);
        await using var launcher = IntegrationHostTestServices.CreateLauncher(
            new LanguageSupportResolver(services, () => [], NullLogger<LanguageSupportResolver>.Instance),
            registry, new ConfigurationBuilder().Build(),
            new ApplicationLifetime(NullLogger<ApplicationLifetime>.Instance), NullLogger<IntegrationHostLauncher>.Instance);
        using var connection = new IntegrationHostTestConnection(
            JsonSerializer.SerializeToElement(new[] { new { id = "test/partial" } }));
        registry.AddIntegrationHost(connection.ServerRpc);
 
        var exception = await Assert.ThrowsAsync<TimeoutException>(() =>
            launcher.InitializeHostsAsync(2, TimeSpan.Zero, TimeSpan.FromSeconds(10), TestContext.Current.CancellationToken));
 
        Assert.StartsWith("Only 1 of 2 integration hosts registered", exception.Message);
        Assert.False(registry.IsRegistered("test/partial"));
    }
 
    [Fact]
    public async Task InitializeHostsAsync_DisconnectedRegistrationCannotPublishPartialCapabilities()
    {
        using var services = new ServiceCollection().BuildServiceProvider();
        using var registry = new ExternalCapabilityRegistry(NullLogger<ExternalCapabilityRegistry>.Instance);
        await using var launcher = IntegrationHostTestServices.CreateLauncher(
            new LanguageSupportResolver(services, () => [], NullLogger<LanguageSupportResolver>.Instance),
            registry, new ConfigurationBuilder().Build(),
            new ApplicationLifetime(NullLogger<ApplicationLifetime>.Instance), NullLogger<IntegrationHostLauncher>.Instance);
        using var connection = new IntegrationHostTestConnection(
            JsonSerializer.SerializeToElement(new[] { new { id = "test/partial" } }));
        registry.AddIntegrationHost(connection.ServerRpc);
        registry.MarkHostUnavailable(connection.ServerRpc);
 
        var exception = await Assert.ThrowsAsync<InvalidOperationException>(() =>
            launcher.InitializeHostsAsync(1, TimeSpan.Zero, TimeSpan.FromSeconds(10), TestContext.Current.CancellationToken));
 
        Assert.Equal("Only 0 of 1 integration hosts remain connected for capability discovery.", exception.Message);
        Assert.Throws<InvalidOperationException>(() => registry.AugmentContext(new AtsContext
        {
            Capabilities = [],
            HandleTypes = [],
            DtoTypes = [],
            EnumTypes = [],
            ExportedValues = [],
            Diagnostics = []
        }));
    }
 
    [Fact]
    public async Task InitializeHostsAsync_CancellationStopsWaitingForRegistration()
    {
        using var services = new ServiceCollection().BuildServiceProvider();
        using var registry = new ExternalCapabilityRegistry(NullLogger<ExternalCapabilityRegistry>.Instance);
        await using var launcher = IntegrationHostTestServices.CreateLauncher(
            new LanguageSupportResolver(services, () => [], NullLogger<LanguageSupportResolver>.Instance),
            registry,
            new ConfigurationBuilder().Build(),
            new ApplicationLifetime(NullLogger<ApplicationLifetime>.Instance), NullLogger<IntegrationHostLauncher>.Instance);
        using var cancellation = CancellationTokenSource.CreateLinkedTokenSource(TestContext.Current.CancellationToken);
 
        var initialization = launcher.InitializeHostsAsync(
            1, Timeout.InfiniteTimeSpan, TimeSpan.FromSeconds(10), cancellation.Token);
        Assert.False(initialization.IsCompleted);
        cancellation.Cancel();
 
        await Assert.ThrowsAnyAsync<OperationCanceledException>(() => initialization);
    }
 
    [Fact]
    public void CreateProcessStartInfo_PreservesArgumentBoundariesAndExpandsEntryPoint()
    {
        var entryPoint = Path.GetFullPath(Path.Combine("integration packages", "host entry.mts"));
        var startInfo = IntegrationHostProcessLauncher.CreateProcessStartInfo(
            "integration-runtime",
            [
                "--no-install",
                "tsx",
                "{entryPoint}",
                "--entry={entryPoint}",
                "{entryPoint};{entryPoint}",
                "--label=a \"quoted\" value",
                @"C:\tools with spaces\trailing\",
                "\\\"",
                "",
                "{otherPlaceholder}"
            ],
            entryPoint,
            isWindows: false);
 
        Assert.Equal(
            [
                "--no-install",
                "tsx",
                entryPoint,
                $"--entry={entryPoint}",
                $"{entryPoint};{entryPoint}",
                "--label=a \"quoted\" value",
                @"C:\tools with spaces\trailing\",
                "\\\"",
                "",
                "{otherPlaceholder}"
            ],
            startInfo.ArgumentList.ToArray());
        Assert.Empty(startInfo.Arguments);
        Assert.Equal("integration-runtime", startInfo.FileName);
        Assert.Equal(Path.GetDirectoryName(entryPoint), startInfo.WorkingDirectory);
        Assert.False(startInfo.UseShellExecute);
        Assert.True(startInfo.RedirectStandardOutput);
        Assert.True(startInfo.RedirectStandardError);
    }
 
    [Fact]
    public void CreateProcessStartInfo_AllowsEmptyArgumentList()
    {
        var entryPoint = Path.GetFullPath("host.mts");
 
        var startInfo = IntegrationHostProcessLauncher.CreateProcessStartInfo("integration-runtime", [], entryPoint, isWindows: false);
 
        Assert.Empty(startInfo.ArgumentList);
        Assert.Empty(startInfo.Arguments);
    }
 
    [Theory]
    [InlineData("npx.cmd")]
    [InlineData("npx.CMD")]
    [InlineData("npx.bat")]
    public void CreateProcessStartInfo_WindowsBatchShim_UsesOuterQuotedCommand(string shim)
    {
        var command = $@"C:\Program Files\%TEMP%\nodejs\{shim}";
        var entryPoint = Path.GetFullPath(Path.Combine("integration %TEMP%! packages", "host entry.mts"));
 
        var startInfo = IntegrationHostProcessLauncher.CreateProcessStartInfo(
            command, ["--no-install", "tsx", "{entryPoint}", "", "literal %PATH%! & value"], entryPoint, isWindows: true);
 
        Assert.Equal(Environment.GetEnvironmentVariable("ComSpec") ?? "cmd.exe", startInfo.FileName);
        Assert.Empty(startInfo.ArgumentList);
        Assert.Equal(
            "/D /V:OFF /S /C \"\"%ASPIRE_COMMAND_SHIM_PATH%\" \"%ASPIRE_COMMAND_SHIM_ARGUMENT_0%\" \"%ASPIRE_COMMAND_SHIM_ARGUMENT_1%\" \"%ASPIRE_COMMAND_SHIM_ARGUMENT_2%\" \"\" \"%ASPIRE_COMMAND_SHIM_ARGUMENT_4%\"\"",
            startInfo.Arguments);
        Assert.Equal(command, startInfo.Environment["ASPIRE_COMMAND_SHIM_PATH"]);
        Assert.Equal(
            ["--no-install", "tsx", entryPoint, "literal %PATH%! & value"],
            startInfo.Environment
                .Where(variable => variable.Key.StartsWith("ASPIRE_COMMAND_SHIM_ARGUMENT_", StringComparison.Ordinal))
                .OrderBy(variable => variable.Key, StringComparer.Ordinal)
                .Select(variable => variable.Value));
        Assert.False(startInfo.UseShellExecute);
        Assert.True(startInfo.RedirectStandardOutput);
        Assert.True(startInfo.RedirectStandardError);
    }
}