using Microsoft.Extensions.Logging;
using Microsoft.Extensions.Logging.Abstractions;

namespace ZB.MOM.WW.OtOpcUa.Driver.Galaxy.Runtime;

/// <summary>
///     Coordinates GalaxyDriver's recovery from gateway transport failure. Drives a
///     state machine — <c>Healthy → TransportLost → Reopening → Replaying → Healthy</c>
///     — and exposes the current state through a snapshot + change event so the
///     driver's <c>DriverHealth</c> reflects <c>Degraded</c> while we're not in
///     <c>Healthy</c>.
/// </summary>
/// <remarks>
///     <para>
///         The supervisor doesn't own the session, the subscription registry, or the
///         event pump. It receives transport-failure signals from the rest of the
///         driver (EventPump throws, a gw RPC raises, the heartbeat times out), runs
///         a one-attempt-at-a-time recovery loop, and lets the rest of the driver
///         continue serving cached state during recovery.
///     </para>
///     <para>
///         <b>Reopen</b>: caller-supplied callback that re-opens the gw session +
///         re-Registers the MXAccess client. Throws on failure.
///     </para>
///     <para>
///         <b>Replay</b>: caller-supplied callback that re-establishes every active
///         subscription. Production wraps gw's <c>ReplaySubscriptionsCommand</c>
///         (tracked in the mxaccessgw backlog); when that's not available, the callback falls
///         back to walking the SubscriptionRegistry and re-issuing SubscribeBulk for
///         every tracked tag.
///     </para>
///     <para>
///         Backoff is capped exponential — first retry after
///         <see cref="ReconnectOptions.InitialBackoff"/>, doubled per failed attempt,
///         capped at <see cref="ReconnectOptions.MaxBackoff"/>. Persistent failures
///         hold the supervisor in <c>Reopening</c> indefinitely; the supervisor never
///         gives up on its own — operators / Phase 6.4 soak handle that policy.
///     </para>
/// </remarks>
public sealed class ReconnectSupervisor : IDisposable
{
    /// <summary>Recovery state machine.</summary>
    public enum State
    {
        Healthy,
        TransportLost,
        Reopening,
        Replaying,
    }

    private readonly Func<CancellationToken, Task> _reopen;
    private readonly Func<CancellationToken, Task> _replay;
    private readonly ReconnectOptions _options;
    private readonly ILogger _logger;
    private readonly Func<int, TimeSpan, TimeSpan, TimeSpan>? _backoffDelay;

    private readonly Lock _stateLock = new();
    private State _state = State.Healthy;
    private string? _lastError;
    private DateTime? _lastTransitionUtc;

    private Task? _recoveryLoop;
    private CancellationTokenSource? _loopCts;
    private bool _disposed;
    private int _cyclesStarted;
    private int _cyclesReleased;

    /// <summary>Recovery cycles started so far (one per <see cref="ReportTransportFailure"/> that left Healthy). Test seam.</summary>
    internal int CyclesStarted => Volatile.Read(ref _cyclesStarted);

    /// <summary>Recovery cycles whose <see cref="CancellationTokenSource"/> has been disposed. Equals
    /// <see cref="CyclesStarted"/> once every loop has exited — a cycle's CTS used to be released only
    /// if it happened to be the LAST one when <see cref="Dispose"/> ran. Test seam.</summary>
    internal int CyclesReleased => Volatile.Read(ref _cyclesReleased);

    /// <summary>Fires after every state transition.</summary>
    public event EventHandler<StateTransition>? StateChanged;

    /// <summary>
    /// Initializes a new ReconnectSupervisor with the specified recovery handlers and options.
    /// </summary>
    /// <param name="reopen">Callback to reopen the transport connection.</param>
    /// <param name="replay">Callback to replay subscriptions and state after reopening.</param>
    /// <param name="options">Optional reconnection behavior configuration.</param>
    /// <param name="logger">Optional logger instance; defaults to null logger if not provided.</param>
    /// <param name="backoffDelay">Optional custom backoff delay calculator; uses default if not provided.</param>
    public ReconnectSupervisor(
        Func<CancellationToken, Task> reopen,
        Func<CancellationToken, Task> replay,
        ReconnectOptions? options = null,
        ILogger? logger = null,
        Func<int, TimeSpan, TimeSpan, TimeSpan>? backoffDelay = null)
    {
        _reopen = reopen ?? throw new ArgumentNullException(nameof(reopen));
        _replay = replay ?? throw new ArgumentNullException(nameof(replay));
        _options = options ?? new ReconnectOptions();
        _logger = logger ?? NullLogger.Instance;
        _backoffDelay = backoffDelay;
    }

    /// <summary>Current state. Healthy = fully recovered + subscriptions live.</summary>
    public State CurrentState
    {
        get { lock (_stateLock) return _state; }
    }

    /// <summary>True when CurrentState != Healthy. Drivers map this to DriverState.Degraded.</summary>
    public bool IsDegraded
    {
        get { lock (_stateLock) return _state != State.Healthy; }
    }

    /// <summary>Gets the message of the last reported transport error, if any.</summary>
    public string? LastError
    {
        get { lock (_stateLock) return _lastError; }
    }

    /// <summary>Gets the UTC timestamp of the most recent state transition.</summary>
    public DateTime? LastTransitionUtc
    {
        get { lock (_stateLock) return _lastTransitionUtc; }
    }

    /// <summary>
    ///     Notify the supervisor that a gw transport failure has been observed. Idempotent —
    ///     repeated calls during an in-flight recovery do not start a parallel loop. The
    ///     first call spawns a background task that drives reopen → replay until it
    ///     succeeds or <see cref="Dispose"/> cancels it.
    /// </summary>
    /// <param name="cause">The exception that triggered the failure notification.</param>
    public void ReportTransportFailure(Exception cause)
    {
        ArgumentNullException.ThrowIfNull(cause);
        ObjectDisposedException.ThrowIf(_disposed, this);

        lock (_stateLock)
        {
            _lastError = cause.Message;
            if (_state != State.Healthy)
            {
                // Already recovering — nothing else to do.
                _logger.LogDebug("Transport failure reported during {State}: {Message}", _state, cause.Message);
                return;
            }

            TransitionLocked(State.TransportLost, cause.Message);

            // Each cycle owns its CTS and releases it when its loop exits (see RecoveryLoopAsync's
            // finally). Before, every cycle allocated a fresh CTS and only the last one was ever
            // disposed — one leaked CancellationTokenSource per outage for the driver's lifetime.
            var cts = new CancellationTokenSource();
            _loopCts = cts;
            Interlocked.Increment(ref _cyclesStarted);
            _recoveryLoop = Task.Run(() => RecoveryLoopAsync(cts));
        }
    }

    /// <summary>
    ///     Wait until the current recovery cycle reaches Healthy or the supplied token
    ///     is cancelled. Returns immediately when already Healthy. Useful for tests
    ///     and for orchestration that wants to gate calls on recovery completing.
    /// </summary>
    /// <param name="cancellationToken">Token to cancel the wait operation.</param>
    /// <returns>A task that represents the asynchronous operation.</returns>
    public async Task WaitForHealthyAsync(CancellationToken cancellationToken)
    {
        while (!cancellationToken.IsCancellationRequested && IsDegraded)
        {
            await Task.Delay(50, cancellationToken).ConfigureAwait(false);
        }
    }

    private async Task RecoveryLoopAsync(CancellationTokenSource cycleCts)
    {
        try
        {
            await RecoveryLoopCoreAsync(cycleCts.Token).ConfigureAwait(false);
        }
        finally
        {
            // Release this cycle's CTS once nothing can observe its token any more. Stop/Dispose may
            // race us here: they Cancel (guarded against ODE) and Dispose (idempotent) the same object.
            lock (_stateLock)
            {
                if (ReferenceEquals(_loopCts, cycleCts)) _loopCts = null;
            }
            cycleCts.Dispose();
            Interlocked.Increment(ref _cyclesReleased);
        }
    }

    private async Task RecoveryLoopCoreAsync(CancellationToken ct)
    {
        var attempt = 0;
        while (!ct.IsCancellationRequested)
        {
            attempt++;
            if (attempt > 1)
            {
                var delay = ComputeBackoff(attempt);
                _logger.LogInformation(
                    "Galaxy reconnect attempt {Attempt} — waiting {Delay} before retry", attempt, delay);
                try { await Task.Delay(delay, ct).ConfigureAwait(false); }
                catch (OperationCanceledException) { return; }
            }

            // === Reopening phase ===
            lock (_stateLock) TransitionLocked(State.Reopening, _lastError);

            try
            {
                await _reopen(ct).ConfigureAwait(false);
            }
            catch (OperationCanceledException) when (ct.IsCancellationRequested) { return; }
            catch (Exception ex)
            {
                _logger.LogWarning(ex, "Galaxy reopen failed (attempt {Attempt})", attempt);
                lock (_stateLock) { _lastError = ex.Message; }
                continue; // back to backoff + retry
            }

            // === Replaying phase ===
            lock (_stateLock) TransitionLocked(State.Replaying, _lastError);

            try
            {
                await _replay(ct).ConfigureAwait(false);
            }
            catch (OperationCanceledException) when (ct.IsCancellationRequested) { return; }
            catch (Exception ex)
            {
                _logger.LogWarning(ex, "Galaxy replay failed (attempt {Attempt})", attempt);
                lock (_stateLock) { _lastError = ex.Message; }
                continue; // back to backoff + retry
            }

            // === Done ===
            lock (_stateLock)
            {
                _lastError = null;
                TransitionLocked(State.Healthy, null);
            }
            _logger.LogInformation("Galaxy reconnect succeeded after {Attempt} attempt(s)", attempt);
            return;
        }
    }

    private TimeSpan ComputeBackoff(int attempt)
    {
        if (_backoffDelay is not null)
            return _backoffDelay(attempt, _options.InitialBackoff, _options.MaxBackoff);

        // Standard capped exponential — InitialBackoff * 2^(attempt-2), capped at MaxBackoff.
        // Attempt 2 → InitialBackoff, attempt 3 → 2x, attempt 4 → 4x, etc.
        var multiplier = Math.Min(1L << Math.Max(0, attempt - 2), int.MaxValue);
        var ticks = _options.InitialBackoff.Ticks * multiplier;
        if (ticks <= 0 || ticks > _options.MaxBackoff.Ticks) ticks = _options.MaxBackoff.Ticks;
        return TimeSpan.FromTicks(ticks);
    }

    private void TransitionLocked(State next, string? cause)
    {
        if (next == _state) return;
        var previous = _state;
        _state = next;
        _lastTransitionUtc = DateTime.UtcNow;
        var transition = new StateTransition(previous, next, cause, _lastTransitionUtc.Value);
        try { StateChanged?.Invoke(this, transition); }
        catch (Exception ex)
        {
            _logger.LogWarning(ex,
                "Galaxy reconnect StateChanged handler threw — continuing.");
        }
    }

    /// <summary>
    ///     Cancels any in-flight recovery loop and waits for it to exit, without blocking the caller's
    ///     thread. The async counterpart of <see cref="Dispose"/>; the driver's <c>ShutdownAsync</c>
    ///     uses it so a loop parked inside a reopen round-trip is awaited, not blocked on. Idempotent.
    /// </summary>
    /// <returns>A task that completes once the recovery loop has exited.</returns>
    public async Task StopAsync()
    {
        if (_disposed) return;
        _disposed = true;

        CancellationTokenSource? cts;
        Task? loop;
        lock (_stateLock) { cts = _loopCts; loop = _recoveryLoop; _loopCts = null; _recoveryLoop = null; }

        // The loop's own finally may have released the CTS already — a cancelled-then-disposed CTS
        // throws on Cancel, and a disposed one is a no-op on Dispose.
        try { cts?.Cancel(); } catch (ObjectDisposedException) { /* loop already released it */ }
        if (loop is not null)
        {
            try { await loop.ConfigureAwait(false); } catch { /* shutdown */ }
        }
        cts?.Dispose();
    }

    /// <summary>Disposes the supervisor and cancels any in-flight recovery loop (blocking until it exits).</summary>
    public void Dispose() => StopAsync().GetAwaiter().GetResult();
}

/// <summary>
///     One state transition observed by the supervisor.
/// </summary>
public sealed record StateTransition(
    ReconnectSupervisor.State Previous,
    ReconnectSupervisor.State Next,
    string? Cause,
    DateTime AtUtc);

/// <summary>
///     Knobs for the supervisor's backoff. <see cref="ReconnectOptions"/> on the driver
///     options record (PR 4.0) maps onto this — they're separate types so the supervisor
///     can be exercised in tests without the full driver options surface.
/// </summary>
public sealed record ReconnectOptions(
    TimeSpan? InitialBackoffOverride = null,
    TimeSpan? MaxBackoffOverride = null)
{
    /// <summary>Gets the initial backoff delay for reconnection attempts.</summary>
    public TimeSpan InitialBackoff => InitialBackoffOverride ?? TimeSpan.FromMilliseconds(500);

    /// <summary>Gets the maximum backoff delay for reconnection attempts.</summary>
    public TimeSpan MaxBackoff => MaxBackoffOverride ?? TimeSpan.FromSeconds(30);
}