ScadaBridge/src/ZB.MOM.WW.ScadaBridge.Communication/Grpc/SitePairChannelProvider.cs

using System.Collections.Concurrent;
using Grpc.Core;
using Grpc.Net.Client;
using Microsoft.Extensions.Logging;
using Microsoft.Extensions.Options;

namespace ZB.MOM.WW.ScadaBridge.Communication.Grpc;

/// <summary>
/// Raised when a site has no usable gRPC channel — an unknown site, or a site with neither
/// <c>GrpcNodeAAddress</c> nor <c>GrpcNodeBAddress</c> configured. The gRPC transport treats this
/// the way the Akka path treats "no ClusterClient for site": warn and drop, so the caller's Ask
/// times out (central never buffers).
/// </summary>
public sealed class SiteChannelUnavailableException(string siteId)
    : Exception($"No gRPC channel is configured for site '{siteId}'.")
{
    /// <summary>The site with no configured channel.</summary>
    public string SiteId { get; } = siteId;
}

/// <summary>
/// Shared, per-site gRPC channel PAIR provider with sticky failover + failback (design §3.5).
/// Central holds one <see cref="GrpcChannel"/> per site node (NodeA/NodeB) rather than one channel
/// with re-dial logic; NodeA is the preferred node (config order). Every call goes to the current
/// sticky channel; on connect failure / <see cref="StatusCode.Unavailable"/> / readiness rejection
/// it flips to the other node and stays there. A background probe returns to the preferred node
/// once it answers again. Reconnect probing backs off 1 s → doubling → 60 s cap.
/// </summary>
/// <remarks>
/// <para>
/// <b>Retry safety.</b> <see cref="ExecuteAsync{T}"/> retries on the other node ONLY when the first
/// attempt failed with <see cref="StatusCode.Unavailable"/> (provably never reached a server —
/// connect refused or readiness-rejected before response headers). A
/// <see cref="StatusCode.DeadlineExceeded"/>, or any other status, is rethrown WITHOUT a cross-node
/// retry: a <c>WriteTag</c>/<c>DeployArtifacts</c>/<c>TriggerSiteFailover</c> may already have
/// executed, and the layers above tolerate the one-shot ambiguity exactly as ClusterClient's Ask
/// did — the migration must not turn one write into two.
/// </para>
/// <para>
/// <b>Addresses.</b> Fed by <see cref="UpdateSite"/>/<see cref="RemoveSite"/> from the single DB
/// refresh loop (the <c>GrpcNodeAAddress</c>/<c>GrpcNodeBAddress</c> columns). PSK credentials and
/// the <c>x-scadabridge-site</c> header ride every channel via
/// <see cref="ControlPlaneCredentials.WithSiteCredentials"/>; removing a site invalidates its
/// cached key.
/// </para>
/// </remarks>
public sealed class SitePairChannelProvider : IDisposable
{
    private static readonly TimeSpan InitialBackoff = TimeSpan.FromSeconds(1);
    private static readonly TimeSpan MaxBackoff = TimeSpan.FromSeconds(60);
    private static readonly TimeSpan FailbackSweepInterval = TimeSpan.FromSeconds(1);

    private readonly ISitePskProvider _pskProvider;
    private readonly CommunicationOptions _options;
    private readonly ILogger<SitePairChannelProvider> _logger;
    private readonly Func<string, HttpMessageHandler>? _handlerFactory;
    private readonly Func<GrpcChannel, CancellationToken, Task<bool>> _reachabilityProbe;

    private readonly ConcurrentDictionary<string, SitePair> _sites = new(StringComparer.Ordinal);
    private readonly CancellationTokenSource _shutdown = new();
    private Task? _failbackLoop;
    private readonly object _loopGate = new();
    private bool _disposed;

    /// <summary>Creates the provider.</summary>
    /// <param name="pskProvider">Resolves each site's preshared key for channel credentials.</param>
    /// <param name="options">Communication options (keepalive settings applied to production channels).</param>
    /// <param name="logger">Logger for failover/failback diagnostics.</param>
    /// <param name="handlerFactory">
    /// Test seam mapping an endpoint to the <see cref="HttpMessageHandler"/> its channel should use
    /// (e.g. an in-process <c>TestServer</c> handler). Null in production, where each channel builds
    /// a keepalive-configured <see cref="SocketsHttpHandler"/>.
    /// </param>
    /// <param name="reachabilityProbe">
    /// Test seam deciding whether a preferred node is reachable during a failback sweep. Null in
    /// production, where <see cref="GrpcChannel.ConnectAsync"/> is used.
    /// </param>
    public SitePairChannelProvider(
        ISitePskProvider pskProvider,
        IOptions<CommunicationOptions> options,
        ILogger<SitePairChannelProvider> logger,
        Func<string, HttpMessageHandler>? handlerFactory = null,
        Func<GrpcChannel, CancellationToken, Task<bool>>? reachabilityProbe = null)
    {
        ArgumentNullException.ThrowIfNull(pskProvider);
        ArgumentNullException.ThrowIfNull(options);
        ArgumentNullException.ThrowIfNull(logger);

        _pskProvider = pskProvider;
        _options = options.Value;
        _logger = logger;
        _handlerFactory = handlerFactory;
        _reachabilityProbe = reachabilityProbe ?? DefaultReachabilityProbe;
    }

    /// <summary>
    /// (Re)builds a site's channel pair from its gRPC node endpoints. A no-op when neither endpoint
    /// changed. Disposes and rebuilds a channel whose endpoint changed, and resets stickiness to the
    /// preferred node when the pair is (re)created. Called on the DB refresh tick.
    /// </summary>
    /// <param name="siteId">Site identifier.</param>
    /// <param name="nodeAEndpoint">NodeA gRPC base address, or null when unconfigured.</param>
    /// <param name="nodeBEndpoint">NodeB gRPC base address, or null when unconfigured.</param>
    public void UpdateSite(string siteId, string? nodeAEndpoint, string? nodeBEndpoint)
    {
        ArgumentException.ThrowIfNullOrWhiteSpace(siteId);

        var pair = _sites.GetOrAdd(siteId, id => new SitePair(id));
        lock (pair.Gate)
        {
            var changedA = !string.Equals(pair.NodeAEndpoint, nodeAEndpoint, StringComparison.Ordinal);
            var changedB = !string.Equals(pair.NodeBEndpoint, nodeBEndpoint, StringComparison.Ordinal);
            if (!changedA && !changedB)
            {
                return;
            }

            if (changedA)
            {
                pair.ChannelA?.Dispose();
                pair.ChannelA = string.IsNullOrWhiteSpace(nodeAEndpoint) ? null : BuildChannel(nodeAEndpoint, siteId);
                pair.NodeAEndpoint = nodeAEndpoint;
            }
            if (changedB)
            {
                pair.ChannelB?.Dispose();
                pair.ChannelB = string.IsNullOrWhiteSpace(nodeBEndpoint) ? null : BuildChannel(nodeBEndpoint, siteId);
                pair.NodeBEndpoint = nodeBEndpoint;
            }

            // Addresses changed — return to the preferred node and reset failback backoff.
            pair.CurrentIsA = pair.ChannelA is not null;
            pair.ResetFailback();
            _logger.LogInformation(
                "gRPC channel pair for site {SiteId} refreshed (nodeA={HasA}, nodeB={HasB})",
                siteId, pair.ChannelA is not null, pair.ChannelB is not null);
        }
    }

    /// <summary>Disposes a removed site's channels and invalidates its cached preshared key.</summary>
    /// <param name="siteId">Site identifier to remove.</param>
    public void RemoveSite(string siteId)
    {
        if (string.IsNullOrWhiteSpace(siteId))
        {
            return;
        }

        if (_sites.TryRemove(siteId, out var pair))
        {
            lock (pair.Gate)
            {
                pair.ChannelA?.Dispose();
                pair.ChannelB?.Dispose();
                pair.ChannelA = null;
                pair.ChannelB = null;
            }
            _pskProvider.Invalidate(siteId);
            _logger.LogInformation("gRPC channel pair for site {SiteId} removed", siteId);
        }
    }

    /// <summary>
    /// Runs <paramref name="call"/> against the site's current sticky channel, failing over to the
    /// other node once on <see cref="StatusCode.Unavailable"/> (never on
    /// <see cref="StatusCode.DeadlineExceeded"/>). Ensures the background failback loop is running.
    /// </summary>
    /// <typeparam name="T">The RPC reply type.</typeparam>
    /// <param name="siteId">Target site.</param>
    /// <param name="call">The RPC to run against a given channel.</param>
    /// <param name="ct">Cancellation token.</param>
    /// <returns>The RPC reply.</returns>
    /// <exception cref="SiteChannelUnavailableException">The site has no configured channel.</exception>
    public async Task<T> ExecuteAsync<T>(
        string siteId, Func<GrpcChannel, CancellationToken, Task<T>> call, CancellationToken ct)
    {
        ArgumentException.ThrowIfNullOrWhiteSpace(siteId);
        ArgumentNullException.ThrowIfNull(call);
        EnsureFailbackLoop();

        if (!_sites.TryGetValue(siteId, out var pair))
        {
            throw new SiteChannelUnavailableException(siteId);
        }

        GrpcChannel primary;
        GrpcChannel? secondary;
        bool primaryIsA;
        lock (pair.Gate)
        {
            primaryIsA = pair.CurrentIsA;
            primary = (primaryIsA ? pair.ChannelA : pair.ChannelB)
                ?? pair.ChannelA ?? pair.ChannelB
                ?? throw new SiteChannelUnavailableException(siteId);
            // Recompute which node `primary` actually is (the current side may be null).
            primaryIsA = ReferenceEquals(primary, pair.ChannelA);
            secondary = primaryIsA ? pair.ChannelB : pair.ChannelA;
        }

        try
        {
            return await call(primary, ct).ConfigureAwait(false);
        }
        catch (RpcException ex) when (ex.StatusCode == StatusCode.Unavailable && secondary is not null)
        {
            // Provably never reached a server → safe to try the other node, even for a write.
            _logger.LogWarning(
                "Site {SiteId} node {FromNode} unavailable; failing over to node {ToNode}",
                siteId, primaryIsA ? "A" : "B", primaryIsA ? "B" : "A");
            FlipTo(pair, toIsA: !primaryIsA);
            return await call(secondary, ct).ConfigureAwait(false);
        }
    }

    /// <summary>
    /// Probes a site's preferred node (NodeA) and, if reachable, returns stickiness to it. The
    /// background loop calls this per due site; tests call it to drive failback deterministically.
    /// </summary>
    /// <param name="siteId">Site identifier.</param>
    /// <param name="ct">Cancellation token.</param>
    /// <returns>True when the site is now (or already) pointed at its preferred node.</returns>
    internal async Task<bool> TryFailbackAsync(string siteId, CancellationToken ct)
    {
        if (!_sites.TryGetValue(siteId, out var pair))
        {
            return false;
        }

        GrpcChannel? preferred;
        lock (pair.Gate)
        {
            if (pair.CurrentIsA || pair.ChannelA is null)
            {
                return true; // already on preferred (or no preferred channel to fail back to)
            }
            preferred = pair.ChannelA;
        }

        var reachable = await _reachabilityProbe(preferred, ct).ConfigureAwait(false);
        lock (pair.Gate)
        {
            if (reachable)
            {
                pair.CurrentIsA = true;
                pair.ResetFailback();
                _logger.LogInformation("Site {SiteId} preferred node A reachable again; failing back", siteId);
                return true;
            }

            pair.BumpFailbackBackoff();
            return false;
        }
    }

    /// <summary>Test/diagnostic accessor: true when the site currently targets its preferred node (A).</summary>
    /// <param name="siteId">Site identifier.</param>
    /// <returns>True when on NodeA (or when the site is unknown).</returns>
    internal bool IsOnPreferredNode(string siteId)
        => !_sites.TryGetValue(siteId, out var pair) || pair.CurrentIsA;

    /// <summary>Starts the background failback loop if not already running. Idempotent.</summary>
    public void EnsureFailbackLoop()
    {
        if (_failbackLoop is not null || _disposed)
        {
            return;
        }
        lock (_loopGate)
        {
            if (_failbackLoop is null && !_disposed)
            {
                _failbackLoop = Task.Run(() => FailbackLoopAsync(_shutdown.Token));
            }
        }
    }

    private async Task FailbackLoopAsync(CancellationToken ct)
    {
        using var timer = new PeriodicTimer(FailbackSweepInterval);
        try
        {
            while (await timer.WaitForNextTickAsync(ct).ConfigureAwait(false))
            {
                var now = DateTime.UtcNow;
                foreach (var kvp in _sites)
                {
                    var pair = kvp.Value;
                    bool due;
                    lock (pair.Gate)
                    {
                        due = !pair.CurrentIsA && pair.ChannelA is not null && now >= pair.NextProbeUtc;
                    }
                    if (due)
                    {
                        try
                        {
                            await TryFailbackAsync(kvp.Key, ct).ConfigureAwait(false);
                        }
                        catch (OperationCanceledException)
                        {
                            throw;
                        }
                        catch (Exception ex)
                        {
                            _logger.LogDebug(ex, "Failback probe for site {SiteId} faulted", kvp.Key);
                        }
                    }
                }
            }
        }
        catch (OperationCanceledException)
        {
            // Shutting down.
        }
    }

    private void FlipTo(SitePair pair, bool toIsA)
    {
        lock (pair.Gate)
        {
            pair.CurrentIsA = toIsA;
            if (!toIsA)
            {
                // Failed over off the preferred node → schedule the first failback probe.
                pair.ScheduleFailback(InitialBackoff);
            }
            else
            {
                pair.ResetFailback();
            }
        }
    }

    private GrpcChannel BuildChannel(string endpoint, string siteId)
    {
        var channelOptions = new GrpcChannelOptions
        {
            HttpHandler = _handlerFactory is not null
                ? _handlerFactory(endpoint)
                : new SocketsHttpHandler
                {
                    KeepAlivePingDelay = _options.GrpcKeepAlivePingDelay,
                    KeepAlivePingTimeout = _options.GrpcKeepAlivePingTimeout,
                    KeepAlivePingPolicy = HttpKeepAlivePingPolicy.Always,
                    EnableMultipleHttp2Connections = true
                }
        }.WithSiteCredentials(_pskProvider, siteId);

        return GrpcChannel.ForAddress(endpoint, channelOptions);
    }

    private static async Task<bool> DefaultReachabilityProbe(GrpcChannel channel, CancellationToken ct)
    {
        try
        {
            using var timeout = CancellationTokenSource.CreateLinkedTokenSource(ct);
            timeout.CancelAfter(TimeSpan.FromSeconds(5));
            await channel.ConnectAsync(timeout.Token).ConfigureAwait(false);
            return true;
        }
        catch
        {
            return false;
        }
    }

    /// <inheritdoc />
    public void Dispose()
    {
        if (_disposed)
        {
            return;
        }
        _disposed = true;

        _shutdown.Cancel();
        try
        {
            _failbackLoop?.Wait(TimeSpan.FromSeconds(2));
        }
        catch
        {
            // Best-effort loop drain on shutdown.
        }
        _shutdown.Dispose();

        foreach (var pair in _sites.Values)
        {
            lock (pair.Gate)
            {
                pair.ChannelA?.Dispose();
                pair.ChannelB?.Dispose();
            }
        }
        _sites.Clear();
    }

    /// <summary>Per-site mutable channel state, guarded by <see cref="Gate"/>.</summary>
    private sealed class SitePair(string siteId)
    {
        public string SiteId { get; } = siteId;
        public readonly object Gate = new();
        public string? NodeAEndpoint;
        public string? NodeBEndpoint;
        public GrpcChannel? ChannelA;
        public GrpcChannel? ChannelB;

        /// <summary>Sticky current node; NodeA is preferred. True = pointing at NodeA.</summary>
        public bool CurrentIsA = true;

        /// <summary>Next time a failback probe is due (while failed over). MaxValue = not scheduled.</summary>
        public DateTime NextProbeUtc = DateTime.MaxValue;

        private TimeSpan _backoff = InitialBackoff;

        public void ScheduleFailback(TimeSpan delay)
        {
            _backoff = delay;
            NextProbeUtc = DateTime.UtcNow + _backoff;
        }

        public void BumpFailbackBackoff()
        {
            var next = TimeSpan.FromTicks(Math.Min(_backoff.Ticks * 2, MaxBackoff.Ticks));
            _backoff = next;
            NextProbeUtc = DateTime.UtcNow + _backoff;
        }

        public void ResetFailback()
        {
            _backoff = InitialBackoff;
            NextProbeUtc = DateTime.MaxValue;
        }
    }
}