using Mbproxy.Options;
using Mbproxy.Proxy.Cache;
using Mbproxy.Proxy.Multiplexing;
using Polly;

namespace Mbproxy.Proxy.Supervision;

/// <summary>
/// Wraps one <see cref="PlcListener"/> in a Polly-backed recovery loop.
///
/// <para><b>State machine</b>:
/// <list type="bullet">
///   <item><description><b>Bound</b>: listener is accepting connections; <see cref="PlcListener.RunAsync"/> is awaiting.</description></item>
///   <item><description><b>Recovering</b>: bind failed or RunAsync faulted; in Polly's delay window before the next attempt.</description></item>
///   <item><description><b>Stopped</b>: terminal. <see cref="StopAsync"/> was called; no further retries.</description></item>
/// </list>
/// </para>
///
/// <para><b>RecoveryAttempts</b>: the counter accumulates over the lifetime of the
/// supervisor. It is never reset after a successful re-bind so operators can see
/// "this listener has flapped N times since the service started." See also
/// <see cref="SupervisorSnapshot"/> doc comment.</para>
///
/// <para>The supervisor does NOT swallow exceptions from <see cref="PlcListener.RunAsync"/>
/// except <see cref="OperationCanceledException"/>. Every other fault is logged at Warning
/// with the exception message so operators can see WHY the listener was restarted.</para>
/// </summary>
internal sealed partial class PlcListenerSupervisor : IAsyncDisposable
{
    private readonly PlcOptions _plc;
    private readonly ConnectionOptions _connectionOptions;
    private readonly IPduPipeline _pipeline;
    private readonly ILogger<PlcListener> _listenerLogger;
    private readonly ILogger<PlcMultiplexer> _multiplexerLogger;
    private readonly ILogger _pipeLogger;
    private readonly PerPlcContext? _perPlcContext;
    private readonly ResiliencePipeline _recoveryPipeline;
    private readonly ILogger<PlcListenerSupervisor> _logger;
    private readonly ResiliencePipeline? _backendConnectPipeline;
    private readonly Func<ReadCoalescingOptions>? _coalescingOptions;
    private readonly Func<KeepaliveOptions>? _keepaliveOptions;

    // ── Mutable state ────────────────────────────────────────────────────────────────────

    // Volatile so Snapshot() reads are coherent without locking.
    private volatile SupervisorState _state = SupervisorState.Stopped;
    private volatile string? _lastBindError;
    private int _recoveryAttempts; // Interlocked

    // Current active listener for status-page pair enumeration.
    private volatile PlcListener? _currentListener;

    // _perPlcContext is mutable so ReplaceContextAsync can swap it. Access from the accept
    // loop (RunAsync) and from ReplaceContextAsync must be coherent; we use a volatile
    // reference so the accept loop always reads the latest context without locking. The
    // PlcListener created on each Polly attempt holds its own copy of the context at
    // construction time; existing in-flight connections keep their old reference until they
    // complete.
    private volatile PerPlcContext? _currentContext;

    /// <summary>
    /// Per-supervisor CTS: cancelling it stops both the Polly delay and the inner
    /// <see cref="PlcListener.RunAsync"/> loop.
    /// </summary>
    private CancellationTokenSource _supervisorCts = new();

    private Task _supervisorTask = Task.CompletedTask;

    private bool _disposed;

    // Completes when the supervisor has transitioned out of Stopped for the first time
    // (reached Bound or Recovering). Used by WaitForInitialBindAttemptAsync to avoid
    // racing fast Stopped→Bound→Stopped transitions or hanging if the supervisor task
    // throws inside Polly.
    //
    // Non-readonly so StartAsync can re-arm it for a re-Started supervisor. Without
    // re-arming, a restart-after-stop scenario would have WaitForInitialBindAttemptAsync
    // return immediately on the previous run's signal, never observing the new run's
    // bind status.
    private TaskCompletionSource _firstAttemptCompleted = new(
        TaskCreationOptions.RunContinuationsAsynchronously);

    // ── Public surface ────────────────────────────────────────────────────────────────────

    public string PlcName => _plc.Name;

    public PlcListenerSupervisor(
        PlcOptions plc,
        ConnectionOptions connectionOptions,
        IPduPipeline pipeline,
        ILogger<PlcListener> listenerLogger,
        ILogger<PlcMultiplexer> multiplexerLogger,
        ILogger pipeLogger,
        PerPlcContext? perPlcContext,
        ResiliencePipeline recoveryPipeline,
        ILogger<PlcListenerSupervisor> logger,
        ResiliencePipeline? backendConnectPipeline = null,
        Func<ReadCoalescingOptions>? coalescingOptions = null,
        Func<KeepaliveOptions>? keepaliveOptions = null)
    {
        _plc                    = plc;
        _connectionOptions      = connectionOptions;
        _pipeline               = pipeline;
        _listenerLogger         = listenerLogger;
        _multiplexerLogger      = multiplexerLogger;
        _pipeLogger             = pipeLogger;
        _perPlcContext          = perPlcContext;
        _currentContext         = perPlcContext;   // live context slot
        _recoveryPipeline       = recoveryPipeline;
        _logger                 = logger;
        _backendConnectPipeline = backendConnectPipeline;
        _coalescingOptions      = coalescingOptions;
        _keepaliveOptions       = keepaliveOptions;
    }

    /// <summary>
    /// Returns the current <see cref="ProxyCounters"/> for this PLC.
    /// Used by <see cref="Configuration.ConfigReconciler"/> when building a reseat context
    /// so that counters are preserved across a tag-map swap.
    /// </summary>
    public ProxyCounters CurrentCounters => _currentContext?.Counters ?? new ProxyCounters();

    /// <summary>
    /// Live collection of active <see cref="UpstreamPipe"/> instances attached to this
    /// PLC's multiplexer. Returns an empty collection when the listener is not bound.
    /// Consumed by the status page.
    /// </summary>
    public IReadOnlyCollection<UpstreamPipe> ActiveUpstreams
        => _currentListener?.ActiveUpstreams ?? Array.Empty<UpstreamPipe>();

    /// <summary>
    /// Launches the supervisor task. The task tries to bind immediately; if binding
    /// fails it enters the Polly recovery loop. The method returns as soon as the
    /// background task is started (it does NOT wait for the listener to reach
    /// <see cref="SupervisorState.Bound"/>).
    ///
    /// <para>Call <see cref="WaitForInitialBindAttemptAsync"/> after this to block until the
    /// supervisor has transitioned out of <see cref="SupervisorState.Stopped"/>.</para>
    /// </summary>
    public Task StartAsync(CancellationToken ct)
    {
        // Refuse to re-Start an already-running or already-disposed supervisor. After
        // Stop the state machine returns to Stopped and StartAsync can re-arm; the per-
        // Start state (CTS, TCS) is refreshed below so no leak or stale signal carries
        // across cycles.
        if (_disposed)
            throw new ObjectDisposedException(nameof(PlcListenerSupervisor));
        if (_state != SupervisorState.Stopped || !_supervisorTask.IsCompleted)
            throw new InvalidOperationException(
                $"Supervisor for Plc='{_plc.Name}' has already been started.");

        // Dispose the previous CTS before reassigning so a re-Start cycle does not leak
        // the prior CTS (and any registrations linked to it). Idempotent: the
        // ObjectDisposed catch covers the very-first-Start case where the field-init CTS
        // is still fresh.
        try { _supervisorCts.Dispose(); } catch (ObjectDisposedException) { /* fresh */ }
        _supervisorCts = CancellationTokenSource.CreateLinkedTokenSource(ct);

        // Re-arm the first-attempt TCS so a re-Started supervisor doesn't immediately
        // observe the previous run's signal in WaitForInitialBindAttemptAsync.
        _firstAttemptCompleted = new TaskCompletionSource(
            TaskCreationOptions.RunContinuationsAsynchronously);

        _supervisorTask = Task.Run(() => RunSupervisorAsync(_supervisorCts.Token), CancellationToken.None);
        return Task.CompletedTask;
    }

    /// <summary>
    /// Waits until the supervisor has completed its first bind attempt
    /// (transitioned to <see cref="SupervisorState.Bound"/> or
    /// <see cref="SupervisorState.Recovering"/>).
    /// Returns immediately if the supervisor is already past that point.
    ///
    /// <para>Backed by a <see cref="TaskCompletionSource"/> set when the supervisor task
    /// first transitions out of <see cref="SupervisorState.Stopped"/>. This avoids both
    /// racing fast bind+stop sequences and hanging if the supervisor task throws before
    /// any state write happens.</para>
    /// </summary>
    public async Task WaitForInitialBindAttemptAsync(CancellationToken ct)
    {
        if (_firstAttemptCompleted.Task.IsCompleted) return;
        try
        {
            await _firstAttemptCompleted.Task.WaitAsync(ct).ConfigureAwait(false);
        }
        catch (OperationCanceledException)
        {
            // Caller cancelled; not a fault.
        }
    }

    /// <summary>
    /// Signals the supervisor to stop, cancels the current Polly delay (if in
    /// <see cref="SupervisorState.Recovering"/>) or the <see cref="PlcListener.RunAsync"/>
    /// loop (if in <see cref="SupervisorState.Bound"/>), and waits for the background
    /// task to complete.
    ///
    /// <para>Completes within ~1 s regardless of backoff window size because Polly's
    /// <c>ExecuteAsync(ct)</c> honours the cancellation token.</para>
    /// </summary>
    public async Task StopAsync(CancellationToken ct)
    {
        _state = SupervisorState.Stopped;

        await _supervisorCts.CancelAsync().ConfigureAwait(false);

        try
        {
            await _supervisorTask.WaitAsync(ct).ConfigureAwait(false);
        }
        catch (OperationCanceledException)
        {
            // ct fired before the task completed — supervisor task will terminate
            // asynchronously. Acceptable at shutdown.
        }
        catch (Exception)
        {
            // Supervisor task faulted — already logged inside RunSupervisorAsync.
        }
    }

    /// <summary>
    /// Returns a point-in-time snapshot of this supervisor's state.
    ///
    /// <para>Reads the three observable fields under a single lock so the status page
    /// can never report inconsistent triples like
    /// <c>(State=Bound, LastBindError=&lt;previous&gt;, RecoveryAttempts&gt;0)</c>. The
    /// supervisor task uses <see cref="TransitionTo"/> which takes the same lock, so a
    /// snapshot reads a transition-consistent view.</para>
    /// </summary>
    public SupervisorSnapshot Snapshot()
    {
        lock (_snapshotLock)
        {
            return new SupervisorSnapshot(
                State:            _state,
                LastBindError:    _lastBindError,
                RecoveryAttempts: _recoveryAttempts);
        }
    }

    private readonly object _snapshotLock = new();

    /// <summary>
    /// Atomic three-field transition. State, lastBindError, and (optionally) the
    /// recoveryAttempts increment all happen under one lock so a concurrent
    /// <see cref="Snapshot"/> never sees a half-applied transition.
    /// </summary>
    private void TransitionTo(SupervisorState newState, string? lastBindError, bool incrementRecoveryAttempt)
    {
        lock (_snapshotLock)
        {
            _state         = newState;
            _lastBindError = lastBindError;
            if (incrementRecoveryAttempt)
                _recoveryAttempts++;
        }
    }

    /// <summary>
    /// Atomically swaps the per-PLC context (tag map + optional response cache) on the
    /// running listener AND its live multiplexer. The swap propagates into the running
    /// mux via <see cref="PlcMultiplexer.ReplaceContext"/>, so the very next PDU sees
    /// the new tag map / new cache. Counters are preserved (the new context carries the
    /// same <c>ProxyCounters</c> instance) so operator history is not reset.
    ///
    /// <para><b>Old cache lifecycle</b>: the supervisor disposes the outgoing context's
    /// cache AFTER the multiplexer has been swapped to the new context. By that point no
    /// more reads or writes can land on the old cache. Per the design contract, any
    /// tag-list change drops the entire PLC cache.</para>
    /// </summary>
    public Task ReplaceContextAsync(PerPlcContext newCtx, CancellationToken ct)
    {
        var oldCache = _currentContext?.Cache;

        // Volatile write: the next PlcListener created in RunSupervisorAsync (on any
        // subsequent fault recovery) will pick up newCtx through this slot.
        _currentContext = newCtx;

        // Push the swap into the running multiplexer so existing connections see the new
        // tag map / new cache on their next PDU. _currentListener may be null between
        // Polly retry attempts; in that case the next listener built inside the Polly loop
        // will pick up newCtx through _currentContext above.
        _currentListener?.Multiplexer?.ReplaceContext(newCtx);

        // Drop the outgoing cache AFTER the swap so the running multiplexer can no longer
        // reach it. Clear() snapshots the entry count for the mbproxy.cache.flushed log
        // event before disposing the cache (which stops the eviction loop and releases
        // the timer).
        if (oldCache is not null && !ReferenceEquals(oldCache, newCtx.Cache))
        {
            int dropped = oldCache.Clear();
            CacheLogEvents.Flushed(_logger, _plc.Name, "tag-list-reload", dropped);
            oldCache.Dispose();
        }

        return Task.CompletedTask;
    }

    // ── Supervisor loop ───────────────────────────────────────────────────────────────────

    private async Task RunSupervisorAsync(CancellationToken ct)
    {
        bool firstBind = true;

        try
        {
            // The recovery pipeline wraps the entire try-bind-and-run block.
            // When RunAsync returns or throws, the pipeline delays and retries.
            // Cancellation of ct exits the pipeline with OperationCanceledException.
            await _recoveryPipeline.ExecuteAsync(async token =>
            {
                // ── Instantiate a fresh listener ─────────────────────────────────
                // A faulted listener's TcpListener socket must be disposed before
                // re-binding. We create a new PlcListener on each attempt.
                //
                // Use _currentContext (volatile) so that a ReplaceContextAsync call
                // between Polly retry attempts is picked up here. Each listener captures
                // the context at construction time; existing in-flight pairs keep their
                // own reference. See ReplaceContextAsync for the transition window
                // documentation.
                var listener = new PlcListener(
                    _plc,
                    _connectionOptions,
                    _pipeline,
                    _listenerLogger,
                    _multiplexerLogger,
                    _pipeLogger,
                    _currentContext,
                    _backendConnectPipeline,
                    _coalescingOptions,
                    _keepaliveOptions);

                // Expose the current listener for status-page pair enumeration.
                _currentListener = listener;

                try
                {
                    // ── Bind ─────────────────────────────────────────────────────
                    listener.StartAsync();
                }
                catch (Exception bindEx)
                {
                    // Dispose the listener before entering the recovery delay
                    // so the socket is released and the port can be reused.
                    _currentListener = null;
                    await listener.DisposeAsync().ConfigureAwait(false);

                    string truncated = Truncate(bindEx.Message, 256);
                    TransitionTo(SupervisorState.Recovering, truncated, incrementRecoveryAttempt: true);
                    // Signal the first transition out of Stopped.
                    _firstAttemptCompleted.TrySetResult();

                    // Also update the per-PLC counters if available (status page reads these).
                    _currentContext?.Counters.IncrementRecoveryAttempt(truncated);

                    LogBindFailed(_logger, _plc.Name, _plc.ListenPort, truncated);

                    // Re-throw so the Polly pipeline can delay and retry.
                    throw;
                }

                // ── Bind succeeded ───────────────────────────────────────────────
                if (firstBind)
                {
                    firstBind = false;
                    LogBound(_logger, _plc.Name, _plc.ListenPort);
                }
                else
                {
                    // Re-bind after a recovery — emit the "recovered" event once.
                    int totalAttempts = Interlocked.CompareExchange(ref _recoveryAttempts, 0, 0);
                    LogListenerRecovered(_logger, _plc.Name, _plc.ListenPort, totalAttempts);
                }

                // Clear the last bind error on a successful bind.
                TransitionTo(SupervisorState.Bound, lastBindError: null, incrementRecoveryAttempt: false);
                _currentContext?.Counters.ClearLastBindError();
                // Signal the first transition out of Stopped.
                _firstAttemptCompleted.TrySetResult();

                // ── Run the accept loop ──────────────────────────────────────────
                // RunAsync returns when: (a) token is cancelled (normal shutdown),
                // (b) the listener faults (OS reclaims port, transient network reset).
                // In both cases we fall through to the Polly retry handler.
                try
                {
                    await listener.RunAsync(token).ConfigureAwait(false);
                }
                catch (OperationCanceledException)
                {
                    // Normal shutdown path — do not enter recovery loop.
                    _currentListener = null;
                    await listener.DisposeAsync().ConfigureAwait(false);
                    throw; // Propagate to exit the Polly pipeline.
                }
                catch (Exception runEx)
                {
                    // Listener faulted at runtime (port stolen, OS network reset, etc.).
                    // Log at Warning — operators must see WHY the listener was restarted.
                    LogListenerFaulted(_logger, _plc.Name, _plc.ListenPort, runEx, runEx.Message);
                    _currentListener = null;
                    await listener.DisposeAsync().ConfigureAwait(false);

                    string truncated = Truncate(runEx.Message, 256);
                    TransitionTo(SupervisorState.Recovering, truncated, incrementRecoveryAttempt: true);
                    // Also signal first-attempt-completed in case the very first
                    // listener.RunAsync faulted before the bind-success path signalled it.
                    _firstAttemptCompleted.TrySetResult();

                    // Also update the per-PLC counters if available.
                    _currentContext?.Counters.IncrementRecoveryAttempt(truncated);

                    // Re-throw so Polly can delay and retry.
                    throw;
                }

                // RunAsync returned normally (token was cancelled or listener closed).
                // If we got here without an exception, the loop ended cleanly.
                _currentListener = null;
                await listener.DisposeAsync().ConfigureAwait(false);

                // If cancellation is requested, throw so Polly exits cleanly.
                token.ThrowIfCancellationRequested();

                // Otherwise (listener closed without cancellation — e.g., OS event),
                // treat as a fault and re-enter recovery.
                const string unexpectedEnd = "Listener accept loop ended unexpectedly";
                TransitionTo(SupervisorState.Recovering, unexpectedEnd, incrementRecoveryAttempt: true);
                _currentContext?.Counters.IncrementRecoveryAttempt(unexpectedEnd);
                LogListenerEnded(_logger, _plc.Name, _plc.ListenPort);
                throw new InvalidOperationException(unexpectedEnd);

            }, ct).ConfigureAwait(false);
        }
        catch (OperationCanceledException)
        {
            // Normal: StopAsync cancelled the token.
        }
        catch (Exception ex)
        {
            // Polly pipeline exhausted (should not happen for listener recovery since
            // MaxRetryAttempts = int.MaxValue) or an unexpected fault.
            _logger.LogError(ex, "Supervisor for Plc={Plc} exited unexpectedly: {Message}",
                _plc.Name, ex.Message);
        }
        finally
        {
            // Snapshot consistency: state goes back to Stopped without changing the last
            // bind error so operators can still see WHY the supervisor exited.
            lock (_snapshotLock)
            {
                _state = SupervisorState.Stopped;
            }
            _currentListener = null;
            // Defensive: if RunSupervisorAsync exits before any bind attempt fired
            // (e.g. construction-time fault), unblock any awaiting
            // WaitForInitialBindAttemptAsync caller so it doesn't hang.
            _firstAttemptCompleted.TrySetResult();
        }
    }

    /// <summary>
    /// Single helper for the truncate-exception-message pattern shared across the
    /// supervisor's bind/run/end recovery paths.
    /// </summary>
    private static string Truncate(string s, int max) => s.Length > max ? s[..max] : s;

    // ── IAsyncDisposable ─────────────────────────────────────────────────────────────────

    public async ValueTask DisposeAsync()
    {
        if (_disposed) return;
        _disposed = true;

        using var stopCts = new CancellationTokenSource(TimeSpan.FromSeconds(5));
        try
        {
            await StopAsync(stopCts.Token).ConfigureAwait(false);
        }
        catch
        {
            // Best-effort cleanup.
        }

        // Dispose the response cache (if any) — its eviction timer would otherwise
        // outlive the supervisor.
        _currentContext?.Cache?.Dispose();

        _supervisorCts.Dispose();
    }

    // ── Logging ───────────────────────────────────────────────────────────────────────────

    [LoggerMessage(EventId = 40, EventName = "mbproxy.startup.bind",
        Level = LogLevel.Information,
        Message = "Listener bound: Plc={Plc} Port={Port}")]
    private static partial void LogBound(ILogger logger, string plc, int port);

    [LoggerMessage(EventId = 41, EventName = "mbproxy.startup.bind.failed",
        Level = LogLevel.Error,
        Message = "Failed to bind listener: Plc={Plc} Port={Port} Reason={Reason}")]
    private static partial void LogBindFailed(ILogger logger, string plc, int port, string reason);

    [LoggerMessage(EventId = 42, EventName = "mbproxy.listener.recovered",
        Level = LogLevel.Information,
        Message = "Listener recovered: Plc={Plc} Port={Port} AttemptCount={AttemptCount}")]
    private static partial void LogListenerRecovered(ILogger logger, string plc, int port, int attemptCount);

    [LoggerMessage(EventId = 43, EventName = "mbproxy.listener.faulted",
        Level = LogLevel.Warning,
        Message = "Listener faulted (will recover): Plc={Plc} Port={Port} Reason={Reason}")]
    private static partial void LogListenerFaulted(ILogger logger, string plc, int port, Exception ex, string reason);

    [LoggerMessage(EventId = 44, EventName = "mbproxy.listener.ended",
        Level = LogLevel.Warning,
        Message = "Listener accept loop ended unexpectedly (will recover): Plc={Plc} Port={Port}")]
    private static partial void LogListenerEnded(ILogger logger, string plc, int port);
}
