// Licensed to the .NET Foundation under one or more agreements.
// The .NET Foundation licenses this file to you under the MIT license.
 
using System.Collections.Concurrent;
using System.Linq;
using System.Net;
using System.Net.Sockets;
using Microsoft.AspNetCore.Builder;
using Microsoft.AspNetCore.Http;
using Microsoft.Extensions.DependencyInjection;
using Yarp.ReverseProxy.Configuration;
 
namespace Microsoft.AspNetCore.Components.Testing.Infrastructure;
 
/// <summary>
/// Non-generic base class for an E2E server factory.
/// </summary>
public abstract class ServerFactory : IAsyncDisposable
{
    private protected ServerFactory()
    {
    }
 
    internal abstract Task<ServerInstance> StartServerAsync(
        string appName,
        Action<ServerStartOptions>? configure,
        ITestArtifactManager? artifactManager);
 
    internal abstract ValueTask DisposeAsyncCore();
 
    ValueTask IAsyncDisposable.DisposeAsync() => DisposeAsyncCore();
}
 
/// <summary>
/// Factory + registry that manages the full E2E test lifecycle:
/// loads the E2E manifest, runs a YARP reverse proxy on a dynamic port,
/// and starts app instances on demand through
/// <see cref="Microsoft.AspNetCore.Components.Testing.Playwright.UITest"/>.
/// <para>
/// One instance per test assembly. Owned by the consumer (typically a static field
/// populated from an <c>[AssemblyInitialize]</c> method and disposed from
/// <c>[AssemblyCleanup]</c>).
/// </para>
/// </summary>
/// <typeparam name="TTestAssembly">
/// A marker type defined in the test project (e.g., <c>public class E2ETestAssembly;</c>).
/// Used to locate the test assembly, which determines the manifest filename
/// and the assembly path for <c>DOTNET_STARTUP_HOOKS</c> and <c>ASPNETCORE_HOSTINGSTARTUPASSEMBLIES</c>.
/// </typeparam>
/// <remarks>
/// <para>
/// Tests derive from <see cref="Microsoft.AspNetCore.Components.Testing.Playwright.UITest"/>
/// and use its server-start helper. Starting the same app with the same options is idempotent.
/// </para>
/// <para>
/// Usage:
/// <code>
/// public class E2ETestAssembly;
///
/// public static class TestRoot
/// {
///     public static ServerFactory&lt;E2ETestAssembly&gt; Servers { get; private set; } = null!;
///
///     [AssemblyInitialize]
///     public static async Task Init(TestContext _)
///     {
///         Servers = new ServerFactory&lt;E2ETestAssembly&gt;();
///         await Servers.InitializeAsync();
///     }
///
///     [AssemblyCleanup]
///     public static Task Cleanup() =&gt; Servers.DisposeAsync().AsTask();
/// }
///
/// [UITest]
/// public partial class MyTests : Microsoft.AspNetCore.Components.Testing.Playwright.PageTest
/// {
///     private ServerInstance _server = null!;
///
///     protected override async Task InitializeCoreAsync()
///     {
///         await base.InitializeCoreAsync();
///         _server = await StartServerAsync&lt;App&gt;(TestRoot.Servers);
///     }
/// }
/// </code>
/// </para>
/// </remarks>
public sealed class ServerFactory<TTestAssembly> : ServerFactory where TTestAssembly : class
{
    private WebApplication? _proxyApp;
    private InMemoryConfigProvider? _configProvider;
    private readonly ConcurrentDictionary<string, ServerInstance> _instances = new();
    private readonly ConcurrentDictionary<string, TaskCompletionSource> _readinessSignals = new();
    private readonly SemaphoreSlim _startLock = new(1, 1);
    private string _testAssemblyLocation = "";
    private string _testAssemblyName = "";
 
    internal E2EManifest Manifest { get; private set; } = default!;
 
    /// <summary>
    /// Base URL of the YARP reverse proxy (e.g., <c>http://localhost:12345</c>).
    /// </summary>
    public string ProxyUrl { get; private set; } = "";
 
    /// <inheritdoc/>
    public async ValueTask InitializeAsync()
    {
        // Resolve test assembly info via the TTestAssembly marker type
        var assembly = typeof(TTestAssembly).Assembly;
        _testAssemblyName = assembly.GetName().Name!;
        _testAssemblyLocation = assembly.Location;
 
        // Load the manifest generated by Microsoft.AspNetCore.Components.Testing.targets
        Manifest = E2EManifest.Load(_testAssemblyName);
 
        // Start the YARP reverse proxy on a dynamic port
        var port = GetAvailablePort();
        ProxyUrl = $"http://localhost:{port}";
 
        var builder = WebApplication.CreateSlimBuilder();
        builder.Services.AddRoutingCore();
        builder.Services.AddReverseProxy().LoadFromMemory([], []);
 
        _proxyApp = builder.Build();
 
        // Readiness callback endpoint: app processes POST here when they're ready.
        // Each token maps to a TaskCompletionSource created in StartServerAsync.
        _proxyApp.MapPost("/_ready/{token}", (string token) =>
        {
            if (_readinessSignals.TryRemove(token, out var tcs))
            {
                tcs.TrySetResult();
                return Results.Ok();
            }
            return Results.NotFound();
        });
 
        _proxyApp.MapReverseProxy();
        _proxyApp.Urls.Add(ProxyUrl);
 
        await _proxyApp.StartAsync().ConfigureAwait(false);
 
        _configProvider = _proxyApp.Services.GetRequiredService<InMemoryConfigProvider>();
    }
 
    internal override async Task<ServerInstance> StartServerAsync(
        string appName,
        Action<ServerStartOptions>? configure,
        ITestArtifactManager? artifactManager)
    {
        var options = new ServerStartOptions();
        configure?.Invoke(options);
 
        var key = ServerInstance.ComputeKey(appName, options);
 
        if (_instances.TryGetValue(key, out var existing))
        {
            return existing;
        }
 
        await _startLock.WaitAsync().ConfigureAwait(false);
        try
        {
            // Double-check after acquiring lock
            if (_instances.TryGetValue(key, out existing))
            {
                return existing;
            }
 
            var appEntry = Manifest.GetApp(appName)
                ?? throw new InvalidOperationException(
                    $"App '{appName}' not found in E2E manifest. " +
                    $"Available apps: {string.Join(", ", Manifest.Apps.Keys)}");
 
            var instance = new ServerInstance(appName, key, ProxyUrl, instanceKey =>
            {
                _instances.TryRemove(instanceKey, out _);
                UpdateProxyConfig();
            });
 
            // Create a readiness signal: the app will POST to /_ready/{token}
            // when it is listening, completing this TCS.
            var readyToken = Guid.NewGuid().ToString("N");
            var readyTcs = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);
            _readinessSignals[readyToken] = readyTcs;
 
            var readyUrl = $"{ProxyUrl}/_ready/{readyToken}";
 
            try
            {
                await instance.LaunchAsync(appEntry, options, _testAssemblyLocation, _testAssemblyName, readyUrl, readyTcs.Task).ConfigureAwait(false);
            }
            catch (Exception launchException)
            {
                var exceptions = new List<Exception> { launchException };
 
                if (artifactManager is not null)
                {
                    try
                    {
                        var directory = artifactManager.CreateArtifactDirectory("server-output");
                        var paths = instance.WriteStartupFailureArtifacts(directory, launchException);
                        artifactManager.AddArtifacts(paths);
                    }
                    catch (Exception artifactException)
                    {
                        exceptions.Add(artifactException);
                    }
                }
 
                try
                {
                    await instance.DisposeAsync().ConfigureAwait(false);
                }
                catch (Exception disposeException)
                {
                    exceptions.Add(disposeException);
                }
 
                if (exceptions.Count > 1)
                {
                    throw new AggregateException(exceptions);
                }
 
                throw;
            }
            finally
            {
                // Clean up stale signal if launch failed before the app called back
                _readinessSignals.TryRemove(readyToken, out _);
            }
 
            _instances[key] = instance;
            UpdateProxyConfig();
 
            return instance;
        }
        finally
        {
            _startLock.Release();
        }
    }
 
    /// <inheritdoc/>
    public async ValueTask DisposeAsync()
    {
        // Dispose all running instances
        foreach (var instance in _instances.Values)
        {
            try
            {
                await instance.DisposeAsync().ConfigureAwait(false);
            }
 
            catch
            {
                // Best-effort cleanup
            }
        }
        _instances.Clear();
 
        // Stop the YARP proxy
        if (_proxyApp is not null)
        {
            await _proxyApp.StopAsync().ConfigureAwait(false);
            await _proxyApp.DisposeAsync().ConfigureAwait(false);
        }
 
        _startLock.Dispose();
    }
 
    internal override ValueTask DisposeAsyncCore() => DisposeAsync();
 
    void UpdateProxyConfig()
    {
        if (_configProvider is null)
        {
            return;
        }
 
        var routes = _instances.Values.Select(i => new RouteConfig
        {
            RouteId = i.Id,
            ClusterId = i.Id,
            Match = new RouteMatch
            {
                Path = "{**catch-all}",
                Headers =
                [
                    new RouteHeader
                    {
                        Name = "X-Test-Backend",
                        Values = [i.Id],
                        Mode = HeaderMatchMode.ExactHeader,
                    }
                ]
            }
        }).ToArray();
 
        var clusters = _instances.Values.Select(i => new ClusterConfig
        {
            ClusterId = i.Id,
            Destinations = new Dictionary<string, DestinationConfig>
            {
                ["default"] = new() { Address = i.AppUrl }
            }
        }).ToArray();
 
        _configProvider.Update(routes, clusters);
    }
 
    static int GetAvailablePort()
    {
        var listener = new TcpListener(IPAddress.Loopback, 0);
        listener.Start();
        var port = ((IPEndPoint)listener.LocalEndpoint).Port;
        listener.Stop();
        return port;
    }
}