using System.Diagnostics;
using Mbproxy.Admin;
using Mbproxy.Options;
using Mbproxy.Proxy;
using Mbproxy.Proxy.Multiplexing;
using Mbproxy.Proxy.Supervision;
using Microsoft.Extensions.Options;

namespace Mbproxy.Diagnostics;

// ── Testability interfaces ────────────────────────────────────────────────────────────────────

/// <summary>
/// Abstraction over a supervisor's stop operation and its multiplexer's in-flight count.
/// Introduced so <see cref="ShutdownCoordinator"/> unit tests can inject fakes
/// without needing a real <see cref="PlcListenerSupervisor"/>.
///
/// <para><b>Phase 9:</b> in-flight tracking is now per-multiplexer (the
/// <see cref="CorrelationMap"/>) rather than per-pair. <see cref="InFlightCount"/>
/// replaces <c>ActivePairs.IsProcessing</c> from the 1:1 model.</para>
/// </summary>
internal interface ISupervisorHandle
{
    Task StopAsync(CancellationToken ct);

    /// <summary>
    /// Current number of in-flight Modbus requests on this PLC's multiplexed backend.
    /// Zero if the multiplexer has no in-flight requests (idle).
    /// </summary>
    int InFlightCount { get; }
}

/// <summary>
/// Abstraction over the admin endpoint stop operation.
/// </summary>
internal interface IAdminEndpointHandle
{
    Task StopAsync(CancellationToken ct);
}

/// <summary>
/// Adapts a concrete <see cref="PlcListenerSupervisor"/> to <see cref="ISupervisorHandle"/>.
/// </summary>
internal sealed class PlcSupervisorHandle : ISupervisorHandle
{
    private readonly PlcListenerSupervisor _supervisor;
    public PlcSupervisorHandle(PlcListenerSupervisor supervisor) => _supervisor = supervisor;
    public Task StopAsync(CancellationToken ct) => _supervisor.StopAsync(ct);

    public int InFlightCount
    {
        get
        {
            // CurrentCounters.Snapshot pulls live values from the multiplexer's
            // IMultiplexCountersProvider hook; InFlightCount is point-in-time.
            return (int)_supervisor.CurrentCounters.Snapshot().InFlightCount;
        }
    }
}

/// <summary>
/// Adapts <see cref="AdminEndpointHost"/> to <see cref="IAdminEndpointHandle"/>.
/// </summary>
internal sealed class AdminEndpointHandle : IAdminEndpointHandle
{
    private readonly AdminEndpointHost _host;
    public AdminEndpointHandle(AdminEndpointHost host) => _host = host;
    public Task StopAsync(CancellationToken ct) => _host.StopAsync(ct);
}

// ── ShutdownCoordinator ───────────────────────────────────────────────────────────────────────

/// <summary>
/// Orchestrates graceful shutdown of the proxy service.
///
/// <para>Shutdown sequence:</para>
/// <list type="number">
///   <item>Stop accepting new upstream connections on all supervisors.</item>
///   <item>Wait for in-flight Modbus requests to drain (polls
///         <see cref="ISupervisorHandle.InFlightCount"/> across all supervisors) until
///         <see cref="ConnectionOptions.GracefulShutdownTimeoutMs"/> expires.</item>
///   <item>Stop the admin endpoint.</item>
///   <item>Log <c>mbproxy.shutdown.complete</c> with <c>InFlightAtCancel</c> and <c>ElapsedMs</c>.</item>
/// </list>
///
/// <para>This type is internal. It is registered in DI as a singleton and wired to
/// <see cref="IHostApplicationLifetime.ApplicationStopping"/> in <c>Program.cs</c>.</para>
/// </summary>
internal sealed partial class ShutdownCoordinator
{
    private readonly IReadOnlyList<ISupervisorHandle> _supervisors;
    private readonly IAdminEndpointHandle _adminEndpoint;
    private readonly IOptions<MbproxyOptions> _options;
    private readonly ILogger<ShutdownCoordinator> _logger;

    /// <summary>
    /// Production constructor — wraps concrete types in their adapter handles.
    /// </summary>
    public ShutdownCoordinator(
        IEnumerable<PlcListenerSupervisor> supervisors,
        AdminEndpointHost adminEndpoint,
        IOptions<MbproxyOptions> options,
        ILogger<ShutdownCoordinator> logger)
        : this(
            supervisors.Select(s => (ISupervisorHandle)new PlcSupervisorHandle(s)).ToList(),
            new AdminEndpointHandle(adminEndpoint),
            options,
            logger)
    {
    }

    /// <summary>
    /// Testability constructor — accepts abstractions so unit tests can inject fakes.
    /// </summary>
    internal ShutdownCoordinator(
        IReadOnlyList<ISupervisorHandle> supervisors,
        IAdminEndpointHandle adminEndpoint,
        IOptions<MbproxyOptions> options,
        ILogger<ShutdownCoordinator> logger)
    {
        _supervisors   = supervisors;
        _adminEndpoint = adminEndpoint;
        _options       = options;
        _logger        = logger;
    }

    /// <summary>
    /// Runs the graceful shutdown sequence.
    /// </summary>
    /// <param name="timeoutMs">
    /// Override the configured <c>Connection.GracefulShutdownTimeoutMs</c> (use -1 to
    /// read from options, which is the normal runtime path). Tests pass an explicit value.
    /// </param>
    /// <param name="hostCt">
    /// The host lifetime cancellation token. Not used to gate the drain loop — the
    /// coordinator manages its own deadline so it can log completion regardless.
    /// </param>
    public async Task ShutdownAsync(int timeoutMs = -1, CancellationToken hostCt = default)
    {
        int deadline = timeoutMs >= 0
            ? timeoutMs
            : _options.Value.Connection.GracefulShutdownTimeoutMs;

        var sw = Stopwatch.StartNew();

        // ── Step 1: stop accepting new connections ────────────────────────────────────
        using var stopCts = new CancellationTokenSource(TimeSpan.FromSeconds(5));
        var stopTasks = _supervisors
            .Select(s => s.StopAsync(stopCts.Token))
            .ToArray();

        try
        {
            await Task.WhenAll(stopTasks).ConfigureAwait(false);
        }
        catch
        {
            // Best-effort: individual supervisor failures must not abort shutdown.
        }

        // ── Step 2: wait for in-flight PDUs to drain ──────────────────────────────────
        int inFlightAtCancel = 0;

        using var drainCts = new CancellationTokenSource(TimeSpan.FromMilliseconds(deadline));
        try
        {
            while (!drainCts.Token.IsCancellationRequested)
            {
                int inFlight = CountInFlight(_supervisors);
                if (inFlight == 0) break;

                await Task.Delay(10, drainCts.Token).ConfigureAwait(false);
            }
        }
        catch (OperationCanceledException)
        {
            // Deadline expired — count remaining in-flight and proceed.
            inFlightAtCancel = CountInFlight(_supervisors);
        }

        // ── Step 3: stop the admin endpoint ──────────────────────────────────────────
        // Admin is stopped AFTER listeners to preserve ordering guarantee:
        // supervisors stop → drain → admin stops.
        try
        {
            using var adminCts = new CancellationTokenSource(TimeSpan.FromSeconds(2));
            await _adminEndpoint.StopAsync(adminCts.Token).ConfigureAwait(false);
        }
        catch
        {
            // Best-effort.
        }

        // ── Step 4: log completion ────────────────────────────────────────────────────
        LogShutdownComplete(_logger, inFlightAtCancel, sw.ElapsedMilliseconds);
    }

    private static int CountInFlight(IReadOnlyList<ISupervisorHandle> supervisors)
    {
        int count = 0;
        foreach (var supervisor in supervisors)
        {
            count += supervisor.InFlightCount;
        }
        return count;
    }

    [LoggerMessage(EventId = 80, EventName = "mbproxy.shutdown.complete",
        Level = LogLevel.Information,
        Message = "Graceful shutdown complete: InFlightAtCancel={InFlightAtCancel} ElapsedMs={ElapsedMs}")]
    private static partial void LogShutdownComplete(ILogger logger, int inFlightAtCancel, long elapsedMs);
}
