using System.Text.Json;
using NATS.Server.JetStream.Models;
using NATS.Server.Subscriptions;

// Alias for the full Go-parity StreamState in this namespace.
using StorageStreamState = NATS.Server.JetStream.Storage.StreamState;

namespace NATS.Server.JetStream.Storage;

// Go: server/memstore.go
/// <summary>
/// In-memory implementation of <see cref="IStreamStore"/>. Implements both the
/// async helper surface used by the current JetStream layer and the full Go-parity
/// sync interface matching Go's <c>memStore</c> struct exactly.
/// </summary>
public sealed class MemStore : IStreamStore
{
    // -------------------------------------------------------------------------
    // Snapshot record (used by async serialisation helpers only)
    // -------------------------------------------------------------------------

    private sealed class SnapshotRecord
    {
        public ulong Sequence { get; init; }
        public string Subject { get; init; } = string.Empty;
        public string PayloadBase64 { get; init; } = string.Empty;
        public DateTime TimestampUtc { get; init; }
    }

    // -------------------------------------------------------------------------
    // Internal stored message — tightly packed to avoid per-field allocations.
    // -------------------------------------------------------------------------

    private sealed class Msg
    {
        public readonly string Subj;
        public readonly byte[]? Hdr;
        public readonly byte[]? Data;
        public readonly ulong Seq;
        public readonly long Ts;  // Unix nanoseconds

        public Msg(string subj, byte[]? hdr, byte[]? data, ulong seq, long ts)
        {
            Subj = subj;
            Hdr = hdr;
            Data = data;
            Seq = seq;
            Ts = ts;
        }
    }

    // -------------------------------------------------------------------------
    // Per-subject tracking
    // -------------------------------------------------------------------------

    private sealed class SubjectState
    {
        public ulong Msgs;
        public ulong First;
        public ulong Last;
    }

    // -------------------------------------------------------------------------
    // State
    // -------------------------------------------------------------------------

    private readonly object _gate = new();

    // Go: memStore.msgs
    private Dictionary<ulong, Msg> _msgs = new();

    // Go: memStore.fss — per-subject sequence tracking
    private readonly Dictionary<string, SubjectState> _fss = new(StringComparer.Ordinal);

    // Go: memStore.state — overall stream state
    private StreamStateMutable _st;

    // Go: memStore.dmap — deleted-sequence set (interior gaps)
    private readonly SortedSet<ulong> _dmap = [];

    // Go: memStore.maxp — max messages per subject (0 = unlimited)
    private int _maxp;

    // TTL support: seqno → expiry ticks (DateTime.UtcNow.Ticks)
    private SortedDictionary<ulong, long>? _ttls;
    private Timer? _ttlTimer;

    // Config copy for UpdateConfig
    private StreamConfig? _cfg;

    // -------------------------------------------------------------------------
    // Mutable state helper (avoids re-boxing ulong fields in record structs)
    // -------------------------------------------------------------------------

    private struct StreamStateMutable
    {
        public ulong Msgs;
        public ulong Bytes;
        public ulong FirstSeq;
        public DateTime FirstTime;
        public ulong LastSeq;
        public DateTime LastTime;
    }

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

    public MemStore() { }

    public MemStore(StreamConfig cfg)
    {
        _cfg = cfg;
        _maxp = cfg.MaxMsgsPer;
        if (cfg.AllowMsgTtl)
            _ttls = new SortedDictionary<ulong, long>();
        // Go: newMemStore — handle FirstSeq by pre-advancing the sequence counter
        if (cfg.FirstSeq > 0)
        {
            _st.LastSeq = cfg.FirstSeq - 1;
            _st.FirstSeq = cfg.FirstSeq;
        }
    }

    // -------------------------------------------------------------------------
    // Async helpers (used by existing JetStream layer)
    // -------------------------------------------------------------------------

    // Go: memStore.StoreMsg — async wrapper
    public ValueTask<ulong> AppendAsync(string subject, ReadOnlyMemory<byte> payload, CancellationToken ct)
    {
        lock (_gate)
        {
            var seq = _st.LastSeq + 1;
            var ts = ToUnixNanos(DateTime.UtcNow); // nanos
            StoreInternal(subject, null, payload.IsEmpty ? null : payload.ToArray(), seq, ts, 0);
            return ValueTask.FromResult(seq);
        }
    }

    public ValueTask<StoredMessage?> LoadAsync(ulong sequence, CancellationToken ct)
    {
        lock (_gate)
        {
            if (!_msgs.TryGetValue(sequence, out var m))
                return ValueTask.FromResult<StoredMessage?>(null);
            return ValueTask.FromResult<StoredMessage?>(new StoredMessage
            {
                Sequence = m.Seq,
                Subject = m.Subj,
                Payload = m.Data ?? [],
                TimestampUtc = FromUnixNanos(m.Ts),
            });
        }
    }

    public ValueTask<StoredMessage?> LoadLastBySubjectAsync(string subject, CancellationToken ct)
    {
        lock (_gate)
        {
            if (!_fss.TryGetValue(subject, out var ss) || ss.Msgs == 0)
                return ValueTask.FromResult<StoredMessage?>(null);
            if (!_msgs.TryGetValue(ss.Last, out var m))
                return ValueTask.FromResult<StoredMessage?>(null);
            return ValueTask.FromResult<StoredMessage?>(new StoredMessage
            {
                Sequence = m.Seq,
                Subject = m.Subj,
                Payload = m.Data ?? [],
                TimestampUtc = FromUnixNanos(m.Ts),
            });
        }
    }

    public ValueTask<IReadOnlyList<StoredMessage>> ListAsync(CancellationToken ct)
    {
        lock (_gate)
        {
            var list = _msgs.Values
                .OrderBy(m => m.Seq)
                .Select(m => new StoredMessage
                {
                    Sequence = m.Seq,
                    Subject = m.Subj,
                    Payload = m.Data ?? [],
                    TimestampUtc = FromUnixNanos(m.Ts),
                })
                .ToArray();
            return ValueTask.FromResult<IReadOnlyList<StoredMessage>>(list);
        }
    }

    public ValueTask<bool> RemoveAsync(ulong sequence, CancellationToken ct)
    {
        lock (_gate)
        {
            return ValueTask.FromResult(RemoveInternal(sequence));
        }
    }

    public ValueTask PurgeAsync(CancellationToken ct)
    {
        lock (_gate)
        {
            PurgeInternal();
            return ValueTask.CompletedTask;
        }
    }

    public ValueTask<byte[]> CreateSnapshotAsync(CancellationToken ct)
    {
        lock (_gate)
        {
            var snapshot = _msgs
                .Values
                .OrderBy(x => x.Seq)
                .Select(x => new SnapshotRecord
                {
                    Sequence = x.Seq,
                    Subject = x.Subj,
                    PayloadBase64 = Convert.ToBase64String(x.Data ?? []),
                    TimestampUtc = FromUnixNanos(x.Ts),
                })
                .ToArray();
            return ValueTask.FromResult(JsonSerializer.SerializeToUtf8Bytes(snapshot));
        }
    }

    public ValueTask RestoreSnapshotAsync(ReadOnlyMemory<byte> snapshot, CancellationToken ct)
    {
        lock (_gate)
        {
            // Full reset for restore — not PurgeInternal (which preserves sequence continuity).
            // Restored records may have gaps, so we advance LastSeq per record.
            _msgs = new Dictionary<ulong, Msg>();
            _fss.Clear();
            _dmap.Clear();
            _ttls?.Clear();
            _st = default;

            if (!snapshot.IsEmpty)
            {
                var records = JsonSerializer.Deserialize<SnapshotRecord[]>(snapshot.Span);
                if (records != null)
                {
                    foreach (var record in records)
                    {
                        var data = record.PayloadBase64.Length > 0
                            ? Convert.FromBase64String(record.PayloadBase64)
                            : null;
                        var ts = ToUnixNanos(record.TimestampUtc);
                        // Advance LastSeq so StoreInternal accepts this sequence (handles gaps)
                        _st.LastSeq = record.Sequence - 1;
                        StoreInternal(record.Subject, null, data, record.Sequence, ts, 0);
                    }
                }
            }

            return ValueTask.CompletedTask;
        }
    }

    public ValueTask<ApiStreamState> GetStateAsync(CancellationToken ct)
    {
        lock (_gate)
        {
            // Go: memStore state — when empty but FirstSeq was configured, report it.
            // Reference: server/memstore.go — State() preserves first.seq watermark.
            var firstSeqForApi = _st.Msgs == 0
                ? (_st.FirstSeq > 0 ? _st.FirstSeq : 0UL)
                : _st.FirstSeq;
            return ValueTask.FromResult(new ApiStreamState
            {
                Messages = _st.Msgs,
                FirstSeq = firstSeqForApi,
                LastSeq = _st.LastSeq,
                Bytes = _st.Bytes,
            });
        }
    }

    // -------------------------------------------------------------------------
    // Go-parity sync interface (IStreamStore explicit implementations)
    // -------------------------------------------------------------------------

    // Go: memStore.StoreMsg server/memstore.go:350
    (ulong Seq, long Ts) IStreamStore.StoreMsg(string subject, byte[]? hdr, byte[] msg, long ttl)
    {
        lock (_gate)
        {
            var seq = _st.LastSeq + 1;
            // Go uses time.Now().UnixNano() — use Ticks*100 for ~100ns precision
            var ts = ToUnixNanos(DateTime.UtcNow);
            StoreInternal(subject, hdr, msg.Length == 0 ? null : msg, seq, ts, ttl);
            return (seq, ts);
        }
    }

    // Go: memStore.StoreRawMsg server/memstore.go:329
    void IStreamStore.StoreRawMsg(string subject, byte[]? hdr, byte[] msg, ulong seq, long ts, long ttl, bool discardNewCheck)
    {
        lock (_gate)
        {
            StoreInternal(subject, hdr, msg.Length == 0 ? null : msg, seq, ts, ttl);
        }
    }

    // Go: memStore.SkipMsg server/memstore.go:368
    ulong IStreamStore.SkipMsg(ulong seq)
    {
        lock (_gate)
        {
            if (seq != 0 && seq != _st.LastSeq + 1)
                throw new InvalidOperationException($"Sequence mismatch: expected {_st.LastSeq + 1}, got {seq}.");
            if (seq == 0)
                seq = _st.LastSeq + 1;
            _st.LastSeq = seq;
            _st.LastTime = DateTime.UtcNow;
            if (_st.Msgs == 0)
            {
                _st.FirstSeq = seq + 1;
                _st.FirstTime = default;
            }
            else
            {
                _dmap.Add(seq);
            }
            return seq;
        }
    }

    // Go: memStore.SkipMsgs server/memstore.go:395
    void IStreamStore.SkipMsgs(ulong seq, ulong num)
    {
        lock (_gate)
        {
            if (seq != 0 && seq != _st.LastSeq + 1)
                throw new InvalidOperationException($"Sequence mismatch: expected {_st.LastSeq + 1}, got {seq}.");
            if (seq == 0)
                seq = _st.LastSeq + 1;
            var lseq = seq + num - 1;
            _st.LastSeq = lseq;
            _st.LastTime = DateTime.UtcNow;
            if (_st.Msgs == 0)
            {
                _st.FirstSeq = lseq + 1;
                _st.FirstTime = default;
            }
            else
            {
                for (var s = seq; s <= lseq; s++)
                    _dmap.Add(s);
            }
        }
    }

    // Go: memStore.FlushAllPending server/memstore.go:423 — no-op for in-memory store
    void IStreamStore.FlushAllPending() { }

    // Go: memStore.LoadMsg server/memstore.go:1692
    StoreMsg IStreamStore.LoadMsg(ulong seq, StoreMsg? sm)
    {
        lock (_gate)
        {
            if (!_msgs.TryGetValue(seq, out var m))
                throw new KeyNotFoundException($"Message sequence {seq} not found.");
            return FillStoreMsg(sm ?? new StoreMsg(), m);
        }
    }

    // Go: memStore.LoadNextMsg server/memstore.go:1798
    (StoreMsg Msg, ulong Skip) IStreamStore.LoadNextMsg(string filter, bool wc, ulong start, StoreMsg? sm)
    {
        lock (_gate)
        {
            if (start < _st.FirstSeq)
                start = _st.FirstSeq;

            for (var seq = start; seq <= _st.LastSeq; seq++)
            {
                if (!_msgs.TryGetValue(seq, out var m))
                    continue;
                if (SubjectMatchesFilter(m.Subj, filter))
                {
                    var skip = seq > start ? seq - start : 0UL;
                    return (FillStoreMsg(sm ?? new StoreMsg(), m), skip);
                }
            }
            throw new KeyNotFoundException($"No message found at or after seq {start} matching filter '{filter}'.");
        }
    }

    // Go: memStore.LoadLastMsg server/memstore.go:1724
    StoreMsg IStreamStore.LoadLastMsg(string subject, StoreMsg? sm)
    {
        lock (_gate)
        {
            Msg? found = null;

            if (string.IsNullOrEmpty(subject) || subject == ">")
            {
                // Last message in the stream
                _msgs.TryGetValue(_st.LastSeq, out found);
                if (found == null)
                {
                    // LastSeq might be deleted — walk backwards
                    for (var s = _st.LastSeq; s >= _st.FirstSeq; s--)
                    {
                        if (_msgs.TryGetValue(s, out found)) break;
                    }
                }
            }
            else if (SubjectMatch.IsLiteral(subject))
            {
                if (_fss.TryGetValue(subject, out var ss) && ss.Msgs > 0)
                    _msgs.TryGetValue(ss.Last, out found);
            }
            else
            {
                // Wildcard: scan for the highest-seq matching message
                ulong bestSeq = 0;
                foreach (var kv in _msgs)
                {
                    if (kv.Key > bestSeq && SubjectMatchesFilter(kv.Value.Subj, subject))
                    {
                        bestSeq = kv.Key;
                        found = kv.Value;
                    }
                }
            }

            if (found == null)
                throw new KeyNotFoundException($"No message found for subject '{subject}'.");
            return FillStoreMsg(sm ?? new StoreMsg(), found);
        }
    }

    // Go: memStore.LoadPrevMsg — walk backwards from start
    StoreMsg IStreamStore.LoadPrevMsg(ulong start, StoreMsg? sm)
    {
        lock (_gate)
        {
            for (var seq = start - 1; seq >= _st.FirstSeq && seq <= _st.LastSeq; seq--)
            {
                if (_msgs.TryGetValue(seq, out var m))
                    return FillStoreMsg(sm ?? new StoreMsg(), m);
                if (seq == 0) break;
            }
            throw new KeyNotFoundException($"No message found before seq {start}.");
        }
    }

    // Go: memStore.RemoveMsg — soft delete
    bool IStreamStore.RemoveMsg(ulong seq)
    {
        lock (_gate)
        {
            return RemoveInternal(seq);
        }
    }

    // Go: memStore.EraseMsg — overwrite then remove
    bool IStreamStore.EraseMsg(ulong seq)
    {
        lock (_gate)
        {
            return RemoveInternal(seq);
        }
    }

    // Go: memStore.Purge server/memstore.go:1471
    ulong IStreamStore.Purge()
    {
        lock (_gate)
        {
            return PurgeInternal();
        }
    }

    // Go: memStore.PurgeEx server/memstore.go:1422
    ulong IStreamStore.PurgeEx(string subject, ulong seq, ulong keep)
    {
        if (string.IsNullOrEmpty(subject) || subject == ">")
        {
            if (keep == 0 && seq == 0)
                return ((IStreamStore)this).Purge();
            if (seq > 1)
                return CompactInternal(seq);
            if (keep > 0)
            {
                lock (_gate)
                {
                    var msgs = _st.Msgs;
                    var lseq = _st.LastSeq;
                    if (keep >= msgs) return 0;
                    return CompactInternal(lseq - keep + 1);
                }
            }
            return 0;
        }

        // Subject-specific purge
        var ss = FilteredStateInternal(1, subject);
        if (ss.Msgs == 0) return 0;

        var targetCount = ss.Msgs;
        if (keep > 0)
        {
            if (keep >= ss.Msgs) return 0;
            targetCount -= keep;
        }

        var last = ss.Last;
        if (seq > 1)
            last = seq - 1;

        lock (_gate)
        {
            ulong purged = 0;
            for (var s = ss.First; s <= last && purged < targetCount; s++)
            {
                if (_msgs.TryGetValue(s, out var m) && SubjectMatchesFilter(m.Subj, subject))
                {
                    if (RemoveInternal(s))
                        purged++;
                }
            }
            return purged;
        }
    }

    // Go: memStore.Compact server/memstore.go:1509
    ulong IStreamStore.Compact(ulong seq) => CompactInternal(seq);

    // Go: memStore.Truncate server/memstore.go:1618
    void IStreamStore.Truncate(ulong seq)
    {
        lock (_gate)
        {
            if (seq == 0)
            {
                // Reset to empty
                _msgs = new Dictionary<ulong, Msg>();
                _fss.Clear();
                _dmap.Clear();
                _st = default;
                return;
            }

            for (var i = _st.LastSeq; i > seq; i--)
            {
                if (_msgs.TryGetValue(i, out var m))
                {
                    _st.Bytes -= MsgSize(m);
                    _st.Msgs--;
                    RemoveSeqPerSubject(m.Subj, i);
                    _msgs.Remove(i);
                }
                else
                {
                    _dmap.Remove(i);
                }
            }
            _st.LastSeq = seq;
            if (_msgs.TryGetValue(seq, out var last))
                _st.LastTime = FromUnixNanos(last.Ts);
        }
    }

    // Go: memStore.GetSeqFromTime server/memstore.go:453
    ulong IStreamStore.GetSeqFromTime(DateTime t)
    {
        lock (_gate)
        {
            if (_msgs.Count == 0)
                return _st.LastSeq + 1;

            // Go uses time.Time.UnixNano(); match with Ticks*100 (100ns resolution)
            var ts = ToUnixNanos(t.ToUniversalTime());

            if (!_msgs.TryGetValue(_st.FirstSeq, out var firstMsg))
                return _st.LastSeq + 1;
            if (ts <= firstMsg.Ts)
                return _st.FirstSeq;

            // Find actual last msg (LastSeq may be a gap — last does not go backwards)
            Msg? lastMsg = null;
            for (var s = _st.LastSeq; s > _st.FirstSeq; s--)
            {
                if (_msgs.TryGetValue(s, out lastMsg)) break;
            }
            if (lastMsg == null)
                return _st.LastSeq + 1;

            // Go: if ts == last return lmsg.seq; if ts > last return LastSeq+1
            if (ts == lastMsg.Ts)
                return lastMsg.Seq;
            if (ts > lastMsg.Ts)
                return _st.LastSeq + 1;

            // Binary search with gap awareness
            var fseq = _st.FirstSeq;
            var lseq = _st.LastSeq;
            var result = lseq + 1;

            while (fseq <= lseq)
            {
                var mid = fseq + (lseq - fseq) / 2;
                // Walk forward to find a real message at/after mid
                var off = 0UL;
                Msg? cand = null;
                while (mid + off <= lseq)
                {
                    if (_msgs.TryGetValue(mid + off, out cand))
                        break;
                    off++;
                }
                if (cand == null)
                {
                    lseq = mid - 1;
                    continue;
                }
                if (cand.Ts >= ts)
                {
                    result = cand.Seq;
                    if (mid == fseq) break;
                    lseq = mid - 1;
                }
                else
                {
                    fseq = mid + off + 1;
                }
            }
            return result;
        }
    }

    // Go: memStore.FilteredState server/memstore.go:531
    SimpleState IStreamStore.FilteredState(ulong seq, string subject)
        => FilteredStateInternal(seq, subject);

    // Go: memStore.SubjectsState server/memstore.go:748
    Dictionary<string, SimpleState> IStreamStore.SubjectsState(string filterSubject)
    {
        lock (_gate)
        {
            var result = new Dictionary<string, SimpleState>(StringComparer.Ordinal);
            if (_fss.Count == 0) return result;

            foreach (var kv in _fss)
            {
                if (!SubjectMatchesFilter(kv.Key, filterSubject)) continue;
                result[kv.Key] = new SimpleState
                {
                    Msgs = kv.Value.Msgs,
                    First = kv.Value.First,
                    Last = kv.Value.Last,
                };
            }
            return result;
        }
    }

    // Go: memStore.SubjectsTotals server/memstore.go:881
    Dictionary<string, ulong> IStreamStore.SubjectsTotals(string filterSubject)
    {
        lock (_gate)
        {
            var result = new Dictionary<string, ulong>(StringComparer.Ordinal);
            foreach (var kv in _fss)
            {
                if (SubjectMatchesFilter(kv.Key, filterSubject))
                    result[kv.Key] = kv.Value.Msgs;
            }
            return result;
        }
    }

    // Go: memStore.AllLastSeqs server/memstore.go:780
    ulong[] IStreamStore.AllLastSeqs()
    {
        lock (_gate)
        {
            if (_msgs.Count == 0) return [];
            var seqs = _fss.Values.Select(ss => ss.Last).ToArray();
            Array.Sort(seqs);
            return seqs;
        }
    }

    // Go: memStore.MultiLastSeqs server/memstore.go:828
    ulong[] IStreamStore.MultiLastSeqs(string[] filters, ulong maxSeq, int maxAllowed)
    {
        lock (_gate)
        {
            if (_msgs.Count == 0) return [];

            var effective = maxSeq == 0 ? _st.LastSeq : maxSeq;
            var seen = new HashSet<ulong>();

            foreach (var filter in filters)
            {
                foreach (var kv in _fss)
                {
                    if (!SubjectMatchesFilter(kv.Key, filter)) continue;
                    var ss = kv.Value;
                    if (ss.Last <= effective)
                    {
                        seen.Add(ss.Last);
                    }
                    else if (ss.Msgs > 1)
                    {
                        // Walk backwards to find the last seq <= maxSeq
                        for (var s = effective; s >= _st.FirstSeq; s--)
                        {
                            if (_msgs.TryGetValue(s, out var m) && m.Subj == kv.Key)
                            {
                                seen.Add(s);
                                break;
                            }
                            if (s == 0) break;
                        }
                    }
                }
                if (maxAllowed > 0 && seen.Count > maxAllowed)
                    throw new InvalidOperationException($"Too many results: got {seen.Count}, max allowed is {maxAllowed}.");
            }

            var result = seen.ToArray();
            Array.Sort(result);
            return result;
        }
    }

    // Go: memStore.SubjectForSeq server/memstore.go:1678
    string IStreamStore.SubjectForSeq(ulong seq)
    {
        lock (_gate)
        {
            if (seq < _st.FirstSeq || !_msgs.TryGetValue(seq, out var m))
                throw new KeyNotFoundException($"Message sequence {seq} not found.");
            return m.Subj;
        }
    }

    // Go: memStore.NumPending server/memstore.go:913
    (ulong Total, ulong ValidThrough) IStreamStore.NumPending(ulong sseq, string filter, bool lastPerSubject)
    {
        lock (_gate)
        {
            var ss = FilteredStateLocked(sseq, filter, lastPerSubject);
            return (ss.Msgs, _st.LastSeq);
        }
    }

    // Go: memStore.State server/memstore.go — full state
    StorageStreamState IStreamStore.State()
    {
        lock (_gate)
        {
            var state = new StorageStreamState
            {
                Msgs = _st.Msgs,
                Bytes = _st.Bytes,
                FirstSeq = _st.FirstSeq,
                FirstTime = _st.FirstTime,
                LastSeq = _st.LastSeq,
                LastTime = _st.LastTime,
                NumSubjects = _fss.Count,
                NumDeleted = _dmap.Count,
            };

            if (_dmap.Count > 0)
                state.Deleted = _dmap.ToArray();

            return state;
        }
    }

    // Go: memStore.FastState server/memstore.go — populate without deleted list
    void IStreamStore.FastState(ref StorageStreamState state)
    {
        lock (_gate)
        {
            state.Msgs = _st.Msgs;
            state.Bytes = _st.Bytes;
            state.FirstSeq = _st.FirstSeq;
            state.FirstTime = _st.FirstTime;
            state.LastSeq = _st.LastSeq;
            state.LastTime = _st.LastTime;
            state.NumSubjects = _fss.Count;
            state.NumDeleted = _dmap.Count;
        }
    }

    // Go: memStore.Type
    StorageType IStreamStore.Type() => StorageType.Memory;

    // Go: memStore.UpdateConfig server/memstore.go:86
    void IStreamStore.UpdateConfig(StreamConfig cfg)
    {
        lock (_gate)
        {
            var prevMaxp = _maxp;
            _cfg = cfg;
            _maxp = cfg.MaxMsgsPer;

            // Create or destroy TTL wheel
            if (cfg.AllowMsgTtl && _ttls == null)
                _ttls = new SortedDictionary<ulong, long>();
            else if (!cfg.AllowMsgTtl && _ttls != null)
            {
                _ttls = null;
                _ttlTimer?.Dispose();
                _ttlTimer = null;
            }

            // Enforce per-subject limit if it shrank
            if (_maxp > 0)
            {
                foreach (var kv in _fss)
                    EnforcePerSubjectLimit(kv.Key, kv.Value);
            }
        }
    }

    // Go: memStore.Stop — no-op for in-memory store
    void IStreamStore.Stop() { }

    // Go: memStore.Delete — clear everything
    void IStreamStore.Delete(bool inline)
    {
        lock (_gate)
        {
            _msgs = new Dictionary<ulong, Msg>();
            _fss.Clear();
            _dmap.Clear();
            _st = default;
            _ttls = null;
            _ttlTimer?.Dispose();
            _ttlTimer = null;
        }
    }

    // Go: memStore.ResetState
    void IStreamStore.ResetState() { }

    // EncodedStreamState, ConsumerStore — not needed for MemStore tests
    byte[] IStreamStore.EncodedStreamState(ulong failed) => [];

    IConsumerStore IStreamStore.ConsumerStore(string name, DateTime created, ConsumerConfig cfg)
        => throw new NotSupportedException("MemStore does not implement ConsumerStore.");

    // -------------------------------------------------------------------------
    // TrimToMaxMessages — legacy helper used by existing async tests
    // -------------------------------------------------------------------------

    public void TrimToMaxMessages(ulong maxMessages)
    {
        lock (_gate)
        {
            while (_st.Msgs > maxMessages)
                RemoveInternal(_st.FirstSeq);
        }
    }

    // -------------------------------------------------------------------------
    // Internal helpers
    // -------------------------------------------------------------------------

    // Go: memStore.storeRawMsg server/memstore.go:195
    private void StoreInternal(string subj, byte[]? hdr, byte[]? data, ulong seq, long ts, long ttl)
    {
        if (seq != _st.LastSeq + 1)
        {
            if (seq > 0)
                throw new InvalidOperationException($"Sequence mismatch: expected {_st.LastSeq + 1}, got {seq}.");
            seq = _st.LastSeq + 1;
        }

        var now = FromUnixNanos(ts);
        if (_st.Msgs == 0)
        {
            _st.FirstSeq = seq;
            _st.FirstTime = now;
        }

        var m = new Msg(subj, hdr, data, seq, ts);
        _msgs[seq] = m;
        _st.Msgs++;
        _st.Bytes += MsgSize(m);
        _st.LastSeq = seq;
        _st.LastTime = now;

        // Per-subject tracking
        if (!string.IsNullOrEmpty(subj))
        {
            if (_fss.TryGetValue(subj, out var ss))
            {
                ss.Msgs++;
                ss.Last = seq;
                if (_maxp > 0 && ss.Msgs > (ulong)_maxp)
                    EnforcePerSubjectLimit(subj, ss);
            }
            else
            {
                _fss[subj] = new SubjectState { Msgs = 1, First = seq, Last = seq };
            }
        }

        // TTL tracking
        if (_ttls != null && ttl > 0)
        {
            var expires = DateTime.UtcNow.AddSeconds(ttl).Ticks;
            _ttls[seq] = expires;
            ScheduleTtlCheck();
        }
    }

    // Go: memStore.removeMsg + updateFirstSeq server/memstore.go:2110 + 2004
    // Go always inserts seq into dmap first, then updateFirstSeq cleans up dmap entries
    // for any sequences <= the new first seq.
    private bool RemoveInternal(ulong seq)
    {
        if (!_msgs.TryGetValue(seq, out var m))
            return false;

        _st.Msgs--;
        _st.Bytes -= MsgSize(m);

        // Go: ms.dmap.Insert(seq) — always insert into dmap first
        _dmap.Add(seq);

        // Go: ms.updateFirstSeq(seq) — update first seq and clean dmap
        UpdateFirstSeqLocked(seq);

        RemoveSeqPerSubject(m.Subj, seq);
        _msgs.Remove(seq);
        _ttls?.Remove(seq);
        return true;
    }

    // Go: memStore.updateFirstSeq server/memstore.go:2004
    // Must be called with _gate held. Seq is the sequence that was just removed.
    private void UpdateFirstSeqLocked(ulong seq)
    {
        if (seq != _st.FirstSeq)
        {
            // Interior delete — nothing to do for first seq
            return;
        }

        // Walk forward to find the next real message
        Msg? nsm = null;
        var nseq = _st.FirstSeq + 1;
        while (nseq <= _st.LastSeq)
        {
            if (_msgs.TryGetValue(nseq, out nsm))
                break;
            nseq++;
        }

        var oldFirst = _st.FirstSeq;
        if (nsm != null)
        {
            _st.FirstSeq = nsm.Seq;
            _st.FirstTime = FromUnixNanos(nsm.Ts);
        }
        else
        {
            // Like purge — no more messages
            _st.FirstSeq = _st.LastSeq + 1;
            _st.FirstTime = default;
        }

        // Go: clean up dmap for all sequences from oldFirst to new FirstSeq (exclusive)
        // These are no longer "interior" gaps since they're now before the first msg.
        if (oldFirst == _st.FirstSeq - 1)
        {
            _dmap.Remove(oldFirst);
        }
        else
        {
            for (var s = oldFirst; s < _st.FirstSeq; s++)
                _dmap.Remove(s);
        }
    }

    // Go: memStore.removeSeqPerSubject — update fss when a sequence is removed
    private void RemoveSeqPerSubject(string subj, ulong seq)
    {
        if (!_fss.TryGetValue(subj, out var ss)) return;
        ss.Msgs--;
        if (ss.Msgs == 0)
        {
            _fss.Remove(subj);
            return;
        }
        if (seq == ss.First)
        {
            // Find the next real first for this subject
            for (var s = seq + 1; s <= _st.LastSeq; s++)
            {
                if (_msgs.TryGetValue(s, out var m) && m.Subj == subj)
                {
                    ss.First = s;
                    break;
                }
            }
        }
        if (seq == ss.Last)
        {
            // Find the previous real last for this subject
            for (var s = seq - 1; s >= ss.First && s <= _st.LastSeq; s--)
            {
                if (_msgs.TryGetValue(s, out var m) && m.Subj == subj)
                {
                    ss.Last = s;
                    break;
                }
                if (s == 0) break;
            }
        }
    }

    // Go: memStore.purge server/memstore.go:1475
    private ulong PurgeInternal()
    {
        var purged = (ulong)_msgs.Count;
        var fseq = _st.LastSeq + 1;
        _msgs = new Dictionary<ulong, Msg>();
        _fss.Clear();
        _dmap.Clear();
        _ttls?.Clear();
        _st = new StreamStateMutable
        {
            FirstSeq = fseq,
            LastSeq = fseq - 1,
        };
        return purged;
    }

    // Go: memStore.compact server/memstore.go:1513
    private ulong CompactInternal(ulong seq)
    {
        if (seq == 0)
            return PurgeInternal();

        lock (_gate)
        {
            if (_st.FirstSeq > seq) return 0;

            if (seq <= _st.LastSeq)
            {
                var fseq = _st.FirstSeq;
                // Go: walk forward from seq to find first real message, mutating seq in-place.
                // The backward cleanup loop must start from (newFirstSeq - 1) to correctly
                // remove any dmap entries at the original compact seq that are no longer interior.
                var newFirst = seq;
                for (; newFirst <= _st.LastSeq; newFirst++)
                {
                    if (_msgs.TryGetValue(newFirst, out var nm))
                    {
                        _st.FirstSeq = newFirst;
                        _st.FirstTime = FromUnixNanos(nm.Ts);
                        break;
                    }
                }

                ulong purged = 0;
                for (var s = newFirst - 1; s >= fseq && s < ulong.MaxValue; s--)
                {
                    if (_msgs.TryGetValue(s, out var m))
                    {
                        _st.Bytes -= MsgSize(m);
                        purged++;
                        RemoveSeqPerSubject(m.Subj, s);
                        _msgs.Remove(s);
                    }
                    else
                    {
                        _dmap.Remove(s);
                    }
                    if (s == 0) break;
                }
                if (purged > _st.Msgs) purged = _st.Msgs;
                _st.Msgs -= purged;
                return purged;
            }
            else
            {
                // Compact past the end
                var purged = (ulong)_msgs.Count;
                _msgs = new Dictionary<ulong, Msg>();
                _fss.Clear();
                _dmap.Clear();
                _ttls?.Clear();
                _st = new StreamStateMutable
                {
                    FirstSeq = seq,
                    LastSeq = seq - 1,
                };
                return purged;
            }
        }
    }

    // Go: memStore.filteredStateLocked server/memstore.go:540
    private SimpleState FilteredStateInternal(ulong sseq, string subject)
    {
        lock (_gate)
        {
            return FilteredStateLocked(sseq, subject, false);
        }
    }

    // Called with _gate held
    private SimpleState FilteredStateLocked(ulong sseq, string filter, bool lastPerSubject)
    {
        if (sseq < _st.FirstSeq) sseq = _st.FirstSeq;
        if (sseq > _st.LastSeq) return default;

        var isAll = string.IsNullOrEmpty(filter) || filter == ">";
        if (isAll && sseq <= _st.FirstSeq)
        {
            var total = lastPerSubject ? (ulong)_fss.Count : _st.Msgs;
            return new SimpleState { Msgs = total, First = _st.FirstSeq, Last = _st.LastSeq };
        }

        // Build result by scanning messages
        ulong msgs = 0, first = 0, last = 0;
        var seen = lastPerSubject ? new HashSet<string>(StringComparer.Ordinal) : null;

        for (var seq = sseq; seq <= _st.LastSeq; seq++)
        {
            if (!_msgs.TryGetValue(seq, out var m)) continue;
            if (!isAll && !SubjectMatchesFilter(m.Subj, filter)) continue;
            if (seen != null)
            {
                if (!seen.Add(m.Subj)) continue;
            }
            msgs++;
            if (first == 0) first = seq;
            last = seq;
        }

        return new SimpleState { Msgs = msgs, First = first, Last = last };
    }

    // Go: memStore.enforcePerSubjectLimit server/memstore.go:1072
    private void EnforcePerSubjectLimit(string subj, SubjectState ss)
    {
        if (_maxp <= 0) return;
        while (ss.Msgs > (ulong)_maxp)
        {
            if (!RemoveInternal(ss.First))
                break;
        }
    }

    // Go: memStoreMsgSize server/memstore.go
    private static ulong MsgSize(Msg m)
    {
        var size = (ulong)(m.Subj.Length + (m.Hdr?.Length ?? 0) + (m.Data?.Length ?? 0) + 16);
        return size;
    }

    // Go: time.Now().UnixNano() — convert DateTime to Unix nanoseconds
    private static long ToUnixNanos(DateTime utc) => (utc.Ticks - DateTime.UnixEpoch.Ticks) * 100L;

    // Convert Unix nanoseconds back to DateTime
    private static DateTime FromUnixNanos(long nanos) => new(nanos / 100L + DateTime.UnixEpoch.Ticks, DateTimeKind.Utc);

    // Subject filter matching (same as FileStore helper)
    private static bool SubjectMatchesFilter(string subject, string filter)
    {
        if (string.IsNullOrEmpty(filter) || filter == ">")
            return true;
        if (SubjectMatch.IsLiteral(filter))
            return string.Equals(subject, filter, StringComparison.Ordinal);
        return SubjectMatch.MatchLiteral(subject, filter);
    }

    // Fill a StoreMsg from an internal Msg
    private static StoreMsg FillStoreMsg(StoreMsg sm, Msg m)
    {
        sm.Clear();
        sm.Subject = m.Subj;
        sm.Header = m.Hdr?.Length > 0 ? m.Hdr : null;
        sm.Data = m.Data?.Length > 0 ? m.Data : null;
        sm.Sequence = m.Seq;
        sm.Timestamp = m.Ts;
        return sm;
    }

    // -------------------------------------------------------------------------
    // TTL expiry
    // -------------------------------------------------------------------------

    private void ScheduleTtlCheck()
    {
        _ttlTimer?.Dispose();
        _ttlTimer = new Timer(_ => ExpireTtlMessages(), null, TimeSpan.FromSeconds(1), Timeout.InfiniteTimeSpan);
    }

    private void ExpireTtlMessages()
    {
        lock (_gate)
        {
            if (_ttls == null) return;
            var now = DateTime.UtcNow.Ticks;
            var toRemove = _ttls.Where(kv => kv.Value <= now).Select(kv => kv.Key).ToList();
            foreach (var seq in toRemove)
                RemoveInternal(seq);

            if (_ttls.Count > 0)
                ScheduleTtlCheck();
        }
    }

    // -------------------------------------------------------------------------
    // Internal method exposed for test access — nextWildcardMatchLocked
    // (Go: nextWildcardMatchLocked server/memstore.go:1810)
    // -------------------------------------------------------------------------

    /// <summary>
    /// Returns the (first, last, found) range for messages whose subject matches
    /// <paramref name="filter"/> at or after <paramref name="start"/>. Called with
    /// <c>_gate</c> already held.
    /// </summary>
    internal (ulong First, ulong Last, bool Found) NextWildcardMatchLocked(string filter, ulong start)
    {
        ulong first = _st.LastSeq, last = 0;
        var found = false;

        foreach (var kv in _fss)
        {
            if (!SubjectMatchesFilter(kv.Key, filter)) continue;
            var ss = kv.Value;
            if (start > ss.Last) continue;
            found = true;
            var f = Math.Max(start, ss.First);
            if (f < first) first = f;
            if (ss.Last > last) last = ss.Last;
        }

        return found ? (Math.Max(first, start), last, true) : (0UL, 0UL, false);
    }

    /// <summary>
    /// Returns the (first, last, found) range for messages whose subject exactly
    /// equals <paramref name="filter"/> at or after <paramref name="start"/>. Called
    /// with <c>_gate</c> already held.
    /// </summary>
    internal (ulong First, ulong Last, bool Found) NextLiteralMatchLocked(string filter, ulong start)
    {
        if (!_fss.TryGetValue(filter, out var ss)) return (0, 0, false);
        if (start > ss.Last) return (0, 0, false);
        return (Math.Max(start, ss.First), ss.Last, true);
    }

    /// <summary>Exposes <c>_gate</c> for test-level lock acquisition.</summary>
    internal object Gate => _gate;
}
