using System.Text.Json;
using System.Threading.Channels;
using Microsoft.Extensions.Logging;
using Microsoft.Extensions.Logging.Abstractions;
using NATS.Server.Raft;

namespace NATS.Server.JetStream.Cluster;

/// <summary>
/// Background loop consuming meta RAFT entries and dispatching cluster state changes.
/// Reads <see cref="RaftLogEntry"/> items from a channel, parses the JSON command payload,
/// and applies the corresponding mutation to the <see cref="JetStreamMetaGroup"/>.
/// Go reference: jetstream_cluster.go:1455-1825 (monitorCluster).
/// </summary>
public sealed class JetStreamClusterMonitor
{
    private readonly JetStreamMetaGroup _meta;
    private readonly ChannelReader<RaftLogEntry> _entries;
    private readonly ILogger<JetStreamClusterMonitor> _logger;

    // Monotonic counter incremented after each entry (including malformed ones).
    // Protected by _processedLock so WaitForProcessedAsync can wait for a target
    // without races between the count read and the wait.
    private readonly object _processedLock = new();
    private int _processedCount;

    /// <summary>
    /// Total number of entries dequeued from the channel and processed (including
    /// malformed entries that were skipped). Useful for test synchronisation.
    /// </summary>
    public int ProcessedCount { get { lock (_processedLock) return _processedCount; } }

    public JetStreamClusterMonitor(JetStreamMetaGroup meta, ChannelReader<RaftLogEntry> entries)
        : this(meta, entries, NullLogger<JetStreamClusterMonitor>.Instance)
    {
    }

    public JetStreamClusterMonitor(
        JetStreamMetaGroup meta,
        ChannelReader<RaftLogEntry> entries,
        ILogger<JetStreamClusterMonitor> logger)
    {
        _meta = meta;
        _entries = entries;
        _logger = logger;
    }

    /// <summary>
    /// Starts consuming entries from the channel until the token is cancelled.
    /// Each entry is applied synchronously before the next is read.
    /// Returns normally (without throwing) when <paramref name="ct"/> is cancelled.
    /// </summary>
    public async Task StartAsync(CancellationToken ct)
    {
        try
        {
            await foreach (var entry in _entries.ReadAllAsync(ct))
            {
                ApplyMetaEntry(entry);
                lock (_processedLock)
                {
                    _processedCount++;
                    Monitor.PulseAll(_processedLock);
                }
            }
        }
        catch (OperationCanceledException) when (ct.IsCancellationRequested)
        {
            _logger.LogDebug("JetStreamClusterMonitor stopped via cancellation.");
        }
    }

    /// <summary>
    /// Waits until the cumulative <see cref="ProcessedCount"/> reaches at least
    /// <paramref name="targetCount"/>. Returns immediately when the target is already met.
    /// Used by tests to synchronise without sleeping.
    /// </summary>
    public Task WaitForProcessedAsync(int targetCount, CancellationToken ct)
    {
        // Fast path — already done.
        lock (_processedLock)
        {
            if (_processedCount >= targetCount)
                return Task.CompletedTask;
        }

        // Slow path — offload the blocking Monitor.Wait to a thread-pool thread so the
        // calling async context is not blocked. Monitor.Wait releases the lock atomically
        // while waiting, eliminating the TOCTOU race between reading the count and waiting.
        return Task.Run(() =>
        {
            lock (_processedLock)
            {
                while (_processedCount < targetCount)
                {
                    ct.ThrowIfCancellationRequested();
                    // Wait up to 50 ms so we can check cancellation periodically.
                    Monitor.Wait(_processedLock, millisecondsTimeout: 50);
                }
            }
        }, ct);
    }

    // ---------------------------------------------------------------
    // Entry dispatch
    // ---------------------------------------------------------------

    private void ApplyMetaEntry(RaftLogEntry entry)
    {
        try
        {
            using var doc = JsonDocument.Parse(entry.Command);
            var root = doc.RootElement;

            if (!root.TryGetProperty("Op", out var opElement))
                return;

            switch (opElement.GetString())
            {
                case "assignStream":
                    ProcessStreamAssignment(root);
                    break;
                case "removeStream":
                    ProcessStreamRemoval(root);
                    break;
                case "assignConsumer":
                    ProcessConsumerAssignment(root);
                    break;
                case "removeConsumer":
                    ProcessConsumerRemoval(root);
                    break;
                case "snapshot":
                    ApplyMetaSnapshot(root);
                    break;
                // Unknown ops are silently ignored — forward compatibility.
            }
        }
        catch (JsonException ex)
        {
            _logger.LogWarning(
                ex,
                "Skipping malformed meta RAFT entry at index {Index}: {Message}",
                entry.Index,
                ex.Message);
        }
    }

    // ---------------------------------------------------------------
    // Per-op processors
    // ---------------------------------------------------------------

    private void ProcessStreamAssignment(JsonElement root)
    {
        var streamName = root.GetProperty("StreamName").GetString()!;
        var peers = root.GetProperty("Peers").EnumerateArray()
            .Select(p => p.GetString()!)
            .ToList();
        var config = root.TryGetProperty("Config", out var cfg) ? cfg.GetString() ?? "{}" : "{}";

        var sa = new StreamAssignment
        {
            StreamName = streamName,
            Group = new RaftGroup { Name = streamName, Peers = peers },
            ConfigJson = config,
        };
        _meta.AddStreamAssignment(sa);
        _logger.LogDebug("Applied stream assignment for {StreamName}", streamName);
    }

    private void ProcessStreamRemoval(JsonElement root)
    {
        var streamName = root.GetProperty("StreamName").GetString()!;
        _meta.RemoveStreamAssignment(streamName);
        _logger.LogDebug("Applied stream removal for {StreamName}", streamName);
    }

    private void ProcessConsumerAssignment(JsonElement root)
    {
        var streamName = root.GetProperty("StreamName").GetString()!;
        var consumerName = root.GetProperty("ConsumerName").GetString()!;
        var peers = root.GetProperty("Peers").EnumerateArray()
            .Select(p => p.GetString()!)
            .ToList();

        var ca = new ConsumerAssignment
        {
            ConsumerName = consumerName,
            StreamName = streamName,
            Group = new RaftGroup { Name = consumerName, Peers = peers },
        };
        _meta.AddConsumerAssignment(streamName, ca);
        _logger.LogDebug("Applied consumer assignment {ConsumerName} on {StreamName}", consumerName, streamName);
    }

    private void ProcessConsumerRemoval(JsonElement root)
    {
        var streamName = root.GetProperty("StreamName").GetString()!;
        var consumerName = root.GetProperty("ConsumerName").GetString()!;
        _meta.RemoveConsumerAssignment(streamName, consumerName);
        _logger.LogDebug("Applied consumer removal {ConsumerName} from {StreamName}", consumerName, streamName);
    }

    private void ApplyMetaSnapshot(JsonElement root)
    {
        var dataB64 = root.GetProperty("Data").GetString()!;
        var data = Convert.FromBase64String(dataB64);
        var assignments = MetaSnapshotCodec.Decode(data);
        _meta.ReplaceAllAssignments(assignments);
        _logger.LogInformation("Applied meta snapshot: {StreamCount} streams restored", assignments.Count);
    }
}
