// Copyright 2019-2026 The NATS Authors
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//
// Adapted from server/memstore.go (consumerMemStore)

namespace ZB.MOM.NatsNet.Server;

/// <summary>
/// In-memory implementation of <see cref="IConsumerStore"/>.
/// Stores consumer delivery and ack state in memory only.
/// </summary>
public sealed class ConsumerMemStore : IConsumerStore
{
    // -----------------------------------------------------------------------
    // Fields
    // -----------------------------------------------------------------------

    private readonly object _mu = new();
    private readonly JetStreamMemStore _ms;
    private ConsumerConfig _cfg;
    private ConsumerState _state = new();
    private bool _closed;

    // -----------------------------------------------------------------------
    // Constructor
    // -----------------------------------------------------------------------

    /// <summary>
    /// Creates a new consumer memory store backed by the given stream store.
    /// </summary>
    public ConsumerMemStore(JetStreamMemStore ms, ConsumerConfig cfg)
    {
        _ms = ms;
        _cfg = cfg;
    }

    // -----------------------------------------------------------------------
    // IConsumerStore — starting sequence
    // -----------------------------------------------------------------------

    /// <inheritdoc/>
    public void SetStarting(ulong sseq)
    {
        lock (_mu)
        {
            _state.Delivered.Stream = sseq;
            _state.AckFloor.Stream = sseq;
        }
    }

    /// <inheritdoc/>
    public void UpdateStarting(ulong sseq)
    {
        lock (_mu)
        {
            if (sseq > _state.Delivered.Stream)
            {
                _state.Delivered.Stream = sseq;
                // For AckNone just update delivered and ackfloor at the same time.
                if (_cfg.AckPolicy == AckPolicy.AckNone)
                    _state.AckFloor.Stream = sseq;
            }
        }
    }

    /// <inheritdoc/>
    public void Reset(ulong sseq)
    {
        lock (_mu)
        {
            _state = new ConsumerState();
        }
        SetStarting(sseq);
    }

    // -----------------------------------------------------------------------
    // IConsumerStore — state query
    // -----------------------------------------------------------------------

    /// <inheritdoc/>
    public bool HasState()
    {
        lock (_mu)
        {
            return _state.Delivered.Consumer != 0 || _state.Delivered.Stream != 0;
        }
    }

    // -----------------------------------------------------------------------
    // IConsumerStore — delivery tracking
    // -----------------------------------------------------------------------

    /// <inheritdoc/>
    public void UpdateDelivered(ulong dseq, ulong sseq, ulong dc, long ts)
    {
        lock (_mu)
        {
            if (dc != 1 && _cfg.AckPolicy == AckPolicy.AckNone)
                throw StoreErrors.ErrNoAckPolicy;

            // Replay from old leader — ignore outdated updates.
            if (dseq <= _state.AckFloor.Consumer)
                return;

            if (_cfg.AckPolicy != AckPolicy.AckNone)
            {
                _state.Pending ??= new Dictionary<ulong, Pending>();

                if (sseq <= _state.Delivered.Stream)
                {
                    // Update to a previously delivered message.
                    if (_state.Pending.TryGetValue(sseq, out var p) && p != null)
                        p.Timestamp = ts;
                }
                else
                {
                    _state.Pending[sseq] = new Pending { Sequence = dseq, Timestamp = ts };
                }

                if (dseq > _state.Delivered.Consumer)
                    _state.Delivered.Consumer = dseq;
                if (sseq > _state.Delivered.Stream)
                    _state.Delivered.Stream = sseq;

                if (dc > 1)
                {
                    var maxdc = (ulong)_cfg.MaxDeliver;
                    if (maxdc > 0 && dc > maxdc)
                        _state.Pending.Remove(sseq);

                    _state.Redelivered ??= new Dictionary<ulong, ulong>();
                    if (!_state.Redelivered.TryGetValue(sseq, out var cur) || cur < dc - 1)
                        _state.Redelivered[sseq] = dc - 1;
                }
            }
            else
            {
                // AckNone — update delivered and ackfloor together.
                if (dseq > _state.Delivered.Consumer)
                {
                    _state.Delivered.Consumer = dseq;
                    _state.AckFloor.Consumer = dseq;
                }
                if (sseq > _state.Delivered.Stream)
                {
                    _state.Delivered.Stream = sseq;
                    _state.AckFloor.Stream = sseq;
                }
            }
        }
    }

    /// <inheritdoc/>
    public void UpdateAcks(ulong dseq, ulong sseq)
    {
        lock (_mu)
        {
            if (_cfg.AckPolicy == AckPolicy.AckNone)
                throw StoreErrors.ErrNoAckPolicy;

            // Ignore outdated acks.
            if (dseq <= _state.AckFloor.Consumer)
                return;

            if (_state.Pending == null || !_state.Pending.ContainsKey(sseq))
            {
                _state.Redelivered?.Remove(sseq);
                throw StoreErrors.ErrStoreMsgNotFound;
            }

            if (_cfg.AckPolicy == AckPolicy.AckAll)
            {
                var sgap = sseq - _state.AckFloor.Stream;
                _state.AckFloor.Consumer = dseq;
                _state.AckFloor.Stream = sseq;

                if (sgap > (ulong)_state.Pending.Count)
                {
                    var toRemove = new List<ulong>();
                    foreach (var kv in _state.Pending)
                        if (kv.Key <= sseq)
                            toRemove.Add(kv.Key);
                    foreach (var k in toRemove)
                    {
                        _state.Pending.Remove(k);
                        _state.Redelivered?.Remove(k);
                    }
                }
                else
                {
                    for (var seq = sseq; seq > sseq - sgap && _state.Pending.Count > 0; seq--)
                    {
                        _state.Pending.Remove(seq);
                        _state.Redelivered?.Remove(seq);
                        if (seq == 0) break;
                    }
                }
                return;
            }

            // AckExplicit
            if (_state.Pending.TryGetValue(sseq, out var pending) && pending != null)
            {
                _state.Pending.Remove(sseq);
                if (dseq > pending.Sequence && pending.Sequence > 0)
                    dseq = pending.Sequence; // Use the original delivery sequence.
            }

            if (_state.Pending.Count == 0)
            {
                _state.AckFloor.Consumer = _state.Delivered.Consumer;
                _state.AckFloor.Stream = _state.Delivered.Stream;
            }
            else if (dseq == _state.AckFloor.Consumer + 1)
            {
                _state.AckFloor.Consumer = dseq;
                _state.AckFloor.Stream = sseq;

                if (_state.Delivered.Consumer > dseq)
                {
                    for (var ss = sseq + 1; ss <= _state.Delivered.Stream; ss++)
                    {
                        if (_state.Pending.TryGetValue(ss, out var pp) && pp != null)
                        {
                            if (pp.Sequence > 0)
                            {
                                _state.AckFloor.Consumer = pp.Sequence - 1;
                                _state.AckFloor.Stream = ss - 1;
                            }
                            break;
                        }
                    }
                }
            }

            _state.Redelivered?.Remove(sseq);
        }
    }

    // -----------------------------------------------------------------------
    // IConsumerStore — config update
    // -----------------------------------------------------------------------

    /// <inheritdoc/>
    public void UpdateConfig(ConsumerConfig cfg)
    {
        lock (_mu)
        {
            _cfg = cfg;
        }
    }

    // -----------------------------------------------------------------------
    // IConsumerStore — update state
    // -----------------------------------------------------------------------

    /// <inheritdoc/>
    public void Update(ConsumerState state)
    {
        if (state.AckFloor.Consumer > state.Delivered.Consumer)
            throw new InvalidOperationException("bad ack floor for consumer");
        if (state.AckFloor.Stream > state.Delivered.Stream)
            throw new InvalidOperationException("bad ack floor for stream");

        Dictionary<ulong, Pending>? pending = null;
        Dictionary<ulong, ulong>? redelivered = null;

        if (state.Pending?.Count > 0)
        {
            pending = new Dictionary<ulong, Pending>(state.Pending.Count);
            foreach (var kv in state.Pending)
            {
                if (kv.Key <= state.AckFloor.Stream || kv.Key > state.Delivered.Stream)
                    throw new InvalidOperationException($"bad pending entry, sequence [{kv.Key}] out of range");
                pending[kv.Key] = new Pending { Sequence = kv.Value.Sequence, Timestamp = kv.Value.Timestamp };
            }
        }

        if (state.Redelivered?.Count > 0)
        {
            redelivered = new Dictionary<ulong, ulong>(state.Redelivered);
        }

        lock (_mu)
        {
            // Ignore outdated updates.
            if (state.Delivered.Consumer < _state.Delivered.Consumer ||
                state.AckFloor.Stream < _state.AckFloor.Stream)
                throw new InvalidOperationException("old update ignored");

            _state.Delivered = new SequencePair { Consumer = state.Delivered.Consumer, Stream = state.Delivered.Stream };
            _state.AckFloor = new SequencePair { Consumer = state.AckFloor.Consumer, Stream = state.AckFloor.Stream };
            _state.Pending = pending;
            _state.Redelivered = redelivered;
        }
    }

    // -----------------------------------------------------------------------
    // IConsumerStore — state retrieval
    // -----------------------------------------------------------------------

    /// <inheritdoc/>
    public (ConsumerState? State, Exception? Error) State() => StateWithCopy(doCopy: true);

    /// <inheritdoc/>
    public (ConsumerState? State, Exception? Error) BorrowState() => StateWithCopy(doCopy: false);

    private (ConsumerState? State, Exception? Error) StateWithCopy(bool doCopy)
    {
        lock (_mu)
        {
            if (_closed)
                return (null, StoreErrors.ErrStoreClosed);

            var state = new ConsumerState
            {
                Delivered = new SequencePair { Consumer = _state.Delivered.Consumer, Stream = _state.Delivered.Stream },
                AckFloor = new SequencePair { Consumer = _state.AckFloor.Consumer, Stream = _state.AckFloor.Stream },
            };

            if (_state.Pending?.Count > 0)
            {
                state.Pending = doCopy ? CopyPending() : _state.Pending;
            }

            if (_state.Redelivered?.Count > 0)
            {
                state.Redelivered = doCopy ? CopyRedelivered() : _state.Redelivered;
            }

            return (state, null);
        }
    }

    // -----------------------------------------------------------------------
    // IConsumerStore — encoding
    // -----------------------------------------------------------------------

    /// <inheritdoc/>
    public byte[] EncodedState()
    {
        lock (_mu)
        {
            if (_closed)
                throw StoreErrors.ErrStoreClosed;
            // Session 17 target: encode consumer state to binary form.
            return Array.Empty<byte>();
        }
    }

    // -----------------------------------------------------------------------
    // IConsumerStore — lifecycle
    // -----------------------------------------------------------------------

    /// <inheritdoc/>
    public StorageType Type() => StorageType.MemoryStorage;

    /// <inheritdoc/>
    public void Stop()
    {
        lock (_mu)
        {
            _closed = true;
        }
        _ms.RemoveConsumer(this);
    }

    /// <inheritdoc/>
    public void Delete() => Stop();

    /// <inheritdoc/>
    public void StreamDelete() => Stop();

    // -----------------------------------------------------------------------
    // Private helpers
    // -----------------------------------------------------------------------

    private Dictionary<ulong, Pending> CopyPending()
    {
        var pending = new Dictionary<ulong, Pending>(_state.Pending!.Count);
        foreach (var kv in _state.Pending!)
            pending[kv.Key] = new Pending { Sequence = kv.Value.Sequence, Timestamp = kv.Value.Timestamp };
        return pending;
    }

    private Dictionary<ulong, ulong> CopyRedelivered()
    {
        return new Dictionary<ulong, ulong>(_state.Redelivered!);
    }
}