lmxopcua/src/Drivers/ZB.MOM.WW.OtOpcUa.Driver.Galaxy/Health/PerPlatformProbeWatcher.cs

using System.Collections.Concurrent;
using Microsoft.Extensions.Logging;
using Microsoft.Extensions.Logging.Abstractions;
using ZB.MOM.WW.OtOpcUa.Core.Abstractions;
using ZB.MOM.WW.OtOpcUa.Driver.Galaxy.Runtime;

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

/// <summary>
///     Subscribes the <c>ScanState</c> attribute of every <c>$WinPlatform</c> /
///     <c>$AppEngine</c> object the discoverer surfaced and translates ScanState
///     value-changes into per-host <see cref="HostConnectivityStatus"/> updates.
///     Ports the state machine in
///     <c>Driver.Galaxy.Host/Backend/Stability/GalaxyRuntimeProbeManager.cs</c> onto the
///     gateway subscription path.
/// </summary>
/// <remarks>
///     Address grammar: each platform tag's probe address is
///     <c>{platformTagName}.ScanState</c>. The watcher subscribes that address through
///     <see cref="IGalaxySubscriber"/>; the EventPump (PR 4.4) routes inbound
///     OnDataChange events back via <see cref="OnProbeValueChanged"/>. State decoding:
///     <list type="bullet">
///       <item>Quality &lt; <c>192</c> (Good) → <see cref="HostState.Unknown"/>.</item>
///       <item>Value <c>1</c>, <c>true</c>, or "Running" → <see cref="HostState.Running"/>.</item>
///       <item>Value <c>0</c>, <c>false</c>, or "Stopped" → <see cref="HostState.Stopped"/>.</item>
///       <item>Anything else with Good quality → <see cref="HostState.Faulted"/>.</item>
///     </list>
///     <see cref="SyncPlatformsAsync"/> is idempotent — call it after every
///     Discover / Rediscover. Newly-added platforms are subscribed; removed ones are
///     unsubscribed and dropped from the aggregator.
/// </remarks>
public sealed class PerPlatformProbeWatcher : IDisposable
{
    public const string ProbeSuffix = ".ScanState";

    private readonly IGalaxySubscriber _subscriber;
    private readonly HostStatusAggregator _aggregator;
    private readonly ILogger _logger;
    private readonly int _bufferedUpdateIntervalMs;

    // The driver's fan-out registry. The EventPump resolves every inbound event through it, so
    // the probe handles must be registered there (under SubscriptionRegistry.ProbeSubscriptionId)
    // or every ScanState event is dropped as "stale after unsubscribe". Null when the watcher is
    // driven directly (unit tests of the decoding table).
    private readonly SubscriptionRegistry? _registry;

    // #865: invoked before every probe SubscribeBulk so the driver's event stream is attached when the
    // gateway pushes the initial ScanState value (a fresh StreamEvents does not replay earlier events,
    // and ScanState rarely changes — a lost initial value leaves the host missing indefinitely).
    private readonly Func<CancellationToken, Task>? _ensureEventStream;

    // Tracked platform → gw item handle. Item handle 0 means the gw rejected the subscribe;
    // we keep the entry so SyncPlatformsAsync doesn't try to subscribe it again on every call.
    // Each handle carries the registry session that issued it (#859 I-A): item handles are per gateway
    // session and the next session reuses the numbers, so an old-session handle must never be published,
    // unsubscribed or purged against the current session.
    private readonly ConcurrentDictionary<string, ProbeHandle> _itemHandlesByPlatform =
        new(StringComparer.OrdinalIgnoreCase);
    private readonly Lock _syncLock = new();
    // Written under _syncLock by StopAsync; read under it wherever a write-back must not outlive a stop.
    private volatile bool _disposed;

    /// <summary>Initializes a new instance of the PerPlatformProbeWatcher class.</summary>
    /// <param name="subscriber">The Galaxy subscriber for managing probe subscriptions.</param>
    /// <param name="aggregator">The host status aggregator for tracking platform connectivity.</param>
    /// <param name="logger">Optional logger for diagnostic messages.</param>
    /// <param name="bufferedUpdateIntervalMs">Buffered update interval in milliseconds; must be >= 0.</param>
    public PerPlatformProbeWatcher(
        IGalaxySubscriber subscriber,
        HostStatusAggregator aggregator,
        ILogger? logger = null,
        int bufferedUpdateIntervalMs = 0)
        : this(subscriber, aggregator, registry: null, logger, bufferedUpdateIntervalMs)
    {
    }

    /// <summary>
    ///     Driver-side ctor: the watcher publishes its probe handles into <paramref name="registry"/>
    ///     under <see cref="SubscriptionRegistry.ProbeSubscriptionId"/> so the <see cref="EventPump"/>
    ///     fans ScanState events out to <see cref="OnProbeValueChanged"/>.
    /// </summary>
    /// <param name="subscriber">The Galaxy subscriber for managing probe subscriptions.</param>
    /// <param name="aggregator">The host status aggregator for tracking platform connectivity.</param>
    /// <param name="registry">The driver's subscription registry; null to keep the handles private.</param>
    /// <param name="logger">Optional logger for diagnostic messages.</param>
    /// <param name="bufferedUpdateIntervalMs">Buffered update interval in milliseconds; must be >= 0.</param>
    /// <param name="ensureEventStream">
    ///     Optional: awaited before every probe SubscribeBulk to make sure the event stream that carries
    ///     the initial ScanState value is running and requested (the driver starts its
    ///     <see cref="EventPump"/> and waits for its StreamEvents call to be issued).
    /// </param>
    internal PerPlatformProbeWatcher(
        IGalaxySubscriber subscriber,
        HostStatusAggregator aggregator,
        SubscriptionRegistry? registry,
        ILogger? logger = null,
        int bufferedUpdateIntervalMs = 0,
        Func<CancellationToken, Task>? ensureEventStream = null)
    {
        _subscriber = subscriber ?? throw new ArgumentNullException(nameof(subscriber));
        _aggregator = aggregator ?? throw new ArgumentNullException(nameof(aggregator));
        _registry = registry;
        _ensureEventStream = ensureEventStream;
        _logger = logger ?? NullLogger.Instance;
        if (bufferedUpdateIntervalMs < 0)
        {
            throw new ArgumentOutOfRangeException(nameof(bufferedUpdateIntervalMs),
                "bufferedUpdateIntervalMs must be >= 0; 0 means use the gw's default cadence.");
        }
        _bufferedUpdateIntervalMs = bufferedUpdateIntervalMs;
    }

    /// <summary>Current platform → probe-handle bindings, as the registry sees them. Test seam.</summary>
    internal IReadOnlyList<TagBinding> CurrentBindings => BuildBindings();

    private IReadOnlyList<TagBinding> BuildBindings() =>
        [.. _itemHandlesByPlatform.Select(kvp => new TagBinding(kvp.Key + ProbeSuffix, kvp.Value.ItemHandle))];

    /// <summary>The bindings to publish for <paramref name="session"/>: another session's handles bind as 0 (no events).</summary>
    private IReadOnlyList<TagBinding> BuildBindingsFor(int session) =>
        [.. _itemHandlesByPlatform.Select(kvp => new TagBinding(
            kvp.Key + ProbeSuffix, kvp.Value.Session == session ? kvp.Value.ItemHandle : 0))];

    private int CurrentSession => _registry?.CurrentSession ?? 0;

    /// <summary>
    ///     Mirror the tracked handle set into the driver registry (no-op without one, or once stopped —
    ///     a stopped watcher must never write its entry back into a registry the teardown cleared).
    /// </summary>
    private void PublishBindings()
    {
        lock (_syncLock) PublishBindingsLocked(CurrentSession);
    }

    /// <returns>False when the registry's session is no longer <paramref name="session"/> (nothing published).</returns>
    private bool PublishBindingsLocked(int session)
    {
        if (_registry is null || _disposed) return true;
        if (_itemHandlesByPlatform.IsEmpty)
        {
            _registry.Remove(SubscriptionRegistry.ProbeSubscriptionId);
            return true;
        }
        // Fenced (#859 I-A): only handles issued by `session` are bound, and only while it is current.
        return _registry.TryUpsert(SubscriptionRegistry.ProbeSubscriptionId, BuildBindingsFor(session), session);
    }

    /// <summary>
    ///     Re-issue SubscribeBulk for every tracked platform and replace the handles — the
    ///     reconnect replay path, where the gateway session was recreated and every pre-reconnect
    ///     handle is dead. Platforms the gateway now rejects keep a 0 handle (no events expected).
    /// </summary>
    /// <param name="cancellationToken">A cancellation token to cancel the operation.</param>
    /// <returns>The number of probe references re-subscribed.</returns>
    public async Task<int> ResubscribeAsync(CancellationToken cancellationToken)
    {
        ObjectDisposedException.ThrowIf(_disposed, this);
        var platforms = _itemHandlesByPlatform.Keys.ToArray();
        if (platforms.Length == 0) return 0;

        await SubscribeAndRecordAsync(platforms, cancellationToken).ConfigureAwait(false);
        return platforms.Length;
    }

    /// <summary>
    ///     The reconnect replay: re-issue SubscribeBulk only for platforms whose handle predates the
    ///     registry's current session, then republish. A platform a seed already advised on the new
    ///     session is left alone (re-subscribing it would double-advise and leak a handle), and every
    ///     reconnect calls this regardless of the shared registry entry's session — that entry is
    ///     re-stamped by any probe publish, so it cannot say which platforms are stale (#859).
    /// </summary>
    /// <param name="cancellationToken">A cancellation token to cancel the operation.</param>
    /// <returns>The number of probe references re-subscribed.</returns>
    public async Task<int> ResubscribeStaleAsync(CancellationToken cancellationToken)
    {
        ObjectDisposedException.ThrowIf(_disposed, this);
        string[] stale;
        lock (_syncLock)
        {
            var session = CurrentSession;
            stale = [.. _itemHandlesByPlatform.Where(kvp => kvp.Value.Session != session).Select(kvp => kvp.Key)];
        }
        if (stale.Length == 0) return 0;

        await SubscribeAndRecordAsync(stale, cancellationToken).ConfigureAwait(false);
        return stale.Length;
    }

    private async Task SubscribeAndRecordAsync(IReadOnlyList<string> platforms, CancellationToken cancellationToken)
    {
        // #861 / #859 I-A: a reply that straddled a gateway reconnect carries dead handles; re-issue.
        for (var attempt = 1; ; attempt++)
        {
            if (await TrySubscribeAndRecordAsync(platforms, cancellationToken).ConfigureAwait(false)) return;
            if (attempt >= MaxSessionAttempts)
            {
                _logger.LogWarning(
                    "PerPlatformProbeWatcher probe subscribe straddled {Attempts} gateway reconnects — giving up until the next replay/seed", attempt);
                return;
            }
        }
    }

    private const int MaxSessionAttempts = 3;

    /// <returns>False when the reply straddled a session change (nothing recorded; retry).</returns>
    private async Task<bool> TrySubscribeAndRecordAsync(IReadOnlyList<string> platforms, CancellationToken cancellationToken)
    {
        var probeAddresses = platforms.Select(p => p + ProbeSuffix).ToArray();

        // #865: the stream must be live BEFORE the advise — the initial value is the probe's only
        // event until the platform's scan state changes, and an unattached stream never sees it.
        if (_ensureEventStream is not null) await _ensureEventStream(cancellationToken).ConfigureAwait(false);

        // #859: the initial ScanState value can overtake the SubscribeBulk reply; park it until
        // PublishBindingsLocked registers the handles below (window closes at method exit).
        using var earlyEventWindow = _registry?.BeginPendingSubscribe();
        var session = CurrentSession; // the session this SubscribeBulk is issued on
        // PR 6.3 — use the configured bufferedUpdateIntervalMs (defaults to 0 = gw cadence
        // when the driver hasn't overridden MxAccess.PublishingIntervalMs). Probe ScanState
        // changes are rare so a coarser interval is usually fine; deployments that need
        // tighter health visibility can dial it down through GalaxyDriverOptions.
        var results = await _subscriber.SubscribeBulkAsync(
            probeAddresses, _bufferedUpdateIntervalMs, cancellationToken).ConfigureAwait(false);

        var acquired = new (string Platform, int ItemHandle, string? Error)[platforms.Count];
        for (var i = 0; i < platforms.Count; i++)
        {
            var match = results.FirstOrDefault(r => string.Equals(
                r.TagAddress, probeAddresses[i], StringComparison.OrdinalIgnoreCase));
            acquired[i] = (platforms[i], match is { WasSuccessful: true } ? match.ItemHandle : 0,
                match?.ErrorMessage ?? "<no result returned>");
        }

        // Re-check the stop AFTER the await, atomically with the record + publish: a StopAsync that
        // ran while SubscribeBulk was in flight already cleared the map and the registry entry, so
        // recording now would resurrect the probe entry into the next runtime's registry.
        bool stopped;
        lock (_syncLock)
        {
            stopped = _disposed;
            if (!stopped)
            {
                if (_registry is not null && _registry.CurrentSession != session)
                {
                    // Straddled a reconnect: these handles died with the old session and their numbers may
                    // be reused — neither record, publish nor unsubscribe them.
                    return false;
                }
                foreach (var (platform, itemHandle, _) in acquired)
                    _itemHandlesByPlatform[platform] = new ProbeHandle(itemHandle, session);
                if (!PublishBindingsLocked(session)) return false;
            }
        }

        if (stopped)
        {
            var orphans = acquired.Where(a => a.ItemHandle > 0).Select(a => a.ItemHandle).ToArray();
            if (orphans.Length > 0)
            {
                try { await _subscriber.UnsubscribeBulkAsync(orphans, CancellationToken.None).ConfigureAwait(false); }
                catch (Exception ex)
                {
                    _logger.LogWarning(ex,
                        "PerPlatformProbeWatcher could not release {Count} probe handle(s) acquired after stop", orphans.Length);
                }
            }
            return true;
        }

        foreach (var (platform, itemHandle, error) in acquired)
        {
            if (itemHandle <= 0)
            {
                _logger.LogWarning("PerPlatformProbeWatcher subscribe failed for {Platform}: {Error}", platform, error);
            }
        }
        return true;
    }

    /// <summary>
    ///     Snapshot of platform tag names currently watched. Read under the sync lock so a platform
    ///     is never visible here before <see cref="SubscribeAndRecordAsync"/> has published its probe
    ///     binding into the driver registry (the record and the publish are one locked step): a
    ///     caller that sees the platform and then raises a ScanState event must find it routable,
    ///     not dropped by the pump as "stale after unsubscribe" (#856).
    /// </summary>
    public IReadOnlyCollection<string> WatchedPlatforms
    {
        get { lock (_syncLock) return [.. _itemHandlesByPlatform.Keys]; }
    }

    /// <summary>
    ///     Reconcile the watched platform set against <paramref name="platformTagNames"/>.
    ///     Subscribes new entries, unsubscribes dropped ones. Calling with the same set is
    ///     a no-op.
    /// </summary>
    /// <param name="platformTagNames">The platform tag names to synchronize.</param>
    /// <param name="cancellationToken">A cancellation token to cancel the operation.</param>
    /// <returns>A task that represents the asynchronous operation.</returns>
    public async Task SyncPlatformsAsync(
        IEnumerable<string> platformTagNames, CancellationToken cancellationToken)
    {
        ObjectDisposedException.ThrowIf(_disposed, this);
        ArgumentNullException.ThrowIfNull(platformTagNames);

        var desired = new HashSet<string>(platformTagNames, StringComparer.OrdinalIgnoreCase);

        // Compute deltas under the lock so concurrent SyncPlatformsAsync calls don't
        // race on the membership view.
        List<string> toAdd;
        List<(string Platform, int ItemHandle)> toRemove;
        int removeSession;
        bool droppedWithoutRelease;
        lock (_syncLock)
        {
            toAdd = [.. desired.Where(p => !_itemHandlesByPlatform.ContainsKey(p))];
            // Captured before the RPC (#859 I-B); only handles issued by the CURRENT session are released —
            // an older session's numbers may now belong to unrelated live items (#859 I-A).
            removeSession = CurrentSession;
            toRemove = [.. _itemHandlesByPlatform
                .Where(kvp => !desired.Contains(kvp.Key)
                              && kvp.Value.ItemHandle > 0 && kvp.Value.Session == removeSession)
                .Select(kvp => (kvp.Key, kvp.Value.ItemHandle))];
            var deadRemovals = _itemHandlesByPlatform
                .Where(kvp => !desired.Contains(kvp.Key) && !(kvp.Value.ItemHandle > 0 && kvp.Value.Session == removeSession))
                .Select(kvp => kvp.Key).ToArray();
            // Dropped platforms with no releasable handle (rejected, or an older session's dead number):
            // forget them without any unsubscribe.
            foreach (var platform in deadRemovals)
            {
                _itemHandlesByPlatform.TryRemove(platform, out _);
                _aggregator.Remove(platform);
            }
            droppedWithoutRelease = deadRemovals.Length > 0;

            // Drop removed entries from the membership map up-front so a concurrent
            // OnProbeValueChanged for them is silently ignored. The unsubscribe RPC
            // runs outside the lock.
            foreach (var (platform, _) in toRemove)
            {
                _itemHandlesByPlatform.TryRemove(platform, out _);
                _aggregator.Remove(platform);
            }
        }

        if (droppedWithoutRelease && toRemove.Count == 0) PublishBindings(); // keep the fan-out map in step

        if (toRemove.Count > 0)
        {
            // Drop the dead handles from the fan-out map before the (slow, fallible) unsubscribe.
            PublishBindings();
            try
            {
                await _subscriber.UnsubscribeBulkAsync(
                    [.. toRemove.Select(t => t.ItemHandle)], cancellationToken).ConfigureAwait(false);
            }
            catch (Exception ex)
            {
                _logger.LogWarning(ex,
                    "PerPlatformProbeWatcher unsubscribe failed for {Count} probe(s); aggregator entries already cleared.",
                    toRemove.Count);
            }
            finally
            {
                // #859 I-3 interim: no trailing event for a removed probe handle may be promoted onto
                // a reused handle number.
                _registry?.PurgeParked([.. toRemove.Select(t => t.ItemHandle)], removeSession);
            }
        }

        if (toAdd.Count == 0) return;

        await SubscribeAndRecordAsync(toAdd, cancellationToken).ConfigureAwait(false);
    }

    /// <summary>
    ///     Route an OnDataChange for a probe address into the aggregator. The EventPump
    ///     (PR 4.4) calls this; tests can drive it directly to exercise the state machine
    ///     without spinning a real gw. Foreign references (anything not ending in
    ///     <see cref="ProbeSuffix"/>, or a probe for a platform we're not tracking) are
    ///     silently dropped.
    /// </summary>
    /// <param name="fullReference">The full reference path of the probe attribute.</param>
    /// <param name="value">The probe value to decode.</param>
    /// <param name="qualityByte">The quality byte for the value.</param>
    public void OnProbeValueChanged(string fullReference, object? value, byte qualityByte)
    {
        if (_disposed) return;
        ArgumentNullException.ThrowIfNull(fullReference);

        if (!fullReference.EndsWith(ProbeSuffix, StringComparison.OrdinalIgnoreCase)) return;
        var platform = fullReference[..^ProbeSuffix.Length];
        if (!_itemHandlesByPlatform.ContainsKey(platform)) return;

        var state = DecodeState(value, qualityByte);
        _aggregator.Update(new HostConnectivityStatus(platform, state, DateTime.UtcNow));
    }

    /// <summary>
    ///     Decode a ScanState value + raw quality byte to a <see cref="HostState"/>.
    ///     Public for tests that want to pin the decoding table.
    /// </summary>
    /// <param name="value">The probe value to decode.</param>
    /// <param name="qualityByte">The quality byte for the value.</param>
    /// <returns>The decoded host state.</returns>
    public static HostState DecodeState(object? value, byte qualityByte)
    {
        if (qualityByte < 192) return HostState.Unknown;

        return value switch
        {
            bool b => b ? HostState.Running : HostState.Stopped,
            int i => i == 1 ? HostState.Running : i == 0 ? HostState.Stopped : HostState.Faulted,
            short s => s == 1 ? HostState.Running : s == 0 ? HostState.Stopped : HostState.Faulted,
            long l => l == 1 ? HostState.Running : l == 0 ? HostState.Stopped : HostState.Faulted,
            string str when string.Equals(str, "Running", StringComparison.OrdinalIgnoreCase) => HostState.Running,
            string str when string.Equals(str, "Stopped", StringComparison.OrdinalIgnoreCase) => HostState.Stopped,
            _ => HostState.Faulted,
        };
    }

    /// <summary>
    ///     Stops the watcher and unsubscribes every tracked platform's probe handle on the gateway
    ///     (best-effort — a transport error is logged and swallowed). The async counterpart of
    ///     <see cref="Dispose"/>; the driver's <c>ShutdownAsync</c> uses it so the unsubscribe
    ///     round-trip is awaited rather than blocked on. Idempotent.
    /// </summary>
    /// <param name="cancellationToken">Bounds the unsubscribe round-trip.</param>
    /// <returns>A task that completes once the handles have been released (or the attempt gave up).</returns>
    public async Task StopAsync(CancellationToken cancellationToken)
    {
        int[] liveHandles;
        lock (_syncLock)
        {
            if (_disposed) return;
            _disposed = true;
            var session = CurrentSession; // never release an older session's (reusable) numbers
            liveHandles = [.. _itemHandlesByPlatform.Values
                .Where(h => h.ItemHandle > 0 && h.Session == session).Select(h => h.ItemHandle)];
            _itemHandlesByPlatform.Clear();
            _registry?.Remove(SubscriptionRegistry.ProbeSubscriptionId);
        }
        if (liveHandles.Length > 0)
        {
            try { await _subscriber.UnsubscribeBulkAsync(liveHandles, cancellationToken).ConfigureAwait(false); }
            catch (Exception ex) { _logger.LogWarning(ex, "PerPlatformProbeWatcher stop unsubscribe failed"); }
        }
    }

    /// <summary>Disposes the probe watcher and unsubscribes all tracked platforms (blocking on the round-trip).</summary>
    public void Dispose() => StopAsync(CancellationToken.None).GetAwaiter().GetResult();

    /// <summary>A probe item handle and the registry session that issued it.</summary>
    /// <param name="ItemHandle">The gateway item handle (0 = rejected).</param>
    /// <param name="Session">The registry session the handle belongs to.</param>
    private readonly record struct ProbeHandle(int ItemHandle, int Session);
}