using System.Collections.Concurrent;
using System.Text.Json;
using System.Threading.Channels;
using NATS.Server.JetStream.Models;
using NATS.Server.JetStream.Storage;
using NATS.Server.Subscriptions;

namespace NATS.Server.JetStream.MirrorSource;

// Go reference: server/stream.go:3860-4007 (processInboundSourceMsg)
// Go reference: server/stream.go:3752-3833 (processAllSourceMsgs)
// Go reference: server/stream.go:3474-3720 (setupSourceConsumer, trySetupSourceConsumer)

/// <summary>
/// Coordinates consumption from a source stream into a target stream with support for:
/// - Subject filtering via FilterSubject (Go: StreamSource.FilterSubject)
/// - Subject transform prefix applied before storing (Go: SubjectTransforms)
/// - Account isolation via SourceAccount
/// - Deduplication via Nats-Msg-Id header with configurable window
/// - Lag tracking per source
/// - Background sync loop with exponential backoff retry
/// </summary>
public sealed class SourceCoordinator : IAsyncDisposable
{
    // Go: sourceHealthCheckInterval = 10 * time.Second
    private static readonly TimeSpan HealthCheckInterval = TimeSpan.FromSeconds(10);

    // Go: sourceHealthHB = 1 * time.Second
    private static readonly TimeSpan HeartbeatInterval = TimeSpan.FromSeconds(1);

    private static readonly TimeSpan InitialRetryDelay = TimeSpan.FromMilliseconds(250);
    private static readonly TimeSpan MaxRetryDelay = TimeSpan.FromSeconds(30);
    private const int DefaultBatchSize = 256;

    private readonly IStreamStore _targetStore;
    private readonly StreamSourceConfig _sourceConfig;
    private readonly Channel<StoredMessage> _inbound;
    private readonly Lock _gate = new();
    private CancellationTokenSource? _cts;
    private Task? _syncLoop;
    private int _consecutiveFailures;

    // Go: si.sseq — last stream sequence from origin
    private ulong _expectedOriginSeq;

    // Go: si.dseq — delivery sequence tracking
    private ulong _deliverySeq;

    // Deduplication state: tracks recently seen Nats-Msg-Id values with their timestamps.
    // Go: server/stream.go doesn't have per-source dedup, but the stream's duplicate window
    // (DuplicateWindowMs) applies to publishes. We implement source-level dedup here.
    private readonly ConcurrentDictionary<string, DateTime> _dedupWindow = new(StringComparer.Ordinal);
    private DateTime _lastDedupPrune = DateTime.UtcNow;

    /// <summary>Last sequence number successfully applied from the origin stream.</summary>
    public ulong LastOriginSequence { get; private set; }

    /// <summary>UTC timestamp of the last successful sync operation.</summary>
    public DateTime LastSyncUtc { get; private set; }

    /// <summary>Number of consecutive sync failures (resets on success).</summary>
    public int ConsecutiveFailures
    {
        get { lock (_gate) return _consecutiveFailures; }
    }

    /// <summary>Whether the background sync loop is actively running.</summary>
    public bool IsRunning
    {
        get { lock (_gate) return _syncLoop is not null && !_syncLoop.IsCompleted; }
    }

    /// <summary>
    /// Current lag: origin last sequence minus local last sequence.
    /// Returns 0 when fully caught up.
    /// </summary>
    public ulong Lag { get; private set; }

    /// <summary>Total messages dropped by the subject filter.</summary>
    public long FilteredOutCount { get; private set; }

    /// <summary>Total messages dropped by deduplication.</summary>
    public long DeduplicatedCount { get; private set; }

    /// <summary>The source configuration driving this coordinator.</summary>
    public StreamSourceConfig Config => _sourceConfig;

    /// <summary>
    /// Current delivery sequence counter — number of messages successfully written to the target store.
    /// Go reference: server/stream.go si.dseq field
    /// </summary>
    public ulong GetDeliverySequence => _deliverySeq;

    /// <summary>
    /// When true, the target stream uses CRDT counter semantics and source messages
    /// are aggregated by adding their counter values.
    /// Go reference: server/stream.go — AllowMsgCounter source aggregation.
    /// </summary>
    public bool AllowMsgCounter { get; set; }

    // Per-source counter tracking: maps subject → last known counter value from this source.
    // Used for delta computation during aggregation.
    private readonly Dictionary<string, long> _sourceCounterValues = new(StringComparer.Ordinal);

    public SourceCoordinator(IStreamStore targetStore, StreamSourceConfig sourceConfig)
    {
        _targetStore = targetStore;
        _sourceConfig = sourceConfig;
        _inbound = Channel.CreateUnbounded<StoredMessage>(new UnboundedChannelOptions
        {
            SingleReader = true,
            SingleWriter = false,
        });
    }

    /// <summary>
    /// Builds the consumer configuration that would be sent to the source stream to set up
    /// a push consumer for consumption. This models the consumer create request generated by
    /// Go's setupSourceConsumer / trySetupSourceConsumer.
    /// Go reference: server/stream.go:3474-3720 (setupSourceConsumer, trySetupSourceConsumer)
    /// </summary>
    public ConsumerConfig BuildConsumerCreateRequest()
    {
        // Go: server/stream.go:3597-3598 — if ssi.FilterSubject != _EMPTY_ { req.Config.FilterSubject = ssi.FilterSubject }
        var cfg = new ConsumerConfig
        {
            AckPolicy = AckPolicy.None,
            Push = true,
            FlowControl = true,
            HeartbeatMs = (int)HeartbeatInterval.TotalMilliseconds,
        };

        if (!string.IsNullOrWhiteSpace(_sourceConfig.FilterSubject))
            cfg.FilterSubject = _sourceConfig.FilterSubject;

        // Go: server/stream.go:3573-3582 — resume from LastOriginSequence+1, or DeliverAll when starting fresh
        if (LastOriginSequence == 0)
        {
            cfg.DeliverPolicy = DeliverPolicy.All;
        }
        else
        {
            cfg.DeliverPolicy = DeliverPolicy.ByStartSequence;
            cfg.OptStartSeq = LastOriginSequence + 1;
        }

        return cfg;
    }

    /// <summary>
    /// Returns the JetStream API subject used to create a consumer on the source stream.
    /// Go reference: server/stream.go:3531 — $JS.API.CONSUMER.CREATE.{sourceName}
    /// </summary>
    public string BuildConsumerCreateSubject() =>
        $"$JS.API.CONSUMER.CREATE.{_sourceConfig.Name}";

    /// <summary>
    /// Processes a single inbound message from the origin stream.
    /// This is the direct-call path used when the origin and target are in the same process.
    /// Go reference: server/stream.go:3860-4007 (processInboundSourceMsg)
    /// </summary>
    public async Task OnOriginAppendAsync(StoredMessage message, CancellationToken ct)
    {
        // Account isolation: skip messages from different accounts.
        // Go: This is checked at the subscription level, but we enforce it here for in-process sources.
        if (!string.IsNullOrWhiteSpace(_sourceConfig.SourceAccount)
            && !string.IsNullOrWhiteSpace(message.Account)
            && !string.Equals(_sourceConfig.SourceAccount, message.Account, StringComparison.Ordinal))
        {
            return;
        }

        // Subject filter: only forward messages matching the filter.
        // Go: server/stream.go:3597-3598 — if ssi.FilterSubject != _EMPTY_ { req.Config.FilterSubject = ssi.FilterSubject }
        if (!string.IsNullOrWhiteSpace(_sourceConfig.FilterSubject)
            && !SubjectMatch.MatchLiteral(message.Subject, _sourceConfig.FilterSubject))
        {
            FilteredOutCount++;
            return;
        }

        // Deduplication: check Nats-Msg-Id header against the dedup window.
        if (_sourceConfig.DuplicateWindowMs > 0 && message.MsgId is not null)
        {
            if (IsDuplicate(message.MsgId))
            {
                DeduplicatedCount++;
                return;
            }

            RecordMsgId(message.MsgId);
        }

        // Go: si.sseq <= current — ignore older/duplicate messages
        if (_expectedOriginSeq > 0 && message.Sequence <= _expectedOriginSeq)
            return;

        // Subject transform: apply prefix and/or SubjectTransforms before storing.
        // Go: server/stream.go:3943-3956 (subject transform for the source)
        var subject = message.Subject;
        if (!string.IsNullOrWhiteSpace(_sourceConfig.SubjectTransformPrefix))
            subject = $"{_sourceConfig.SubjectTransformPrefix}{subject}";
        subject = ApplySubjectTransforms(subject);

        // Go: counter aggregation — when target has AllowMsgCounter, aggregate source counter values.
        // Reference: server/stream.go — processInboundSourceMsg counter aggregation path.
        if (AllowMsgCounter)
        {
            await AggregateCounterAsync(subject, message.Payload, ct);
        }
        else
        {
            await _targetStore.AppendAsync(subject, message.Payload, ct);
        }

        _expectedOriginSeq = message.Sequence;
        _deliverySeq++;
        LastOriginSequence = message.Sequence;
        LastSyncUtc = DateTime.UtcNow;
        Lag = 0;
    }

    /// <summary>
    /// Enqueues a message for processing by the background sync loop.
    /// </summary>
    public bool TryEnqueue(StoredMessage message)
    {
        return _inbound.Writer.TryWrite(message);
    }

    /// <summary>
    /// Starts the background sync loop that drains the inbound channel.
    /// Go reference: server/stream.go:3752-3833 (processAllSourceMsgs)
    /// </summary>
    public void StartSyncLoop()
    {
        lock (_gate)
        {
            if (_syncLoop is not null && !_syncLoop.IsCompleted)
                return;

            _cts = new CancellationTokenSource();
            _syncLoop = RunSyncLoopAsync(_cts.Token);
        }
    }

    /// <summary>
    /// Starts a pull-based sync loop that actively fetches from the origin store.
    /// Go reference: server/stream.go:3474-3720 (setupSourceConsumer + trySetupSourceConsumer)
    /// </summary>
    public void StartPullSyncLoop(IStreamStore originStore, int batchSize = DefaultBatchSize)
    {
        lock (_gate)
        {
            if (_syncLoop is not null && !_syncLoop.IsCompleted)
                return;

            _cts = new CancellationTokenSource();
            _syncLoop = RunPullSyncLoopAsync(originStore, batchSize, _cts.Token);
        }
    }

    /// <summary>
    /// Stops the background sync loop.
    /// Go reference: server/stream.go:3438-3469 (cancelSourceConsumer)
    /// </summary>
    public async Task StopAsync()
    {
        CancellationTokenSource? cts;
        Task? loop;
        lock (_gate)
        {
            cts = _cts;
            loop = _syncLoop;
        }

        if (cts is not null)
        {
            await cts.CancelAsync();
            if (loop is not null)
            {
                try { await loop; }
                catch (OperationCanceledException) { }
            }
        }

        lock (_gate)
        {
            _cts?.Dispose();
            _cts = null;
            _syncLoop = null;
        }
    }

    /// <summary>
    /// Reports current health state for monitoring.
    /// Go reference: server/stream.go:2687-2695 (sourcesInfo)
    /// </summary>
    public SourceHealthReport GetHealthReport(ulong? originLastSeq = null)
    {
        var lag = originLastSeq.HasValue && originLastSeq.Value > LastOriginSequence
            ? originLastSeq.Value - LastOriginSequence
            : Lag;

        return new SourceHealthReport
        {
            SourceName = _sourceConfig.Name,
            FilterSubject = _sourceConfig.FilterSubject,
            LastOriginSequence = LastOriginSequence,
            LastSyncUtc = LastSyncUtc,
            Lag = lag,
            ConsecutiveFailures = ConsecutiveFailures,
            IsRunning = IsRunning,
            IsStalled = LastSyncUtc != default
                        && DateTime.UtcNow - LastSyncUtc > HealthCheckInterval,
            FilteredOutCount = FilteredOutCount,
            DeduplicatedCount = DeduplicatedCount,
        };
    }

    /// <summary>
    /// Returns a structured monitoring response for this source.
    /// Go reference: server/stream.go:2687-2736 (sourcesInfo / StreamSourceInfo)
    /// </summary>
    public SourceInfoResponse GetSourceInfo()
    {
        var report = GetHealthReport();
        return new SourceInfoResponse
        {
            Name = _sourceConfig.Name,
            Lag = report.Lag,
            Active = report.LastSyncUtc != default ? (long)(DateTime.UtcNow - report.LastSyncUtc).TotalMilliseconds : -1,
            FilterSubject = _sourceConfig.FilterSubject,
            Error = null, // SourceCoordinator doesn't have error state tracking yet
        };
    }

    public async ValueTask DisposeAsync()
    {
        await StopAsync();
        _inbound.Writer.TryComplete();
    }

    // -------------------------------------------------------------------------
    // Background sync loop: channel-based
    // Go reference: server/stream.go:3752-3833 (processAllSourceMsgs)
    // -------------------------------------------------------------------------
    private async Task RunSyncLoopAsync(CancellationToken ct)
    {
        using var healthTimer = new PeriodicTimer(HealthCheckInterval);
        var reader = _inbound.Reader;

        while (!ct.IsCancellationRequested)
        {
            try
            {
                while (reader.TryRead(out var msg))
                {
                    await ProcessInboundMessageAsync(msg, ct);
                }

                var readTask = reader.WaitToReadAsync(ct).AsTask();
                var healthTask = healthTimer.WaitForNextTickAsync(ct).AsTask();
                await Task.WhenAny(readTask, healthTask);

                if (ct.IsCancellationRequested)
                    break;

                while (reader.TryRead(out var msg2))
                {
                    await ProcessInboundMessageAsync(msg2, ct);
                }
            }
            catch (OperationCanceledException) when (ct.IsCancellationRequested)
            {
                break;
            }
            catch (Exception)
            {
                lock (_gate) _consecutiveFailures++;
                var delay = CalculateBackoff(_consecutiveFailures);
                try { await Task.Delay(delay, ct); }
                catch (OperationCanceledException) { break; }
            }
        }
    }

    // -------------------------------------------------------------------------
    // Background sync loop: pull-based
    // Go reference: server/stream.go:3474-3720 (setupSourceConsumer)
    // -------------------------------------------------------------------------
    private async Task RunPullSyncLoopAsync(IStreamStore originStore, int batchSize, CancellationToken ct)
    {
        while (!ct.IsCancellationRequested)
        {
            try
            {
                var messages = await originStore.ListAsync(ct);
                var applied = 0;

                foreach (var msg in messages)
                {
                    if (ct.IsCancellationRequested) break;

                    if (msg.Sequence <= LastOriginSequence)
                        continue;

                    await ProcessInboundMessageAsync(msg, ct);
                    applied++;

                    if (applied >= batchSize)
                        break;
                }

                // Update lag
                if (messages.Count > 0)
                {
                    var originLast = messages[^1].Sequence;
                    Lag = originLast > LastOriginSequence ? originLast - LastOriginSequence : 0;
                }

                lock (_gate) _consecutiveFailures = 0;

                if (applied == 0)
                {
                    try { await Task.Delay(HeartbeatInterval, ct); }
                    catch (OperationCanceledException) { break; }
                }
            }
            catch (OperationCanceledException) when (ct.IsCancellationRequested)
            {
                break;
            }
            catch (Exception)
            {
                lock (_gate) _consecutiveFailures++;
                var delay = CalculateBackoff(_consecutiveFailures);
                try { await Task.Delay(delay, ct); }
                catch (OperationCanceledException) { break; }
            }
        }
    }

    // Go reference: server/stream.go:3860-4007 (processInboundSourceMsg)
    private async Task ProcessInboundMessageAsync(StoredMessage message, CancellationToken ct)
    {
        // Account isolation
        if (!string.IsNullOrWhiteSpace(_sourceConfig.SourceAccount)
            && !string.IsNullOrWhiteSpace(message.Account)
            && !string.Equals(_sourceConfig.SourceAccount, message.Account, StringComparison.Ordinal))
        {
            return;
        }

        // Subject filter
        if (!string.IsNullOrWhiteSpace(_sourceConfig.FilterSubject)
            && !SubjectMatch.MatchLiteral(message.Subject, _sourceConfig.FilterSubject))
        {
            FilteredOutCount++;
            return;
        }

        // Deduplication
        if (_sourceConfig.DuplicateWindowMs > 0 && message.MsgId is not null)
        {
            if (IsDuplicate(message.MsgId))
            {
                DeduplicatedCount++;
                return;
            }

            RecordMsgId(message.MsgId);
        }

        // Skip already-seen sequences
        if (_expectedOriginSeq > 0 && message.Sequence <= _expectedOriginSeq)
            return;

        // Redelivery check (Go: dc > 1)
        if (message.Redelivered)
            return;

        // Subject transform
        var subject = message.Subject;
        if (!string.IsNullOrWhiteSpace(_sourceConfig.SubjectTransformPrefix))
            subject = $"{_sourceConfig.SubjectTransformPrefix}{subject}";
        subject = ApplySubjectTransforms(subject);

        // Go: counter aggregation — when target has AllowMsgCounter, aggregate source counter values.
        if (AllowMsgCounter)
        {
            await AggregateCounterAsync(subject, message.Payload, ct);
        }
        else
        {
            await _targetStore.AppendAsync(subject, message.Payload, ct);
        }

        _expectedOriginSeq = message.Sequence;
        _deliverySeq++;
        LastOriginSequence = message.Sequence;
        LastSyncUtc = DateTime.UtcNow;

        lock (_gate) _consecutiveFailures = 0;
    }

    // -------------------------------------------------------------------------
    // Subject transforms
    // Go reference: server/stream.go — SubjectTransforms on StreamSource.
    // -------------------------------------------------------------------------

    /// <summary>
    /// Applies the SubjectTransforms list from the source config to the subject.
    /// Each transform in the list is tried; the first match wins.
    /// </summary>
    private string ApplySubjectTransforms(string subject)
    {
        if (_sourceConfig.SubjectTransforms.Count == 0)
            return subject;

        foreach (var stc in _sourceConfig.SubjectTransforms)
        {
            var transform = SubjectTransform.Create(stc.Source, stc.Destination);
            if (transform == null)
                continue;

            var result = transform.Apply(subject);
            if (result != null)
                return result;
        }

        return subject;
    }

    // -------------------------------------------------------------------------
    // Counter aggregation
    // Go reference: server/stream.go — processInboundSourceMsg counter aggregation.
    // -------------------------------------------------------------------------

    /// <summary>
    /// Aggregates a counter value from a source message into the target store.
    /// Computes the delta from this source's last known value and adds it to the
    /// target's current total for the subject.
    /// </summary>
    private async Task AggregateCounterAsync(string subject, ReadOnlyMemory<byte> payload, CancellationToken ct)
    {
        // Parse the source's current counter value from payload.
        var sourceValue = CounterValue.FromPayload(payload.Span);
        var sourceTotal = sourceValue.AsLong();

        // Compute delta: how much this source's value changed since last seen.
        _sourceCounterValues.TryGetValue(subject, out var previousSourceValue);
        var delta = sourceTotal - previousSourceValue;
        _sourceCounterValues[subject] = sourceTotal;

        if (delta == 0)
            return;

        // Load existing target value and add the delta.
        var lastMsg = await _targetStore.LoadLastBySubjectAsync(subject, ct);
        var existing = lastMsg != null
            ? CounterValue.FromPayload(lastMsg.Payload.Span)
            : new CounterValue();

        var newTotal = existing.AsLong() + delta;
        var newPayload = CounterValue.FromLong(newTotal).ToPayload();

        await _targetStore.AppendAsync(subject, newPayload, ct);
    }

    // -------------------------------------------------------------------------
    // Deduplication helpers
    // -------------------------------------------------------------------------

    /// <summary>
    /// Returns true if the given message ID is already present in the dedup window.
    /// Go reference: server/stream.go duplicate window check
    /// </summary>
    public bool IsDuplicate(string msgId)
    {
        PruneDedupWindowIfNeeded();
        return _dedupWindow.ContainsKey(msgId);
    }

    /// <summary>
    /// Records a message ID in the dedup window with the current timestamp.
    /// Go reference: server/stream.go duplicate window tracking
    /// </summary>
    public void RecordMsgId(string msgId)
    {
        _dedupWindow[msgId] = DateTime.UtcNow;
    }

    /// <summary>
    /// Removes all dedup window entries whose timestamps are earlier than the given cutoff.
    /// This is the explicit-cutoff variant intended for testing, as opposed to the internal
    /// time-based pruning done by <see cref="PruneDedupWindowIfNeeded"/>.
    /// Go reference: server/stream.go duplicate window pruning
    /// </summary>
    public void PruneDedupWindow(DateTimeOffset cutoff)
    {
        var cutoffDt = cutoff.UtcDateTime;
        foreach (var kvp in _dedupWindow)
        {
            if (kvp.Value < cutoffDt)
                _dedupWindow.TryRemove(kvp.Key, out _);
        }
    }

    private void PruneDedupWindowIfNeeded()
    {
        if (_sourceConfig.DuplicateWindowMs <= 0)
            return;

        var now = DateTime.UtcNow;
        // Prune at most once per second to avoid overhead
        if ((now - _lastDedupPrune).TotalMilliseconds < 1000)
            return;

        _lastDedupPrune = now;
        var cutoff = now.AddMilliseconds(-_sourceConfig.DuplicateWindowMs);
        foreach (var kvp in _dedupWindow)
        {
            if (kvp.Value < cutoff)
                _dedupWindow.TryRemove(kvp.Key, out _);
        }
    }

    // Go reference: server/stream.go:3478-3505 (calculateRetryBackoff)
    private static TimeSpan CalculateBackoff(int failures)
    {
        var baseDelay = InitialRetryDelay.TotalMilliseconds * Math.Pow(2, Math.Min(failures - 1, 10));
        var capped = Math.Min(baseDelay, MaxRetryDelay.TotalMilliseconds);
        var jitter = Random.Shared.NextDouble() * 0.2 * capped;
        return TimeSpan.FromMilliseconds(capped + jitter);
    }
}

/// <summary>
/// Health report for a source coordinator, used by monitoring endpoints.
/// Go reference: server/stream.go:2687-2736 (sourcesInfo, sourceInfo)
/// </summary>
public sealed record SourceHealthReport
{
    public string SourceName { get; init; } = string.Empty;
    public string? FilterSubject { get; init; }
    public ulong LastOriginSequence { get; init; }
    public DateTime LastSyncUtc { get; init; }
    public ulong Lag { get; init; }
    public int ConsecutiveFailures { get; init; }
    public bool IsRunning { get; init; }
    public bool IsStalled { get; init; }
    public long FilteredOutCount { get; init; }
    public long DeduplicatedCount { get; init; }
}

/// <summary>
/// Structured monitoring response for a source, as returned by the JetStream monitoring API.
/// Go reference: server/stream.go:2687-2736 (sourcesInfo / StreamSourceInfo)
/// </summary>
public sealed record SourceInfoResponse
{
    /// <summary>Name of the origin stream.</summary>
    public string Name { get; init; } = string.Empty;

    /// <summary>Number of messages behind the origin stream.</summary>
    public ulong Lag { get; init; }

    /// <summary>
    /// Milliseconds since the last successful sync.
    /// -1 when the source has never synced.
    /// </summary>
    public long Active { get; init; } = -1;

    /// <summary>Optional subject filter applied by this source coordinator.</summary>
    public string? FilterSubject { get; init; }

    /// <summary>Current error message, or null if no error.</summary>
    public string? Error { get; init; }
}
