using NATS.Server.JetStream.Storage;
using NATS.Server.JetStream.Models;
using NATS.Server.Subscriptions;

namespace NATS.Server.JetStream.Consumers;

/// <summary>
/// Pre-compiled filter for efficient subject matching against consumer filter subjects.
/// For 0 filters: always matches. For 1 filter: uses SubjectMatch.MatchLiteral directly.
/// For N filters: uses a HashSet for exact (literal) subjects and falls back to
/// SubjectMatch.MatchLiteral for wildcard filter patterns.
/// </summary>
public sealed class CompiledFilter
{
    private readonly HashSet<string>? _exactFilters;
    private readonly string[]? _wildcardFilters;
    private readonly string? _singleFilter;
    private readonly bool _matchAll;

    public CompiledFilter(IReadOnlyList<string> filterSubjects)
    {
        if (filterSubjects.Count == 0)
        {
            _matchAll = true;
            return;
        }

        if (filterSubjects.Count == 1)
        {
            _singleFilter = filterSubjects[0];
            return;
        }

        // Separate exact (literal) subjects from wildcard patterns
        var exact = new HashSet<string>(StringComparer.Ordinal);
        var wildcards = new List<string>();

        foreach (var filter in filterSubjects)
        {
            if (SubjectMatch.IsLiteral(filter))
                exact.Add(filter);
            else
                wildcards.Add(filter);
        }

        _exactFilters = exact.Count > 0 ? exact : null;
        _wildcardFilters = wildcards.Count > 0 ? wildcards.ToArray() : null;
    }

    /// <summary>
    /// Returns <c>true</c> if the given subject matches any of the compiled filter patterns.
    /// </summary>
    public bool Matches(string subject)
    {
        if (_matchAll)
            return true;

        if (_singleFilter is not null)
            return SubjectMatch.MatchLiteral(subject, _singleFilter);

        // Multi-filter path: check exact set first, then wildcard patterns
        if (_exactFilters is not null && _exactFilters.Contains(subject))
            return true;

        if (_wildcardFilters is not null)
        {
            foreach (var wc in _wildcardFilters)
            {
                if (SubjectMatch.MatchLiteral(subject, wc))
                    return true;
            }
        }

        return false;
    }

    /// <summary>
    /// Create a <see cref="CompiledFilter"/> from a <see cref="ConsumerConfig"/>.
    /// Uses <see cref="ConsumerConfig.FilterSubjects"/> first, falling back to
    /// <see cref="ConsumerConfig.FilterSubject"/> if the list is empty.
    /// </summary>
    public static CompiledFilter FromConfig(ConsumerConfig config)
    {
        if (config.FilterSubjects.Count > 0)
            return new CompiledFilter(config.FilterSubjects);

        if (!string.IsNullOrWhiteSpace(config.FilterSubject))
            return new CompiledFilter([config.FilterSubject]);

        return new CompiledFilter([]);
    }
}

public sealed class PullConsumerEngine
{
    public async ValueTask<PullFetchBatch> FetchAsync(StreamHandle stream, ConsumerHandle consumer, int batch, CancellationToken ct)
        => await FetchAsync(stream, consumer, new PullFetchRequest { Batch = batch }, ct);

    public async ValueTask<PullFetchBatch> FetchAsync(StreamHandle stream, ConsumerHandle consumer, PullFetchRequest request, CancellationToken ct)
    {
        var batch = Math.Max(request.Batch, 1);
        var messages = new List<StoredMessage>(batch);

        // Go: consumer.go — enforce ExpiresMs timeout on pull fetch requests.
        // When ExpiresMs > 0, create a linked CancellationTokenSource that fires
        // after the timeout. If it fires before the batch is full, return partial
        // results with TimedOut = true.
        using var expiresCts = request.ExpiresMs > 0
            ? CancellationTokenSource.CreateLinkedTokenSource(ct)
            : null;
        if (expiresCts is not null)
            expiresCts.CancelAfter(request.ExpiresMs);

        var effectiveCt = expiresCts?.Token ?? ct;

        if (consumer.NextSequence == 1)
        {
            consumer.NextSequence = await ResolveInitialSequenceAsync(stream, consumer.Config, effectiveCt);
        }

        if (request.NoWait)
        {
            var available = await stream.Store.LoadAsync(consumer.NextSequence, effectiveCt);
            if (available == null)
                return new PullFetchBatch([], timedOut: false);
        }

        if (consumer.Config.AckPolicy == AckPolicy.Explicit)
        {
            if (consumer.AckProcessor.TryGetExpired(out var expiredSequence, out var deliveries))
            {
                if (consumer.Config.MaxDeliver > 0 && deliveries > consumer.Config.MaxDeliver)
                {
                    consumer.AckProcessor.AckAll(expiredSequence);
                    return new PullFetchBatch(messages);
                }

                var backoff = consumer.Config.BackOffMs.Count >= deliveries
                    ? consumer.Config.BackOffMs[deliveries - 1]
                    : consumer.Config.AckWaitMs;
                consumer.AckProcessor.ScheduleRedelivery(expiredSequence, backoff);

                var redelivery = await stream.Store.LoadAsync(expiredSequence, effectiveCt);
                if (redelivery != null)
                {
                    messages.Add(new StoredMessage
                    {
                        Sequence = redelivery.Sequence,
                        Subject = redelivery.Subject,
                        Payload = redelivery.Payload,
                        Redelivered = true,
                    });
                }

                return new PullFetchBatch(messages);
            }

            if (consumer.AckProcessor.HasPending)
                return new PullFetchBatch(messages);
        }

        // Use CompiledFilter for efficient multi-filter matching
        var compiledFilter = CompiledFilter.FromConfig(consumer.Config);
        var sequence = consumer.NextSequence;

        // Go: consumer.go — MaxBytes caps the total byte payload returned in one pull request
        var remainingBytes = request.MaxBytes > 0 ? request.MaxBytes : long.MaxValue;

        try
        {
            for (var i = 0; i < batch; i++)
            {
                StoredMessage? message;

                // Go: consumer.go — when ExpiresMs is set, retry loading until a message
                // appears or the timeout fires. This handles the case where the stream
                // is empty or the consumer has caught up to the end of the stream.
                if (expiresCts is not null)
                {
                    message = await WaitForMessageAsync(stream.Store, sequence, effectiveCt);
                }
                else
                {
                    message = await stream.Store.LoadAsync(sequence, effectiveCt);
                }

                if (message == null)
                    break;

                if (!compiledFilter.Matches(message.Subject))
                {
                    sequence++;
                    i--;
                    continue;
                }

                if (message.Sequence <= consumer.AckProcessor.AckFloor)
                {
                    sequence++;
                    i--;
                    continue;
                }

                // Go: consumer.go — stop delivery if adding this message would exceed MaxBytes
                if (request.MaxBytes > 0)
                {
                    var msgSize = message.Payload.Length + message.Subject.Length;
                    if (msgSize > remainingBytes)
                        break;
                    remainingBytes -= msgSize;
                }

                if (consumer.Config.ReplayPolicy == ReplayPolicy.Original)
                    await Task.Delay(60, effectiveCt);

                messages.Add(message);
                if (consumer.Config.AckPolicy is AckPolicy.Explicit or AckPolicy.All)
                {
                    if (consumer.Config.MaxAckPending > 0 && consumer.AckProcessor.PendingCount >= consumer.Config.MaxAckPending)
                        break;
                    consumer.AckProcessor.Register(message.Sequence, consumer.Config.AckWaitMs);
                }
                sequence++;
            }
        }
        catch (OperationCanceledException) when (expiresCts is not null && expiresCts.IsCancellationRequested && !ct.IsCancellationRequested)
        {
            // ExpiresMs timeout fired — return partial results
            consumer.NextSequence = sequence;
            return new PullFetchBatch(messages, timedOut: true);
        }

        consumer.NextSequence = sequence;
        return new PullFetchBatch(messages);
    }

    /// <summary>
    /// Poll-wait for a message to appear at the given sequence, retrying with a
    /// short delay until the cancellation token fires (typically from ExpiresMs).
    /// </summary>
    private static async ValueTask<StoredMessage?> WaitForMessageAsync(IStreamStore store, ulong sequence, CancellationToken ct)
    {
        while (!ct.IsCancellationRequested)
        {
            var message = await store.LoadAsync(sequence, ct);
            if (message is not null)
                return message;

            // Yield briefly before retrying — the ExpiresMs CTS will cancel when time is up
            await Task.Delay(5, ct).ConfigureAwait(false);
        }

        return null;
    }

    private static async ValueTask<ulong> ResolveInitialSequenceAsync(StreamHandle stream, ConsumerConfig config, CancellationToken ct)
    {
        var state = await stream.Store.GetStateAsync(ct);
        return config.DeliverPolicy switch
        {
            DeliverPolicy.Last when state.LastSeq > 0 => state.LastSeq,
            DeliverPolicy.New when state.LastSeq > 0 => state.LastSeq + 1,
            DeliverPolicy.ByStartSequence when config.OptStartSeq > 0 => config.OptStartSeq,
            DeliverPolicy.ByStartTime when config.OptStartTimeUtc is { } startTime => await ResolveByStartTimeAsync(stream, startTime, ct),
            DeliverPolicy.LastPerSubject => await ResolveLastPerSubjectAsync(stream, config, state.LastSeq, ct),
            // Go: consumer.go — DeliverAll starts from stream's FirstSeq (not always 1).
            // After purge, FirstSeq advances past deleted messages.
            _ => state.FirstSeq > 0 ? state.FirstSeq : 1,
        };
    }

    private static async ValueTask<ulong> ResolveLastPerSubjectAsync(
        StreamHandle stream,
        ConsumerConfig config,
        ulong fallbackSequence,
        CancellationToken ct)
    {
        var subject = config.ResolvePrimaryFilterSubject();
        if (string.IsNullOrWhiteSpace(subject))
            return fallbackSequence > 0 ? fallbackSequence : 1UL;

        var last = await stream.Store.LoadLastBySubjectAsync(subject, ct);
        if (last != null)
            return last.Sequence;

        return fallbackSequence > 0 ? fallbackSequence : 1UL;
    }

    private static async ValueTask<ulong> ResolveByStartTimeAsync(StreamHandle stream, DateTime startTimeUtc, CancellationToken ct)
    {
        var messages = await stream.Store.ListAsync(ct);
        var match = messages.FirstOrDefault(m => m.TimestampUtc >= startTimeUtc);
        return match?.Sequence ?? 1UL;
    }
}

public sealed class PullFetchBatch
{
    public IReadOnlyList<StoredMessage> Messages { get; }
    public bool TimedOut { get; }

    public PullFetchBatch(IReadOnlyList<StoredMessage> messages, bool timedOut = false)
    {
        Messages = messages;
        TimedOut = timedOut;
    }
}

public sealed class PullFetchRequest
{
    public int Batch { get; init; } = 1;
    public bool NoWait { get; init; }
    public int ExpiresMs { get; init; }
    // Go: consumer.go — max_bytes limits total bytes per fetch request
    // Reference: golang/nats-server/server/consumer.go — maxRequestBytes
    public long MaxBytes { get; init; }
}

// Go: consumer.go — pull wait queue for pending pull requests with priority ordering
// Reference: golang/nats-server/server/consumer.go waitQueue + addPrioritized + popAndRequeue
public sealed class PullRequestWaitQueue
{
    private readonly int _maxSize;
    private readonly List<PullWaitingRequest> _items = new();

    public PullRequestWaitQueue(int maxSize = int.MaxValue) => _maxSize = maxSize;

    public int Count => _items.Count;

    /// <summary>
    /// Enqueue a waiting pull request using stable priority ordering (lower Priority value = higher precedence).
    /// Returns false if the queue is at capacity.
    /// Go: consumer.go — waitQueue.addPrioritized with sort.SliceStable semantics.
    /// </summary>
    public bool Enqueue(PullWaitingRequest request)
    {
        if (_maxSize > 0 && _items.Count >= _maxSize)
            return false;

        // Stable insertion sort: find first item with strictly higher priority value, insert before it
        var insertAt = _items.Count;
        for (var i = 0; i < _items.Count; i++)
        {
            if (_items[i].Priority > request.Priority)
            {
                insertAt = i;
                break;
            }
        }
        _items.Insert(insertAt, request);
        return true;
    }

    public PullWaitingRequest? Peek()
        => _items.Count > 0 ? _items[0] : null;

    public PullWaitingRequest? Dequeue()
    {
        if (_items.Count == 0) return null;
        var head = _items[0];
        _items.RemoveAt(0);
        return head;
    }

    /// <summary>
    /// Pop the head item, decrement its RemainingBatch, and re-insert at the END of
    /// its priority group if still > 0. Always returns the item with the decremented count.
    /// Go: consumer.go — waitQueue.popAndRequeue: decrements n, re-queues after same-priority
    /// items (round-robin within priority), returns item.
    /// </summary>
    public PullWaitingRequest? PopAndRequeue()
    {
        if (_items.Count == 0) return null;
        var head = _items[0];
        _items.RemoveAt(0);

        var decremented = head with { RemainingBatch = head.RemainingBatch - 1 };
        if (decremented.RemainingBatch > 0)
        {
            // Re-insert at the end of the same-priority group (round-robin within priority)
            var insertAt = _items.Count;
            for (var i = 0; i < _items.Count; i++)
            {
                if (_items[i].Priority > decremented.Priority)
                {
                    insertAt = i;
                    break;
                }
            }
            _items.Insert(insertAt, decremented);
        }

        return decremented;
    }

    public bool TryDequeue(out PullWaitingRequest? request)
    {
        request = Dequeue();
        return request is not null;
    }
}

// Go: consumer.go — a single queued pull request with batch/bytes/expires params
// Reference: golang/nats-server/server/consumer.go waitingRequest
public sealed record PullWaitingRequest
{
    public int Priority { get; init; }
    public int Batch { get; init; } = 1;
    public int RemainingBatch { get; init; } = 1;
    public long MaxBytes { get; init; }
    public int ExpiresMs { get; init; }
    public string? Reply { get; init; }
}
