ScadaBridge/src/ZB.MOM.WW.ScadaBridge.SiteRuntime/Streaming/SiteStreamManager.cs

using Akka;
using Akka.Actor;
using Akka.Streams;
using Akka.Streams.Dsl;
using Microsoft.Extensions.Logging;
using ZB.MOM.WW.ScadaBridge.Communication.Grpc;
using ZB.MOM.WW.ScadaBridge.Commons.Messages.Streaming;

namespace ZB.MOM.WW.ScadaBridge.SiteRuntime.Streaming;

/// <summary>
/// Site-Wide Akka Stream — manages a broadcast stream for attribute value
/// and alarm state changes. Instance Actors publish events via fire-and-forget Tell.
/// A BroadcastHub fans events out to per-subscriber graphs, each filtered by
/// instance name and bounded by a drop-oldest buffer.
///
/// Filterable by instance name for debug view.
/// Implements ISiteStreamSubscriber so the gRPC server can subscribe actors
/// to instance events without referencing SiteRuntime directly.
/// </summary>
public class SiteStreamManager : ISiteStreamSubscriber
{
    private ActorSystem? _system;
    private IMaterializer? _materializer;
    private readonly int _bufferSize;
    private readonly ILogger<SiteStreamManager> _logger;
    private readonly object _lock = new();

    private IActorRef? _sourceActor;
    private Source<ISiteStreamEvent, NotUsed>? _hubSource;
    private readonly Dictionary<string, SubscriptionInfo> _subscriptions = new();

    /// <summary>Initializes the stream manager with configuration and logger; the Akka stream is not started until <see cref="Initialize"/> is called.</summary>
    /// <param name="options">Site runtime options providing the stream buffer size.</param>
    /// <param name="logger">Logger instance.</param>
    public SiteStreamManager(
        SiteRuntimeOptions options,
        ILogger<SiteStreamManager> logger)
    {
        _bufferSize = options.StreamBufferSize;
        _logger = logger;
    }

    /// <summary>
    /// Initializes the broadcast stream. Must be called after ActorSystem is ready.
    /// The ActorSystem is passed here rather than via the constructor so that
    /// SiteStreamManager can be created by DI before the actor system exists.
    /// </summary>
    /// <param name="system">The running Akka <see cref="ActorSystem"/> used to materialize the broadcast stream.</param>
    public void Initialize(ActorSystem system)
    {
        _system = system;
        _materializer = _system.Materializer();

        var (sourceActor, hubSource) = Source.ActorRef<ISiteStreamEvent>(
                _bufferSize,
                OverflowStrategy.DropHead)
            .ToMaterialized(
                BroadcastHub.Sink<ISiteStreamEvent>(bufferSize: 256),
                Keep.Both)
            .Run(_materializer);

        _sourceActor = sourceActor;
        _hubSource = hubSource;

        _logger.LogInformation(
            "SiteStreamManager initialized with publish buffer size {BufferSize}", _bufferSize);
    }

    /// <summary>
    /// Publishes an attribute value change to the broadcast hub.
    /// Fire-and-forget — never blocks the calling actor.
    /// </summary>
    /// <param name="changed">The attribute value change event to publish.</param>
    public void PublishAttributeValueChanged(AttributeValueChanged changed)
    {
        _sourceActor?.Tell(changed);
    }

    /// <summary>
    /// Publishes an alarm state change to the broadcast hub.
    /// Fire-and-forget — never blocks the calling actor.
    /// </summary>
    /// <param name="changed">The alarm state change event to publish.</param>
    public void PublishAlarmStateChanged(AlarmStateChanged changed)
    {
        _sourceActor?.Tell(changed);
    }

    /// <inheritdoc />
    public string Subscribe(string instanceName, IActorRef subscriber)
    {
        if (_hubSource is null || _materializer is null)
            throw new InvalidOperationException("SiteStreamManager.Initialize must be called before Subscribe");

        var subscriptionId = Guid.NewGuid().ToString();
        var capturedInstance = instanceName;
        var capturedSubscriber = subscriber;

        var killSwitch = _hubSource
            .Where(ev => ev.InstanceUniqueName == capturedInstance)
            .Buffer(_bufferSize, OverflowStrategy.DropHead)
            .ViaMaterialized(KillSwitches.Single<ISiteStreamEvent>(), Keep.Right)
            .To(Sink.ForEach<ISiteStreamEvent>(ev => capturedSubscriber.Tell(ev)))
            .Run(_materializer);

        lock (_lock)
        {
            _subscriptions[subscriptionId] = new SubscriptionInfo(
                instanceName, subscriber, killSwitch, DateTimeOffset.UtcNow);
        }

        _logger.LogDebug(
            "Subscriber {SubscriptionId} registered for instance {Instance}",
            subscriptionId, instanceName);

        return subscriptionId;
    }

    /// <summary>
    /// Unsubscribe from instance events. Shuts down the per-subscriber
    /// stream graph via its KillSwitch.
    /// </summary>
    /// <param name="subscriptionId">The subscription ID returned by <see cref="Subscribe"/>.</param>
    /// <returns><c>true</c> if the subscription was found and removed; <c>false</c> if it was already gone.</returns>
    public bool Unsubscribe(string subscriptionId)
    {
        SubscriptionInfo? info;
        lock (_lock)
        {
            if (!_subscriptions.Remove(subscriptionId, out info))
                return false;
        }

        info.KillSwitch.Shutdown();
        _logger.LogDebug("Subscriber {SubscriptionId} removed", subscriptionId);
        return true;
    }

    /// <inheritdoc />
    public void RemoveSubscriber(IActorRef subscriber)
    {
        List<SubscriptionInfo> toShutdown;
        lock (_lock)
        {
            var matched = _subscriptions
                .Where(kvp => kvp.Value.Subscriber.Equals(subscriber))
                .ToList();
            foreach (var kvp in matched)
                _subscriptions.Remove(kvp.Key);
            toShutdown = matched.Select(kvp => kvp.Value).ToList();
        }

        foreach (var info in toShutdown)
            info.KillSwitch.Shutdown();

        if (toShutdown.Count > 0)
        {
            _logger.LogDebug(
                "Removed {Count} subscriptions for disconnected subscriber", toShutdown.Count);
        }
    }

    /// <summary>
    /// Returns the count of active subscriptions (for diagnostics/testing).
    /// </summary>
    public int SubscriptionCount
    {
        get { lock (_lock) { return _subscriptions.Count; } }
    }

    private record SubscriptionInfo(
        string InstanceName,
        IActorRef Subscriber,
        IKillSwitch KillSwitch,
        DateTimeOffset SubscribedAt);
}