Finlytic/FinlyticCore/Util/ManagedMqttClient.cs

using System;
using System.Collections.Concurrent;
using System.Text;
using System.Text.Json;
using System.Text.Json.Serialization;
using System.Threading;
using System.Threading.Tasks;
using FinlyticCore.Models;
using Microsoft.Extensions.Logging;
using MQTTnet;

namespace FinlyticCore.Util;

/// <summary>
/// An abstract, resilient MQTT client wrapper designed for microservice architectures.
/// Handles automatic reconnection, structured JSON publishing, thread-safe subscription management, and synchronous Request-Reply (RPC).
/// </summary>
public abstract class ManagedMqttClient : IDisposable
{
    private readonly ILogger<ManagedMqttClient> _logger;
    private readonly IMqttClient _mqttClient;
    private CancellationTokenSource? _cts;
    
    // Tracks pending RPC requests waiting for a specific correlation ID reply
    private readonly ConcurrentDictionary<string, TaskCompletionSource<string>> _pendingRequests = new();

    /// <summary>
    /// Gets a value indicating whether the client is currently connected to the MQTT broker.
    /// </summary>
    public bool IsConnected => _mqttClient.IsConnected;

    protected ManagedMqttClient(ILogger<ManagedMqttClient> logger)
    {
        _logger = logger;
        _mqttClient = new MqttClientFactory().CreateMqttClient();
        
        _mqttClient.ApplicationMessageReceivedAsync += HandleIncomingMessageAsync;
        _mqttClient.DisconnectedAsync += HandleDisconnectAsync;
    }

    /// <summary>
    /// Establishes a connection to the MQTT broker and initializes the background auto-reconnection loop.
    /// </summary>
    /// <param name="config">The network and credential configuration options for the broker.</param>
    public async Task ConnectAsync(MqttConfiguration config)
    {
        if (IsConnected)
            throw new InvalidOperationException("MQTT client is already connected.");

        _cts = new CancellationTokenSource();

        var optionsBuilder = new MqttClientOptionsBuilder()
            .WithTcpServer(config.Host, config.Port)
            .WithClientId(config.ClientId)
            .WithCleanSession();

        if (!string.IsNullOrWhiteSpace(config.Username))
        {
            optionsBuilder.WithCredentials(config.Username, config.Password);
        }

        var options = optionsBuilder.Build();

        _logger.LogInformation("Connecting to MQTT broker at {Host}:{Port}...", config.Host, config.Port);
        
        try
        {
            await _mqttClient.ConnectAsync(options, _cts.Token);
            _logger.LogInformation("Successfully connected to MQTT broker.");
            
            await OnConnectedAsync();
        }
        catch (Exception ex)
        {
            _logger.LogError(ex, "Failed to establish initial connection to MQTT broker. Reconnection loop will handle recovery.");
        }
    }

    /// <summary>
    /// Gracefully disconnects from the broker and stops all ongoing background loops.
    /// </summary>
    public async Task DisconnectAsync()
    {
        if (_cts != null)
        {
            await _cts.CancelAsync();
        }

        if (_mqttClient.IsConnected)
        {
            try
            {
                await _mqttClient.DisconnectAsync(new MqttClientDisconnectOptions 
                { 
                    Reason = MqttClientDisconnectOptionsReason.NormalDisconnection 
                });
                _logger.LogInformation("MQTT connection gracefully closed.");
            }
            catch (Exception ex)
            {
                _logger.LogWarning(ex, "An error occurred while disconnecting from the MQTT broker.");
            }
        }
    }

    /// <summary>
    /// Subscribes to a specific MQTT topic filter.
    /// </summary>
    /// <param name="topic">The topic pattern or wildcard to subscribe to.</param>
    /// <param name="noLocal">If set to <c>true</c>, the broker will not forward messages published by this client back to itself.</param>
    protected async Task SubscribeAsync(string topic, bool noLocal = false)
    {
        if (!IsConnected)
        {
            _logger.LogWarning("Subscription to topic '{Topic}' delayed: Client is currently offline.", topic);
            return;
        }

        var filterBuilder = new MqttTopicFilterBuilder().WithTopic(topic);
        if (noLocal)
        {
            filterBuilder.WithNoLocal();
        }

        var subscribeOptions = new MqttClientFactory().CreateSubscribeOptionsBuilder()
            .WithTopicFilter(filterBuilder.Build())
            .Build();

        await _mqttClient.SubscribeAsync(subscribeOptions, CancellationToken.None);
        _logger.LogDebug("Successfully subscribed to topic: {Topic} (NoLocal: {NoLocal})", topic, noLocal);
    }

    /// <summary>
    /// Publishes a raw string message payload to the specified topic.
    /// </summary>
    public async Task PublishAsync(string topic, string payload, bool retain = false)
    {
        if (!IsConnected)
            throw new InvalidOperationException("Cannot publish message: MQTT client is offline.");

        var message = new MqttApplicationMessageBuilder()
            .WithTopic(topic)
            .WithPayload(payload)
            .WithQualityOfServiceLevel(MQTTnet.Protocol.MqttQualityOfServiceLevel.AtLeastOnce)
            .WithRetainFlag(retain)
            .Build();

        await _mqttClient.PublishAsync(message, CancellationToken.None);
    }

    /// <summary>
    /// Serializes a generic object into a structured JSON string and publishes it to the specified topic.
    /// Utilizes .NET 8 JSON Source Generators for zero-reflection overhead, with reflection fallback for unregistered types.
    /// </summary>
    public Task PublishAsync<T>(string topic, T data, bool retain = false)
    {
        byte[] jsonBytes;
        var typeInfo = FinlyticJsonSerializerContext.Default.GetTypeInfo(typeof(T))
                       ?? (data != null ? FinlyticJsonSerializerContext.Default.GetTypeInfo(data.GetType()) : null);

        if (typeInfo != null)
        {
            jsonBytes = JsonSerializer.SerializeToUtf8Bytes(data, typeInfo);
        }
        else
        {
            jsonBytes = JsonSerializer.SerializeToUtf8Bytes(data);
        }

        if (!IsConnected)
            throw new InvalidOperationException("Cannot publish message: MQTT client is offline.");

        var message = new MqttApplicationMessageBuilder()
            .WithTopic(topic)
            .WithPayload(jsonBytes)
            .WithQualityOfServiceLevel(MQTTnet.Protocol.MqttQualityOfServiceLevel.AtLeastOnce)
            .WithRetainFlag(retain)
            .Build();

        return _mqttClient.PublishAsync(message, CancellationToken.None);
    }

    /// <summary>
    /// Sends a generic request payload to an RPC channel and asynchronously blocks until a matching response arrives.
    /// Uses the topic conventions: <c>services/request/{channel}/{correlationId}</c> and <c>services/response/{channel}/{correlationId}</c>.
    /// </summary>
    /// <typeparam name="TResponse">The expected strongly-typed object type of the reply.</typeparam>
    /// <typeparam name="TRequest">The type of the payload being transmitted.</typeparam>
    /// <param name="channel">The target sub-channel or service name (e.g., "sentix", "assets").</param>
    /// <param name="requestData">The object that will be serialized to JSON and sent.</param>
    /// <param name="timeout">Optional. Maximum time to wait before returning null. Defaults to 10 seconds.</param>
    public async Task<TResponse?> SendRpcRequestAsync<TResponse, TRequest>(
        string channel, 
        TRequest requestData, 
        TimeSpan? timeout = null) 
        where TResponse : class 
        where TRequest : class
    {
        if (!IsConnected)
            throw new InvalidOperationException("Cannot execute RPC request: MQTT client is offline.");

        // 1. Generate a unique Correlation ID for this specific transaction
        string correlationId = Guid.NewGuid().ToString("N");
        
        var tcs = new TaskCompletionSource<string>(TaskCreationOptions.RunContinuationsAsynchronously);
        _pendingRequests.TryAdd(correlationId, tcs);

        string requestTopic = $"services/request/{channel}/{correlationId}";

        // 2. Serialize and dispatch via the existing JSON helper
        await PublishAsync(requestTopic, requestData);
        _logger.LogInformation("RPC request published to '{Topic}' [CorrelationId: {Id}]", requestTopic, correlationId);

        try
        {
            // 3. Block asynchronously until the response loop resolves the token
            var effectiveTimeout = timeout ?? TimeSpan.FromSeconds(25);
            var rawJsonResult = await tcs.Task.WaitAsync(effectiveTimeout);

            if (typeof(TResponse) == typeof(string))
            {
                return rawJsonResult as TResponse;
            }

            var respTypeInfo = FinlyticJsonSerializerContext.Default.GetTypeInfo(typeof(TResponse));
            if (respTypeInfo != null)
            {
                return JsonSerializer.Deserialize(rawJsonResult, respTypeInfo) as TResponse;
            }

            return JsonSerializer.Deserialize<TResponse>(rawJsonResult);
        }
        catch (TimeoutException)
        {
            _logger.LogWarning("RPC request timed out on channel '{Channel}' [CorrelationId: {Id}]", channel, correlationId);
            return null;
        }
        finally
        {
            // Always clean up the dictionary to prevent memory leaks
            _pendingRequests.TryRemove(correlationId, out _);
        }
    }

    private async Task HandleIncomingMessageAsync(MqttApplicationMessageReceivedEventArgs e)
    {
        try
        {
            var topic = e.ApplicationMessage.Topic;
            var payload = Encoding.UTF8.GetString(e.ApplicationMessage.Payload);
            _logger.LogInformation("MQTT message received on topic '{Topic}', length={Length}", topic, payload?.Length ?? 0);

            // Intercept message if it belongs to the RPC response convention
            if (topic.StartsWith("services/response/"))
            {
                var lastSlashIndex = topic.LastIndexOf('/');
                if (lastSlashIndex != -1)
                {
                    string correlationId = topic[(lastSlashIndex + 1)..];

                    if (_pendingRequests.TryRemove(correlationId, out var tcs))
                    {
                        tcs.SetResult(payload);
                        return; // Sinks the message, avoiding triggering OnMessageReceivedAsync for active RPC handles
                    }
                }
            }

            // Regular Pub/Sub message propagation
            await OnMessageReceivedAsync(topic, payload);
        }
        catch (Exception ex)
        {
            OnError(ex);
        }
    }

    private async Task HandleDisconnectAsync(MqttClientDisconnectedEventArgs e)
    {
        // Prevent trigger during deliberate connection shutdowns
        if (_cts == null || _cts.IsCancellationRequested) 
            return; 

        _logger.LogWarning("Lost connection to MQTT broker (Reason: {Reason}). Initiating auto-reconnect loop...", e.Reason);

        int attempt = 0;
        while (_cts != null && !_cts.IsCancellationRequested)
        {
            attempt++;
            var delaySeconds = Math.Min(5 * Math.Pow(2, attempt - 1), 60); // 5s, 10s, 20s, 40s, 60s max

            try
            {
                _logger.LogInformation("Reconnect attempt {Attempt} in {Delay}s...", attempt, delaySeconds);
                await Task.Delay(TimeSpan.FromSeconds(delaySeconds), _cts.Token);
                
                await _mqttClient.ReconnectAsync(_cts.Token);
                
                if (_mqttClient.IsConnected)
                {
                    _logger.LogInformation("MQTT client reconnected successfully after {Attempt} attempt(s).", attempt);
                    await OnConnectedAsync();
                    return;
                }
            }
            catch (OperationCanceledException) { return; /* Expected on application shutdown */ }
            catch (Exception ex)
            {
                _logger.LogWarning(ex, "Reconnect attempt {Attempt} to the MQTT broker failed.", attempt);
            }
        }
    }

    /// <summary>
    /// Fired automatically whenever a connection or reconnection is successfully established.
    /// Ideal place to trigger <see cref="SubscribeAsync"/> operations.
    /// </summary>
    protected abstract Task OnConnectedAsync();

    /// <summary>
    /// Fired whenever a new message lands on a registered subscription channel.
    /// </summary>
    /// <param name="topic">The specific topic where the message was broadcasted.</param>
    /// <param name="payload">The deserialized UTF-8 payload string.</param>
    protected abstract Task OnMessageReceivedAsync(string topic, string payload);

    /// <summary>
    /// Virtual fallback method to catch and handle processing level exceptions inside the incoming pipeline.
    /// </summary>
    protected virtual void OnError(Exception ex)
    {
        _logger.LogError(ex, "An unhandled exception occurred within the ManagedMqttClient messaging pipeline.");
    }

    public void Dispose()
    {
        DisconnectAsync().GetAwaiter().GetResult();
        _cts?.Dispose();
        _mqttClient.Dispose();
        GC.SuppressFinalize(this);
    }
}