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

using System.Collections.Concurrent;
using Grpc.Core;
using Grpc.Net.Client;
using Microsoft.Extensions.Logging;
using ZB.MOM.WW.ScadaBridge.Commons.Messages.Streaming;
using ZB.MOM.WW.ScadaBridge.Commons.Types.Alarms;
using ZB.MOM.WW.ScadaBridge.Commons.Types.Enums;
using Google.Protobuf.WellKnownTypes;

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

/// <summary>
/// Per-site gRPC client that manages streaming subscriptions to a site's
/// SiteStreamGrpcServer. The central-side DebugStreamBridgeActor uses this
/// to open server-streaming calls for individual instances.
/// </summary>
public class SiteStreamGrpcClient : IAsyncDisposable, IDisposable
{
    private readonly GrpcChannel? _channel;
    private readonly SiteStreamService.SiteStreamServiceClient? _client;
    private readonly ILogger? _logger;
    private readonly ConcurrentDictionary<string, CancellationTokenSource> _subscriptions = new();

    /// <summary>
    /// The gRPC endpoint (site node address) this client is bound to. The
    /// <see cref="SiteStreamGrpcClientFactory"/> compares this against the requested
    /// endpoint so a NodeA→NodeB failover flip (or a site address edit) is honoured
    /// rather than served stale from cache.
    /// </summary>
    public virtual string Endpoint { get; } = string.Empty;

    /// <summary>
    /// The HTTP/2 keepalive ping delay actually applied to this client's channel.
    /// Exposed for tests verifying that <see cref="CommunicationOptions"/> is honoured.
    /// </summary>
    internal TimeSpan KeepAlivePingDelay { get; }

    /// <summary>
    /// The HTTP/2 keepalive ping timeout actually applied to this client's channel.
    /// Exposed for tests verifying that <see cref="CommunicationOptions"/> is honoured.
    /// </summary>
    internal TimeSpan KeepAlivePingTimeout { get; }

    /// <summary>
    /// Creates a client with default communication options.
    /// </summary>
    /// <param name="endpoint">The gRPC endpoint address for the site.</param>
    /// <param name="logger">Logger for diagnostics and errors.</param>
    public SiteStreamGrpcClient(string endpoint, ILogger logger)
        : this(endpoint, logger, new CommunicationOptions())
    {
    }

    /// <summary>
    /// Creates a client whose HTTP/2 keepalive is taken from <see cref="CommunicationOptions"/>
    /// rather than hard-coded, satisfying the design doc's "gRPC Connection Keepalive"
    /// section which states these values are configurable.
    /// </summary>
    /// <param name="endpoint">The gRPC endpoint address for the site.</param>
    /// <param name="logger">Logger for diagnostics and errors.</param>
    /// <param name="options">Communication options including keepalive settings.</param>
    public SiteStreamGrpcClient(string endpoint, ILogger logger, CommunicationOptions options)
        : this(endpoint, logger, options, pskProvider: null, siteIdentifier: null)
    {
    }

    /// <summary>
    /// Creates a client that authenticates every call with the site's preshared key.
    /// This is the production shape: <c>SiteStreamService</c> is gated by
    /// <c>ControlPlaneAuthInterceptor</c> on the site node, so a client without credentials
    /// gets <see cref="StatusCode.PermissionDenied"/> on every call.
    /// </summary>
    /// <param name="endpoint">The gRPC endpoint address for the site.</param>
    /// <param name="logger">Logger for diagnostics and errors.</param>
    /// <param name="options">Communication options including keepalive settings.</param>
    /// <param name="pskProvider">Resolves the site's preshared key; null leaves the channel unauthenticated.</param>
    /// <param name="siteIdentifier">Site this channel talks to; null leaves the channel unauthenticated.</param>
    public SiteStreamGrpcClient(
        string endpoint,
        ILogger logger,
        CommunicationOptions options,
        ISitePskProvider? pskProvider,
        string? siteIdentifier)
    {
        Endpoint = endpoint;
        KeepAlivePingDelay = options.GrpcKeepAlivePingDelay;
        KeepAlivePingTimeout = options.GrpcKeepAlivePingTimeout;
        _channel = GrpcChannel.ForAddress(endpoint, new GrpcChannelOptions
        {
            HttpHandler = new SocketsHttpHandler
            {
                KeepAlivePingDelay = options.GrpcKeepAlivePingDelay,
                KeepAlivePingTimeout = options.GrpcKeepAlivePingTimeout,
                KeepAlivePingPolicy = HttpKeepAlivePingPolicy.Always
            }
        }.WithSiteCredentials(pskProvider, siteIdentifier));
        _client = new SiteStreamService.SiteStreamServiceClient(_channel);
        _logger = logger;
    }

    /// <summary>
    /// Protected constructor for unit testing without a real gRPC channel.
    /// Allows subclassing for mock implementations.
    /// </summary>
    protected SiteStreamGrpcClient()
    {
    }

    /// <summary>
    /// Protected constructor for unit testing — records the endpoint without
    /// opening a real gRPC channel, so endpoint-aware factory behaviour can be
    /// exercised by test doubles.
    /// </summary>
    /// <param name="endpoint">The gRPC endpoint address for the site.</param>
    protected SiteStreamGrpcClient(string endpoint)
    {
        Endpoint = endpoint;
    }

    /// <summary>
    /// Creates a test-only instance that has no gRPC channel. Used to test
    /// Unsubscribe and Dispose behavior without needing a real endpoint.
    /// </summary>
    /// <returns>A <see cref="SiteStreamGrpcClient"/> with no channel or client, for testing only.</returns>
    internal static SiteStreamGrpcClient CreateForTesting() => new();

    /// <summary>
    /// Registers a CancellationTokenSource for a correlation ID. Test-only.
    /// </summary>
    /// <param name="correlationId">Unique identifier for the subscription.</param>
    /// <param name="cts">CancellationTokenSource for managing the subscription lifecycle.</param>
    internal void AddSubscriptionForTesting(string correlationId, CancellationTokenSource cts)
    {
        _subscriptions[correlationId] = cts;
    }

    /// <summary>
    /// Registers a subscription's CancellationTokenSource for a correlation ID.
    /// If an entry already exists for that correlation ID (a reconnect race where two
    /// <see cref="SubscribeAsync"/> calls briefly share an ID), the prior CTS is
    /// cancelled and disposed so it cannot leak. Internal for testability.
    /// </summary>
    /// <param name="correlationId">Unique identifier for the subscription.</param>
    /// <param name="cts">CancellationTokenSource for managing the subscription lifecycle.</param>
    internal void RegisterSubscription(string correlationId, CancellationTokenSource cts)
    {
        if (_subscriptions.TryGetValue(correlationId, out var prior) && !ReferenceEquals(prior, cts))
        {
            prior.Cancel();
            prior.Dispose();
        }
        _subscriptions[correlationId] = cts;
    }

    /// <summary>
    /// Removes the subscription entry for a correlation ID only if the stored CTS is
    /// exactly the one supplied. A racing replacement stream may already own the slot,
    /// in which case this is a no-op. Internal for testability.
    /// </summary>
    /// <param name="correlationId">Unique identifier for the subscription.</param>
    /// <param name="cts">CancellationTokenSource to match before removing.</param>
    internal void RemoveSubscription(string correlationId, CancellationTokenSource cts)
    {
        _subscriptions.TryRemove(new KeyValuePair<string, CancellationTokenSource>(correlationId, cts));
    }

    /// <summary>
    /// Opens a server-streaming subscription for a specific instance.
    /// This is a long-running async method; the caller launches it as a background task.
    /// The <paramref name="onEvent"/> callback delivers domain events, and
    /// <paramref name="onError"/> lets the caller handle reconnection.
    /// </summary>
    /// <param name="correlationId">Unique identifier for this subscription.</param>
    /// <param name="instanceUniqueName">Unique name of the instance to subscribe to.</param>
    /// <param name="onEvent">Callback invoked for each domain event received from the stream.</param>
    /// <param name="onError">Callback invoked when the subscription encounters an error.</param>
    /// <param name="ct">Cancellation token to stop the subscription.</param>
    /// <returns>A task that represents the asynchronous operation.</returns>
    public virtual async Task SubscribeAsync(
        string correlationId,
        string instanceUniqueName,
        Action<object> onEvent,
        Action<Exception> onError,
        CancellationToken ct)
    {
        if (_client is null)
            throw new InvalidOperationException("Cannot subscribe on a test-only client.");

        var cts = CancellationTokenSource.CreateLinkedTokenSource(ct);
        RegisterSubscription(correlationId, cts);

        var request = new InstanceStreamRequest
        {
            CorrelationId = correlationId,
            InstanceUniqueName = instanceUniqueName
        };

        try
        {
            using var call = _client.SubscribeInstance(request, cancellationToken: cts.Token);

            await foreach (var evt in call.ResponseStream.ReadAllAsync(cts.Token))
            {
                var domainEvent = ConvertToDomainEvent(evt);
                if (domainEvent != null)
                    onEvent(domainEvent);
            }
        }
        catch (RpcException ex) when (ex.StatusCode == StatusCode.Cancelled)
        {
            // Normal cancellation — not an error
        }
        catch (Exception ex)
        {
            onError(ex);
        }
        finally
        {
            // Remove only our own entry -- a racing reconnect may already own the slot.
            RemoveSubscription(correlationId, cts);
        }
    }

    /// <summary>
    /// Opens the site-wide, <b>alarm-only</b> server-streaming subscription
    /// (<c>SubscribeSite</c>) for a whole site rather than a single instance. This is
    /// the central per-site feed that backs the aggregated Alarm Summary live cache
    /// (plan #10): the site runtime's <c>SubscribeSiteAlarms</c> hub drops the
    /// per-instance filter and carries only <see cref="AlarmStateChanged"/> events for
    /// <em>all</em> instances on the site (attributes are deliberately excluded — the
    /// summary never shows them and they are far higher-volume).
    /// <para>
    /// The callback is deliberately <b>typed</b> as <see cref="AlarmStateChanged"/>
    /// rather than the generic <c>Action&lt;object&gt;</c> used by
    /// <see cref="SubscribeAsync"/>: this stream is alarm-only by contract, so Task 4's
    /// per-site cache consumes an alarm delta directly with no downstream type test.
    /// The mapping is still shared via <see cref="ConvertToDomainEvent"/>; a non-alarm
    /// event (which should never appear on this stream) is defensively ignored rather
    /// than delivered or thrown.
    /// </para>
    /// This is a long-running async method; the caller launches it as a background task.
    /// Error handling, cancellation and cleanup mirror <see cref="SubscribeAsync"/>.
    /// </summary>
    /// <param name="correlationId">Unique identifier for this subscription.</param>
    /// <param name="onAlarmEvent">Callback invoked for each alarm delta received from the site-wide stream.</param>
    /// <param name="onError">Callback invoked when the subscription encounters an error.</param>
    /// <param name="ct">Cancellation token to stop the subscription.</param>
    /// <returns>A task that represents the asynchronous operation.</returns>
    public virtual async Task SubscribeSiteAsync(
        string correlationId,
        Action<AlarmStateChanged> onAlarmEvent,
        Action<Exception> onError,
        CancellationToken ct)
    {
        if (_client is null)
            throw new InvalidOperationException("Cannot subscribe on a test-only client.");

        var cts = CancellationTokenSource.CreateLinkedTokenSource(ct);
        RegisterSubscription(correlationId, cts);

        var request = new SiteStreamRequest
        {
            CorrelationId = correlationId
        };

        try
        {
            using var call = _client.SubscribeSite(request, cancellationToken: cts.Token);

            await foreach (var evt in call.ResponseStream.ReadAllAsync(cts.Token))
            {
                // Site-wide stream is alarm-only by contract; defensively ignore anything else.
                if (ConvertToAlarmEvent(evt) is { } alarm)
                    onAlarmEvent(alarm);
            }
        }
        catch (RpcException ex) when (ex.StatusCode == StatusCode.Cancelled)
        {
            // Normal cancellation — not an error
        }
        catch (Exception ex)
        {
            onError(ex);
        }
        finally
        {
            // Remove only our own entry -- a racing reconnect may already own the slot.
            RemoveSubscription(correlationId, cts);
        }
    }

    /// <summary>
    /// Cancels an active subscription by correlation ID.
    /// </summary>
    /// <param name="correlationId">Unique identifier of the subscription to cancel.</param>
    public virtual void Unsubscribe(string correlationId)
    {
        if (_subscriptions.TryRemove(correlationId, out var cts))
        {
            cts.Cancel();
            cts.Dispose();
        }
    }

    /// <summary>
    /// Converts a proto SiteStreamEvent to the corresponding domain message.
    /// Internal for testability.
    /// </summary>
    /// <param name="evt">The protobuf site stream event to convert.</param>
    /// <returns>The converted domain event, or null if the event type is not recognized.</returns>
    internal static object? ConvertToDomainEvent(SiteStreamEvent evt) => evt.EventCase switch
    {
        SiteStreamEvent.EventOneofCase.AttributeChanged => new AttributeValueChanged(
            evt.AttributeChanged.InstanceUniqueName,
            evt.AttributeChanged.AttributePath,
            evt.AttributeChanged.AttributeName,
            evt.AttributeChanged.Value,
            MapQuality(evt.AttributeChanged.Quality),
            evt.AttributeChanged.Timestamp.ToDateTimeOffset()),
        SiteStreamEvent.EventOneofCase.AlarmChanged => new AlarmStateChanged(
            evt.AlarmChanged.InstanceUniqueName,
            evt.AlarmChanged.AlarmName,
            MapAlarmState(evt.AlarmChanged.State),
            evt.AlarmChanged.Priority,
            evt.AlarmChanged.Timestamp.ToDateTimeOffset())
        {
            Level = MapAlarmLevel(evt.AlarmChanged.Level),
            Message = evt.AlarmChanged.Message ?? string.Empty,
            // Native alarm enrichment (additive — computed alarms carry their default condition).
            Kind = ParseAlarmKind(evt.AlarmChanged.Kind),
            Condition = new AlarmConditionState(
                Active: evt.AlarmChanged.Active,
                Acknowledged: evt.AlarmChanged.Acknowledged,
                Confirmed: evt.AlarmChanged.Confirmed,
                Shelve: AlarmShelveStateCodec.Parse(evt.AlarmChanged.ShelveState),
                Suppressed: evt.AlarmChanged.Suppressed,
                Severity: evt.AlarmChanged.Priority),
            SourceReference = evt.AlarmChanged.SourceReference ?? string.Empty,
            AlarmTypeName = evt.AlarmChanged.AlarmTypeName ?? string.Empty,
            Category = evt.AlarmChanged.Category ?? string.Empty,
            OperatorUser = evt.AlarmChanged.OperatorUser ?? string.Empty,
            OperatorComment = evt.AlarmChanged.OperatorComment ?? string.Empty,
            OriginalRaiseTime = evt.AlarmChanged.OriginalRaiseTime?.ToDateTimeOffset(),
            CurrentValue = evt.AlarmChanged.CurrentValue ?? string.Empty,
            LimitValue = evt.AlarmChanged.LimitValue ?? string.Empty,
            NativeSourceCanonicalName = evt.AlarmChanged.NativeSourceCanonicalName ?? string.Empty,
            IsConfiguredPlaceholder = evt.AlarmChanged.IsConfiguredPlaceholder
        },
        _ => null
    };

    /// <summary>
    /// Maps a proto <see cref="SiteStreamEvent"/> to a domain <see cref="AlarmStateChanged"/>,
    /// returning <c>null</c> for any non-alarm event. Used by <see cref="SubscribeSiteAsync"/>
    /// to enforce the alarm-only contract of the site-wide stream without duplicating the
    /// enrichment mapping in <see cref="ConvertToDomainEvent"/>. Internal for testability.
    /// </summary>
    /// <param name="evt">The protobuf site stream event to convert.</param>
    /// <returns>The mapped <see cref="AlarmStateChanged"/>, or <c>null</c> if the event is not an alarm.</returns>
    internal static AlarmStateChanged? ConvertToAlarmEvent(SiteStreamEvent evt) =>
        ConvertToDomainEvent(evt) as AlarmStateChanged;

    /// <summary>Parses the wire "kind" string back to <see cref="AlarmKind"/>; defaults to Computed.</summary>
    /// <param name="kind">The wire "kind" string from the gRPC payload; null or unrecognised defaults to <see cref="AlarmKind.Computed"/>.</param>
    /// <returns>The parsed <see cref="AlarmKind"/>, or <see cref="AlarmKind.Computed"/> when the value is null or unrecognised.</returns>
    internal static AlarmKind ParseAlarmKind(string? kind) =>
        System.Enum.TryParse<AlarmKind>(kind, ignoreCase: true, out var k) ? k : AlarmKind.Computed;

    /// <summary>
    /// Maps proto Quality enum to domain string. Internal for testability.
    /// </summary>
    /// <param name="quality">The protobuf quality value to map.</param>
    /// <returns>The mapped quality as a string ("Good", "Uncertain", "Bad", or "Unknown").</returns>
    internal static string MapQuality(Quality quality) => quality switch
    {
        Quality.Good => "Good",
        Quality.Uncertain => "Uncertain",
        Quality.Bad => "Bad",
        _ => "Unknown"
    };

    /// <summary>
    /// Maps proto AlarmStateEnum to domain AlarmState. Internal for testability.
    /// </summary>
    /// <param name="state">The protobuf alarm state to map.</param>
    /// <returns>The mapped domain alarm state.</returns>
    internal static AlarmState MapAlarmState(AlarmStateEnum state) => state switch
    {
        AlarmStateEnum.AlarmStateNormal => AlarmState.Normal,
        AlarmStateEnum.AlarmStateActive => AlarmState.Active,
        _ => AlarmState.Normal
    };

    /// <summary>
    /// Maps proto AlarmLevelEnum to domain AlarmLevel. Internal for testability.
    /// </summary>
    /// <param name="level">The protobuf alarm level to map.</param>
    /// <returns>The mapped domain alarm level.</returns>
    internal static AlarmLevel MapAlarmLevel(AlarmLevelEnum level) => level switch
    {
        AlarmLevelEnum.AlarmLevelLow => AlarmLevel.Low,
        AlarmLevelEnum.AlarmLevelLowLow => AlarmLevel.LowLow,
        AlarmLevelEnum.AlarmLevelHigh => AlarmLevel.High,
        AlarmLevelEnum.AlarmLevelHighHigh => AlarmLevel.HighHigh,
        _ => AlarmLevel.None
    };

    /// <summary>
    /// Releases all subscription CancellationTokenSources and the underlying
    /// gRPC channel. All teardown here is synchronous (CTS disposal and
    /// <see cref="GrpcChannel.Dispose"/>), so a synchronous <see cref="Dispose"/>
    /// can release everything without sync-over-async blocking.
    /// </summary>
    private void ReleaseResources()
    {
        foreach (var cts in _subscriptions.Values)
        {
            cts.Cancel();
            cts.Dispose();
        }
        _subscriptions.Clear();

        _channel?.Dispose();
    }

    /// <summary>
    /// Asynchronously disposes of the gRPC client and all subscriptions.
    /// </summary>
    /// <returns>A completed <see cref="ValueTask"/> after all subscriptions and the gRPC channel have been released.</returns>
    public virtual ValueTask DisposeAsync()
    {
        ReleaseResources();
        return ValueTask.CompletedTask;
    }

    /// <summary>
    /// Synchronous disposal. All resources held by this client are released
    /// synchronously, so callers (e.g. <see cref="SiteStreamGrpcClientFactory.Dispose"/>)
    /// need not block on the async disposal path.
    /// </summary>
    public virtual void Dispose()
    {
        ReleaseResources();
    }
}