// Copyright 2025 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/jetstream_batching.go in the NATS server Go source.

using System.Numerics;
using System.Text;
using System.Text.Json;
using ZB.MOM.NatsNet.Server.Internal.DataStructures;

namespace ZB.MOM.NatsNet.Server;

internal static class JetStreamBatching
{
    private const string JsStreamSource = "Nats-Stream-Source";
    private const string JsMessageCounterSources = "Nats-Counter-Sources";

    internal static readonly TimeSpan StreamDefaultMaxBatchTimeout = TimeSpan.FromSeconds(10);
    internal const string BatchTimeout = "timeout";

    private static int _globalInflightBatches;

    internal static int GlobalInflightBatches => Volatile.Read(ref _globalInflightBatches);

    internal static (string BatchName, string StoreDir) GetBatchStoreDir(string storeRootDir, string streamName, string batchId)
    {
        if (string.IsNullOrWhiteSpace(storeRootDir))
            throw new ArgumentException("store root directory is required", nameof(storeRootDir));
        if (string.IsNullOrWhiteSpace(streamName))
            throw new ArgumentException("stream name is required", nameof(streamName));
        if (string.IsNullOrWhiteSpace(batchId))
            throw new ArgumentException("batch ID is required", nameof(batchId));

        var batchName = NatsServer.GetHash(batchId);
        var batchPath = Path.Combine(storeRootDir, "_streams", streamName, "batches", batchName);
        return (batchName, batchPath);
    }

    internal static IStreamStore NewBatchStore(string storeRootDir, StreamConfig streamConfig, string batchId)
    {
        ArgumentNullException.ThrowIfNull(streamConfig);

        if (streamConfig.Replicas == 1 && streamConfig.Storage == StorageType.FileStorage)
        {
            var (batchName, storeDir) = GetBatchStoreDir(storeRootDir, streamConfig.Name, batchId);
            Directory.CreateDirectory(storeDir);

            var cfg = new FileStoreConfig
            {
                AsyncFlush = true,
                BlockSize = FileStoreDefaults.DefaultLargeBlockSize,
                StoreDir = storeDir,
            };

            var fsInfo = new FileStreamInfo
            {
                Created = DateTime.UtcNow,
                Config = new StreamConfig
                {
                    Name = batchName,
                    Storage = StorageType.FileStorage,
                },
            };

            return new JetStreamFileStore(cfg, fsInfo);
        }

        return new JetStreamMemStore(new StreamConfig
        {
            Name = string.Empty,
            Storage = StorageType.MemoryStorage,
        });
    }

    internal static void IncrementGlobalInflightBatches() => Interlocked.Increment(ref _globalInflightBatches);

    internal static void DecrementGlobalInflightBatches() => Interlocked.Decrement(ref _globalInflightBatches);

    internal static void ResetGlobalInflightBatchesForTest() => Interlocked.Exchange(ref _globalInflightBatches, 0);

    internal static void IncrementGlobalInflightBatchesForTest() => IncrementGlobalInflightBatches();

    internal static BatchGroup NewBatchGroup(
        Batching batches,
        string storeRootDir,
        StreamConfig streamConfig,
        string batchId,
        Action<string, string>? abandonedAdvisory = null,
        TimeSpan? maxBatchTimeout = null)
    {
        ArgumentNullException.ThrowIfNull(batches);
        return batches.NewBatchGroup(storeRootDir, streamConfig, batchId, abandonedAdvisory, maxBatchTimeout);
    }

    internal static bool ReadyForCommit(BatchGroup batchGroup)
    {
        ArgumentNullException.ThrowIfNull(batchGroup);
        return batchGroup.ReadyForCommit();
    }

    internal static void Cleanup(BatchGroup batchGroup, string batchId, Batching batches)
    {
        ArgumentNullException.ThrowIfNull(batchGroup);
        batchGroup.Cleanup(batchId, batches);
    }

    internal static void CleanupLocked(BatchGroup batchGroup, string batchId, Batching batches)
    {
        ArgumentNullException.ThrowIfNull(batchGroup);
        batchGroup.CleanupLocked(batchId, batches);
    }

    internal static void StopLocked(BatchGroup batchGroup)
    {
        ArgumentNullException.ThrowIfNull(batchGroup);
        batchGroup.StopLocked();
    }

    internal static void Commit(BatchStagedDiff diff, BatchRuntimeState state)
    {
        ArgumentNullException.ThrowIfNull(diff);
        ArgumentNullException.ThrowIfNull(state);

        if (diff.MsgIds is { Count: > 0 })
        {
            var now = DateTimeOffset.UtcNow.ToUnixTimeMilliseconds() * 1_000_000L;
            foreach (var msgId in diff.MsgIds.Keys)
                state.MsgIds[msgId] = new NatsStream.DedupeEntry { Id = msgId, Seq = 0, TimestampNanos = now };
        }

        if (diff.Counter is { Count: > 0 })
        {
            foreach (var entry in diff.Counter)
                state.ClusteredCounterTotal[entry.Key] = entry.Value;
        }

        if (diff.Inflight is { Count: > 0 })
        {
            foreach (var pair in diff.Inflight)
            {
                if (state.Inflight.TryGetValue(pair.Key, out var existing))
                {
                    existing.Bytes += pair.Value.Bytes;
                    existing.Ops += pair.Value.Ops;
                }
                else
                {
                    state.Inflight[pair.Key] = pair.Value;
                }
            }
        }

        if (diff.ExpectedPerSubject is { Count: > 0 })
        {
            foreach (var pair in diff.ExpectedPerSubject)
            {
                state.ExpectedPerSubjectSequence[pair.Value.ClSeq] = pair.Key;
                state.ExpectedPerSubjectInProcess.Add(pair.Key);
            }
        }
    }

    internal static void ClearBatchStateLocked(BatchApply batch)
    {
        ArgumentNullException.ThrowIfNull(batch);
        batch.ClearBatchStateLocked();
    }

    internal static void RejectBatchStateLocked(BatchApply batch, BatchApplyContext context)
    {
        ArgumentNullException.ThrowIfNull(batch);
        ArgumentNullException.ThrowIfNull(context);
        batch.RejectBatchStateLocked(context);
    }

    internal static void RejectBatchState(BatchApply batch, BatchApplyContext context)
    {
        ArgumentNullException.ThrowIfNull(batch);
        ArgumentNullException.ThrowIfNull(context);
        batch.RejectBatchState(context);
    }

    internal static (byte[] Hdr, byte[] Msg, ulong Sequence, JsApiError? ApiError, Exception? Error) CheckMsgHeadersPreClusteredProposal(
        BatchStagedDiff diff,
        BatchHeaderCheckContext context,
        string subject,
        byte[]? hdr,
        byte[]? msg,
        bool sourced,
        string name,
        bool allowRollup,
        bool denyPurge,
        bool allowTtl,
        bool allowMsgCounter,
        bool allowMsgSchedules,
        DiscardPolicy discard,
        bool discardNewPer,
        int maxMsgSize,
        long maxMsgs,
        long maxMsgsPer,
        long maxBytes)
    {
        ArgumentNullException.ThrowIfNull(diff);
        ArgumentNullException.ThrowIfNull(context);
        ArgumentNullException.ThrowIfNull(context.Store);
        ArgumentException.ThrowIfNullOrWhiteSpace(subject);
        ArgumentException.ThrowIfNullOrWhiteSpace(name);

        hdr ??= Array.Empty<byte>();
        msg ??= Array.Empty<byte>();
        BigInteger? incr = null;

        if (hdr.Length > 0)
        {
            if (hdr.Length > ushort.MaxValue)
            {
                var err = new InvalidOperationException($"JetStream header size exceeds limits for '{context.AccountName} > {name}'");
                return (hdr, msg, 0, JsApiErrors.NewJSStreamHeaderExceedsMaximumError(), err);
            }

            var (parsedIncr, okIncr) = NatsStream.GetMessageIncr(hdr);
            if (!okIncr)
            {
                var apiErr = JsApiErrors.NewJSMessageIncrInvalidError();
                return (hdr, msg, 0, apiErr, new InvalidOperationException(apiErr.Description));
            }

            incr = parsedIncr;
            if (incr != null && !sourced)
            {
                if (!allowMsgCounter)
                {
                    var apiErr = JsApiErrors.NewJSMessageIncrDisabledError();
                    return (hdr, msg, 0, apiErr, new InvalidOperationException(apiErr.Description));
                }

                if (msg.Length > 0)
                {
                    var apiErr = JsApiErrors.NewJSMessageIncrPayloadError();
                    return (hdr, msg, 0, apiErr, new InvalidOperationException(apiErr.Description));
                }

                var hasInvalidCounterHeader =
                    NatsStream.GetRollup(hdr) != string.Empty ||
                    NatsStream.GetExpectedStream(hdr) != string.Empty ||
                    NatsStream.GetExpectedLastMsgId(hdr) != string.Empty ||
                    NatsStream.GetExpectedLastSeqPerSubjectForSubject(hdr) != string.Empty ||
                    NatsStream.GetExpectedLastSeq(hdr).Exists ||
                    NatsStream.GetExpectedLastSeqPerSubject(hdr).Exists;

                if (hasInvalidCounterHeader)
                {
                    var apiErr = JsApiErrors.NewJSMessageIncrInvalidError();
                    return (hdr, msg, 0, apiErr, new InvalidOperationException(apiErr.Description));
                }
            }

            var expectedStream = NatsStream.GetExpectedStream(hdr);
            if (!string.IsNullOrEmpty(expectedStream) && !string.Equals(expectedStream, name, StringComparison.Ordinal))
                return (hdr, msg, 0, JsApiErrors.NewJSStreamNotMatchError(), new InvalidOperationException("stream mismatch"));

            var ttlValue = NatsMessageHeaders.GetHeader(NatsHeaderConstants.JsMessageTtl, hdr) is { Length: > 0 } ttlHdr
                ? Encoding.ASCII.GetString(ttlHdr)
                : string.Empty;
            var (ttl, ttlErr) = string.IsNullOrEmpty(ttlValue)
                ? (0L, (Exception?)null)
                : NatsStream.ParseMessageTTL(ttlValue);
            if (!sourced && (ttl != 0 || ttlErr != null))
            {
                if (!allowTtl)
                    return (hdr, msg, 0, JsApiErrors.NewJSMessageTTLDisabledError(), new InvalidOperationException("message ttl disabled"));

                if (ttlErr != null)
                    return (hdr, msg, 0, JsApiErrors.NewJSMessageTTLInvalidError(), ttlErr);
            }

            var msgId = NatsStream.GetMsgId(hdr);
            if (!string.IsNullOrEmpty(msgId))
            {
                diff.MsgIds ??= new Dictionary<string, object?>(StringComparer.Ordinal);
                if (diff.MsgIds.ContainsKey(msgId))
                {
                    var err = new InvalidOperationException("duplicate message id in staged batch");
                    return (hdr, msg, 0, JsApiErrors.NewJSAtomicPublishContainsDuplicateMessageError(), err);
                }

                if (context.MsgIds.TryGetValue(msgId, out var dedupe))
                {
                    var err = new InvalidOperationException("duplicate message id");
                    if (dedupe.Seq > 0)
                        return (hdr, msg, dedupe.Seq, JsApiErrors.NewJSAtomicPublishContainsDuplicateMessageError(), err);

                    return (hdr, msg, 0, JsApiErrors.NewJSStreamDuplicateMessageConflictError(), err);
                }

                diff.MsgIds[msgId] = null;
            }
        }

        if (incr is null && allowMsgCounter)
        {
            var apiErr = JsApiErrors.NewJSMessageIncrMissingError();
            return (hdr, msg, 0, apiErr, new InvalidOperationException(apiErr.Description));
        }

        if (incr != null && allowMsgCounter)
        {
            var initial = BigInteger.Zero;
            Dictionary<string, Dictionary<string, string>>? sources = null;
            MsgCounterRunningTotal? counter = null;

            diff.Counter ??= new Dictionary<string, MsgCounterRunningTotal>(StringComparer.Ordinal);
            if (diff.Counter.TryGetValue(subject, out var stagedCounter))
            {
                initial = stagedCounter.Total;
                sources = stagedCounter.Sources;
                counter = stagedCounter;
            }
            else if (context.ClusteredCounterTotal.TryGetValue(subject, out var committedCounter))
            {
                initial = committedCounter.Total;
                sources = committedCounter.Sources;
                counter = new MsgCounterRunningTotal { Ops = committedCounter.Ops };
            }
            else
            {
                var last = context.Store.LoadLastMsg(subject, new StoreMsg());
                if (last != null)
                {
                    try
                    {
                        var current = JsonSerializer.Deserialize<CounterValue>(last.Msg);
                        if (current?.Value is null || !BigInteger.TryParse(current.Value, out initial))
                            throw new InvalidOperationException("invalid counter payload");

                        var sourceHdr = NatsMessageHeaders.SliceHeader(JsMessageCounterSources, last.Hdr);
                        if (sourceHdr is { Length: > 0 })
                            sources = JsonSerializer.Deserialize<Dictionary<string, Dictionary<string, string>>>(sourceHdr.Value.ToArray());
                    }
                    catch (Exception ex)
                    {
                        return (hdr, msg, 0, JsApiErrors.NewJSMessageCounterBrokenError(), ex);
                    }
                }
            }

            var sourceHeader = NatsMessageHeaders.SliceHeader(JsStreamSource, hdr);
            if (sourceHeader is { Length: > 0 })
            {
                try
                {
                    var fields = Encoding.ASCII.GetString(sourceHeader.Value.Span).Split(' ', StringSplitOptions.RemoveEmptyEntries);
                    var originStream = fields.Length > 0 ? fields[0] : string.Empty;
                    var originSubject = fields.Length >= 5 ? fields[4] : subject;

                    var current = JsonSerializer.Deserialize<CounterValue>(msg);
                    if (current?.Value is null || !BigInteger.TryParse(current.Value, out var sourcedValue))
                        throw new InvalidOperationException("invalid sourced counter payload");

                    sources ??= new Dictionary<string, Dictionary<string, string>>(StringComparer.Ordinal);
                    if (!sources.TryGetValue(originStream, out var bySubject))
                    {
                        bySubject = new Dictionary<string, string>(StringComparer.Ordinal);
                        sources[originStream] = bySubject;
                    }

                    var previousRaw = bySubject.GetValueOrDefault(originSubject);
                    _ = BigInteger.TryParse(previousRaw, out var previousValue);
                    bySubject[originSubject] = sourcedValue.ToString();
                    incr = sourcedValue - previousValue;
                    hdr = NatsMessageHeaders.SetHeader(NatsHeaderConstants.JsMessageIncr, incr.ToString(), hdr);
                }
                catch (Exception ex)
                {
                    return (hdr, msg, 0, JsApiErrors.NewJSMessageCounterBrokenError(), ex);
                }
            }

            initial += incr.Value;
            msg = Encoding.ASCII.GetBytes($"{{\"val\":\"{initial}\"}}");
            if (sources is { Count: > 0 })
                hdr = NatsMessageHeaders.SetHeader(JsMessageCounterSources, JsonSerializer.Serialize(sources), hdr);

            var maxSize = context.MaxPayload;
            if (maxMsgSize >= 0 && maxMsgSize < maxSize)
                maxSize = maxMsgSize;
            if (hdr.Length > maxSize || msg.Length > maxSize - hdr.Length)
                return (hdr, msg, 0, JsApiErrors.NewJSStreamMessageExceedsMaximumError(), ServerErrors.ErrMaxPayload);

            counter ??= new MsgCounterRunningTotal();
            counter.Total = initial;
            counter.Sources = sources;
            counter.Ops++;
            diff.Counter[subject] = counter;
        }

        if (hdr.Length > 0)
        {
            var (lastSeq, hasLastSeq) = NatsStream.GetExpectedLastSeq(hdr);
            var actualLastSeq = context.ClSeq - context.Clfs;
            if (hasLastSeq && lastSeq != actualLastSeq)
                return (hdr, msg, 0, JsApiErrors.NewJSStreamWrongLastSequenceError(actualLastSeq), new InvalidOperationException($"last sequence mismatch: {lastSeq} vs {actualLastSeq}"));
            if (hasLastSeq && diff.Inflight is { Count: > 0 })
                return (hdr, msg, 0, JsApiErrors.NewJSStreamWrongLastSequenceConstantError(), new InvalidOperationException("last sequence mismatch"));

            var (expectedPerSubject, hasExpectedPerSubject) = NatsStream.GetExpectedLastSeqPerSubject(hdr);
            if (hasExpectedPerSubject)
            {
                var seqSubject = subject;
                var overrideSubject = NatsStream.GetExpectedLastSeqPerSubjectForSubject(hdr);
                if (!string.IsNullOrEmpty(overrideSubject))
                    seqSubject = overrideSubject;

                if ((diff.Inflight?.ContainsKey(seqSubject) ?? false) ||
                    context.ExpectedPerSubjectInProcess.Contains(seqSubject) ||
                    context.Inflight.ContainsKey(seqSubject))
                    return (hdr, msg, 0, JsApiErrors.NewJSStreamWrongLastSequenceConstantError(), new InvalidOperationException("last sequence by subject mismatch"));

                diff.ExpectedPerSubject ??= new Dictionary<string, BatchExpectedPerSubject>(StringComparer.Ordinal);
                if (diff.ExpectedPerSubject.TryGetValue(seqSubject, out var existingExpected))
                {
                    if (existingExpected.SSeq != expectedPerSubject)
                        return (hdr, msg, 0, JsApiErrors.NewJSStreamWrongLastSequenceError(existingExpected.SSeq), new InvalidOperationException($"last sequence by subject mismatch: {expectedPerSubject} vs {existingExpected.SSeq}"));
                    existingExpected.ClSeq = context.ClSeq;
                }
                else
                {
                    var stored = context.Store.LoadLastMsg(seqSubject, new StoreMsg());
                    var foundSeq = stored?.Seq ?? 0;
                    if (foundSeq != expectedPerSubject)
                        return (hdr, msg, 0, JsApiErrors.NewJSStreamWrongLastSequenceError(foundSeq), new InvalidOperationException($"last sequence by subject mismatch: {expectedPerSubject} vs {foundSeq}"));

                    diff.ExpectedPerSubject[seqSubject] = new BatchExpectedPerSubject
                    {
                        SSeq = foundSeq,
                        ClSeq = context.ClSeq,
                    };
                }
            }
            else if (!string.IsNullOrEmpty(NatsStream.GetExpectedLastSeqPerSubjectForSubject(hdr)))
            {
                var apiErr = JsApiErrors.NewJSStreamExpectedLastSeqPerSubjectInvalidError();
                return (hdr, msg, 0, apiErr, new InvalidOperationException(apiErr.Description));
            }

            var (schedule, hasValidSchedulePattern) = NatsStream.GetMessageSchedule(hdr);
            if (!hasValidSchedulePattern)
            {
                var apiErr = allowMsgSchedules
                    ? JsApiErrors.NewJSMessageSchedulesPatternInvalidError()
                    : JsApiErrors.NewJSMessageSchedulesDisabledError();
                return (hdr, msg, 0, apiErr, new InvalidOperationException(apiErr.Description));
            }
            if (schedule != default)
            {
                if (!allowMsgSchedules)
                {
                    var apiErr = JsApiErrors.NewJSMessageSchedulesDisabledError();
                    return (hdr, msg, 0, apiErr, new InvalidOperationException(apiErr.Description));
                }

                var (scheduleTtl, scheduleTtlOk) = NatsStream.GetMessageScheduleTTL(hdr);
                if (!scheduleTtlOk)
                {
                    var apiErr = JsApiErrors.NewJSMessageSchedulesTTLInvalidError();
                    return (hdr, msg, 0, apiErr, new InvalidOperationException(apiErr.Description));
                }
                if (!string.IsNullOrEmpty(scheduleTtl) && !allowTtl)
                    return (hdr, msg, 0, JsApiErrors.NewJSMessageTTLDisabledError(), new InvalidOperationException("message ttl disabled"));

                var scheduleTarget = NatsStream.GetMessageScheduleTarget(hdr);
                if (string.IsNullOrEmpty(scheduleTarget) ||
                    !SubscriptionIndex.IsValidPublishSubject(scheduleTarget) ||
                    SubscriptionIndex.SubjectsCollide(scheduleTarget, subject))
                {
                    var apiErr = JsApiErrors.NewJSMessageSchedulesTargetInvalidError();
                    return (hdr, msg, 0, apiErr, new InvalidOperationException(apiErr.Description));
                }

                var scheduleSource = NatsStream.GetMessageScheduleSource(hdr);
                if (!string.IsNullOrEmpty(scheduleSource) &&
                    (string.Equals(scheduleSource, scheduleTarget, StringComparison.Ordinal) ||
                     string.Equals(scheduleSource, subject, StringComparison.Ordinal) ||
                     !SubscriptionIndex.IsValidPublishSubject(scheduleSource)))
                {
                    var apiErr = JsApiErrors.NewJSMessageSchedulesSourceInvalidError();
                    return (hdr, msg, 0, apiErr, new InvalidOperationException(apiErr.Description));
                }

                if (context.StreamSubjects.Length > 0 && !context.StreamSubjects.Any(s => SubscriptionIndex.SubjectsCollide(s, scheduleTarget)))
                {
                    var apiErr = JsApiErrors.NewJSMessageSchedulesTargetInvalidError();
                    return (hdr, msg, 0, apiErr, new InvalidOperationException(apiErr.Description));
                }

                var rollup = NatsStream.GetRollup(hdr);
                if (string.IsNullOrEmpty(rollup))
                {
                    hdr = NatsMessageHeaders.GenHeader(hdr, NatsHeaderConstants.JsMsgRollup, NatsHeaderConstants.JsMsgRollupSubject);
                }
                else if (!string.Equals(rollup, NatsHeaderConstants.JsMsgRollupSubject, StringComparison.Ordinal))
                {
                    var apiErr = JsApiErrors.NewJSMessageSchedulesRollupInvalidError();
                    return (hdr, msg, 0, apiErr, new InvalidOperationException(apiErr.Description));
                }
            }

            var headerRollup = NatsStream.GetRollup(hdr);
            if (!string.IsNullOrEmpty(headerRollup))
            {
                if (!allowRollup || denyPurge)
                {
                    var err = new InvalidOperationException("rollup not permitted");
                    return (hdr, msg, 0, JsApiErrors.NewJSStreamRollupFailedError(err), err);
                }

                if (string.Equals(headerRollup, NatsHeaderConstants.JsMsgRollupSubject, StringComparison.Ordinal))
                {
                    if (diff.Inflight?.ContainsKey(subject) ?? false)
                    {
                        var err = new InvalidOperationException("batch rollup sub invalid");
                        return (hdr, msg, 0, JsApiErrors.NewJSStreamRollupFailedError(err), err);
                    }
                }
                else if (string.Equals(headerRollup, NatsHeaderConstants.JsMsgRollupAll, StringComparison.Ordinal))
                {
                    if (diff.Inflight is { Count: > 0 })
                    {
                        var err = new InvalidOperationException("batch rollup all invalid");
                        return (hdr, msg, 0, JsApiErrors.NewJSStreamRollupFailedError(err), err);
                    }
                }
                else
                {
                    var err = new InvalidOperationException($"rollup value invalid: {headerRollup}");
                    return (hdr, msg, 0, JsApiErrors.NewJSStreamRollupFailedError(err), err);
                }
            }
        }

        diff.Inflight ??= new Dictionary<string, InflightSubjectRunningTotal>(StringComparer.Ordinal);
        var msgSize = context.Store.Type() == StorageType.FileStorage
            ? JetStreamFileStore.FileStoreMsgSizeRaw(subject.Length, hdr.Length, msg.Length)
            : JetStreamMemStore.MemStoreMsgSizeRaw(subject.Length, hdr.Length, msg.Length);

        if (diff.Inflight.TryGetValue(subject, out var inflight))
        {
            inflight.Bytes += msgSize;
            inflight.Ops++;
        }
        else
        {
            inflight = new InflightSubjectRunningTotal { Bytes = msgSize, Ops = 1 };
            diff.Inflight[subject] = inflight;
        }

        if (discard == DiscardPolicy.DiscardNew)
        {
            if (maxMsgs > 0 || maxBytes > 0)
            {
                var state = new StreamState();
                context.Store.FastState(state);
                var totalMsgs = state.Msgs;
                var totalBytes = state.Bytes;
                foreach (var inflightState in context.Inflight.Values)
                {
                    totalMsgs += inflightState.Ops;
                    totalBytes += inflightState.Bytes;
                }
                foreach (var inflightState in diff.Inflight.Values)
                {
                    totalMsgs += inflightState.Ops;
                    totalBytes += inflightState.Bytes;
                }

                Exception? thresholdErr = null;
                if (maxMsgs > 0 && totalMsgs > (ulong)maxMsgs)
                    thresholdErr = StoreErrors.ErrMaxMsgs;
                else if (maxBytes > 0 && totalBytes > (ulong)maxBytes)
                    thresholdErr = StoreErrors.ErrMaxBytes;

                if (thresholdErr != null)
                    return (hdr, msg, 0, JsApiErrors.NewJSStreamStoreFailedError(thresholdErr, JsApiErrors.Unless(thresholdErr)), thresholdErr);
            }

            if (discardNewPer && maxMsgsPer > 0)
            {
                var bySubject = context.Store.SubjectsTotals(subject);
                var totalForSubject = bySubject.GetValueOrDefault(subject) + inflight.Ops;
                if (context.Inflight.TryGetValue(subject, out var streamInflight))
                    totalForSubject += streamInflight.Ops;
                if (totalForSubject > (ulong)maxMsgsPer)
                {
                    var err = StoreErrors.ErrMaxMsgsPerSubject;
                    return (hdr, msg, 0, JsApiErrors.NewJSStreamStoreFailedError(err, JsApiErrors.Unless(err)), err);
                }
            }
        }

        return (hdr, msg, 0, null, null);
    }
}

internal interface IBatchTimer
{
    bool Stop();
}

internal sealed class SystemBatchTimer : IBatchTimer, IDisposable
{
    private readonly Timer _timer;
    private int _stopped;

    public SystemBatchTimer(TimeSpan timeout, Action callback)
    {
        ArgumentNullException.ThrowIfNull(callback);

        _timer = new Timer(_ =>
        {
            if (Interlocked.Exchange(ref _stopped, 1) == 0)
                callback();
        }, null, timeout, Timeout.InfiniteTimeSpan);
    }

    public bool Stop()
    {
        var firstStop = Interlocked.Exchange(ref _stopped, 1) == 0;
        _timer.Change(Timeout.InfiniteTimeSpan, Timeout.InfiniteTimeSpan);
        return firstStop;
    }

    public void Dispose() => _timer.Dispose();
}

/// <summary>
/// Tracks in-progress atomic publish batch groups for a stream.
/// Mirrors the <c>batching</c> struct in server/jetstream_batching.go.
/// </summary>
internal sealed class Batching
{
    private readonly Lock _mu = new();
    private readonly Dictionary<string, BatchGroup> _group = new(StringComparer.Ordinal);

    public Lock Mu => _mu;
    public Dictionary<string, BatchGroup> Group => _group;

    internal BatchGroup NewBatchGroup(string storeRootDir, StreamConfig streamConfig, string batchId, Action<string, string>? abandonedAdvisory = null, TimeSpan? maxBatchTimeout = null)
    {
        var store = JetStreamBatching.NewBatchStore(storeRootDir, streamConfig, batchId);
        var group = new BatchGroup { Store = store };

        var timeout = maxBatchTimeout.GetValueOrDefault(JetStreamBatching.StreamDefaultMaxBatchTimeout);
        group.BatchTimer = new SystemBatchTimer(timeout, () =>
        {
            group.Cleanup(batchId, this);
            abandonedAdvisory?.Invoke(batchId, JetStreamBatching.BatchTimeout);
        });

        _group[batchId] = group;
        JetStreamBatching.IncrementGlobalInflightBatches();
        return group;
    }
}

/// <summary>
/// A single in-progress atomic batch: its temporary store and cleanup timer.
/// Mirrors <c>batchGroup</c> in server/jetstream_batching.go.
/// </summary>
internal sealed class BatchGroup
{
    /// <summary>Last proposed stream sequence for this batch.</summary>
    public ulong LastSeq { get; set; }

    /// <summary>Temporary backing store for the batch's messages.</summary>
    public IStreamStore? Store { get; set; }

    /// <summary>Timer that abandons the batch after the configured timeout.</summary>
    public IBatchTimer? BatchTimer { get; set; }

    /// <summary>
    /// Stops the cleanup timer and flushes pending writes so the batch is
    /// ready to be committed.
    /// Mirrors <c>batchGroup.readyForCommit</c>.
    /// </summary>
    public bool ReadyForCommit()
    {
        if (BatchTimer?.Stop() != true)
            return false;

        Store?.FlushAllPending();
        return true;
    }

    /// <summary>
    /// Deletes underlying resources associated with the batch and unregisters it from the stream's batches.
    /// Mirrors <c>batchGroup.cleanup</c>.
    /// </summary>
    public void Cleanup(string batchId, Batching batches)
    {
        ArgumentNullException.ThrowIfNull(batches);
        lock (batches.Mu)
        {
            CleanupLocked(batchId, batches);
        }
    }

    /// <summary>
    /// Deletes underlying resources associated with the batch and unregisters it from the stream's batches.
    /// Mirrors <c>batchGroup.cleanupLocked</c>.
    /// </summary>
    public void CleanupLocked(string batchId, Batching batches)
    {
        ArgumentNullException.ThrowIfNull(batches);

        JetStreamBatching.DecrementGlobalInflightBatches();
        _ = BatchTimer?.Stop();
        Store?.Delete(true);
        _ = batches.Group.Remove(batchId);
    }

    /// <summary>
    /// Stops underlying resources associated with the batch.
    /// Mirrors <c>batchGroup.stopLocked</c>.
    /// </summary>
    public void StopLocked()
    {
        JetStreamBatching.DecrementGlobalInflightBatches();
        _ = BatchTimer?.Stop();
        Store?.Stop();
    }
}

/// <summary>
/// Stages consistency-check data for a single atomic batch before it is committed.
/// Mirrors <c>batchStagedDiff</c> in server/jetstream_batching.go.
/// </summary>
internal sealed class BatchStagedDiff
{
    /// <summary>Message IDs seen in this batch, for duplicate detection.</summary>
    public Dictionary<string, object?>? MsgIds { get; set; }

    /// <summary>Running counter totals, keyed by subject.</summary>
    public Dictionary<string, MsgCounterRunningTotal>? Counter { get; set; }

    /// <summary>Inflight subject byte/op totals for DiscardNew checks.</summary>
    public Dictionary<string, InflightSubjectRunningTotal>? Inflight { get; set; }

    /// <summary>Expected-last-seq-per-subject checks staged in this batch.</summary>
    public Dictionary<string, BatchExpectedPerSubject>? ExpectedPerSubject { get; set; }
}

/// <summary>
/// Cached expected-last-sequence-per-subject result for a single subject within a batch.
/// Mirrors <c>batchExpectedPerSubject</c> in server/jetstream_batching.go.
/// </summary>
internal sealed class BatchExpectedPerSubject
{
    /// <summary>Stream sequence of the last message on this subject at proposal time.</summary>
    public ulong SSeq { get; set; }

    /// <summary>Clustered proposal sequence at which this check was computed.</summary>
    public ulong ClSeq { get; set; }
}

/// <summary>
/// Tracks the in-progress application of a committed batch on the Raft apply path.
/// Mirrors <c>batchApply</c> in server/jetstream_batching.go.
/// </summary>
internal sealed class BatchApply
{
    private readonly Lock _mu = new();

    /// <summary>ID of the current batch.</summary>
    public string Id { get; set; } = string.Empty;

    /// <summary>Number of entries expected in the batch (for consistency checks).</summary>
    public ulong Count { get; set; }

    /// <summary>Raft committed entries that make up this batch.</summary>
    public List<ICommittedEntry>? Entries { get; set; }

    /// <summary>Index within an entry indicating the first message of the batch.</summary>
    public int EntryStart { get; set; }

    /// <summary>Applied value before the entry containing the first batch message.</summary>
    public ulong MaxApplied { get; set; }

    public Lock Mu => _mu;

    /// <summary>
    /// Clears in-memory apply-batch state.
    /// Mirrors <c>batchApply.clearBatchStateLocked</c>.
    /// </summary>
    public void ClearBatchStateLocked()
    {
        Id = string.Empty;
        Count = 0;
        Entries = null;
        EntryStart = 0;
        MaxApplied = 0;
    }

    public void RejectBatchStateLocked(BatchApplyContext context)
    {
        ArgumentNullException.ThrowIfNull(context);

        context.Clfs += Count;
        if (Entries is { Count: > 0 })
        {
            foreach (var entry in Entries)
                entry.ReturnToPool();
        }

        ClearBatchStateLocked();
    }

    public void RejectBatchState(BatchApplyContext context)
    {
        lock (_mu)
        {
            RejectBatchStateLocked(context);
        }
    }
}

internal interface ICommittedEntry
{
    void ReturnToPool();
}

internal sealed class BatchApplyContext
{
    public ulong Clfs { get; set; }
}

internal sealed class MsgCounterRunningTotal
{
    public ulong Ops { get; set; }
    public BigInteger Total { get; set; }
    public Dictionary<string, Dictionary<string, string>>? Sources { get; set; }
}

internal sealed class InflightSubjectRunningTotal
{
    public ulong Bytes { get; set; }
    public ulong Ops { get; set; }
}

internal sealed class BatchRuntimeState
{
    public Dictionary<string, NatsStream.DedupeEntry> MsgIds { get; } = new(StringComparer.Ordinal);
    public Dictionary<string, MsgCounterRunningTotal> ClusteredCounterTotal { get; } = new(StringComparer.Ordinal);
    public Dictionary<string, InflightSubjectRunningTotal> Inflight { get; } = new(StringComparer.Ordinal);
    public Dictionary<ulong, string> ExpectedPerSubjectSequence { get; } = new();
    public HashSet<string> ExpectedPerSubjectInProcess { get; } = new(StringComparer.Ordinal);
}

internal sealed class BatchHeaderCheckContext
{
    public string AccountName { get; set; } = string.Empty;
    public required IStreamStore Store { get; set; }
    public Dictionary<string, NatsStream.DedupeEntry> MsgIds { get; } = new(StringComparer.Ordinal);
    public Dictionary<string, MsgCounterRunningTotal> ClusteredCounterTotal { get; } = new(StringComparer.Ordinal);
    public Dictionary<string, InflightSubjectRunningTotal> Inflight { get; } = new(StringComparer.Ordinal);
    public HashSet<string> ExpectedPerSubjectInProcess { get; } = new(StringComparer.Ordinal);
    public string[] StreamSubjects { get; set; } = Array.Empty<string>();
    public ulong ClSeq { get; set; }
    public ulong Clfs { get; set; }
    public int MaxPayload { get; set; } = int.MaxValue;
}

internal sealed class CounterValue
{
    public string Value { get; set; } = "0";
}
