using System.Collections.Concurrent;

namespace NATS.Server.JetStream.Publish;

// Go: server/jetstream_batching.go — streamBatches / batchGroup
// Handles the staging and commit protocol for atomic batch publishing.
// Messages within a batch are buffered until Nats-Batch-Commit is received,
// then committed atomically to the store.

/// <summary>
/// Error codes for atomic batch publish operations.
/// Go reference: server/jetstream_errors_generated.go
/// </summary>
public static class AtomicBatchPublishErrorCodes
{
    /// <summary>JSAtomicPublishDisabledErr (10174) — stream has AllowAtomicPublish=false.</summary>
    public const int Disabled = 10174;

    /// <summary>JSAtomicPublishMissingSeqErr (10175) — Nats-Batch-Id present but Nats-Batch-Sequence missing.</summary>
    public const int MissingSeq = 10175;

    /// <summary>JSAtomicPublishIncompleteBatchErr (10176) — sequence gap, timeout, or over-limit batch.</summary>
    public const int IncompleteBatch = 10176;

    /// <summary>JSAtomicPublishUnsupportedHeaderBatchErr (10177) — unsupported header in batch message.</summary>
    public const int UnsupportedHeader = 10177;

    /// <summary>JSAtomicPublishInvalidBatchIDErr (10179) — batch ID too long (&gt;64 chars).</summary>
    public const int InvalidBatchId = 10179;

    /// <summary>JSAtomicPublishTooLargeBatchErrF (10199) — batch size exceeds maximum.</summary>
    public const int TooLargeBatch = 10199;

    /// <summary>JSAtomicPublishInvalidBatchCommitErr (10200) — Nats-Batch-Commit header has invalid value.</summary>
    public const int InvalidCommit = 10200;

    /// <summary>JSAtomicPublishContainsDuplicateMessageErr (10201) — batch contains a duplicate msg ID.</summary>
    public const int ContainsDuplicate = 10201;

    /// <summary>JSMirrorWithAtomicPublishErr (10198) — mirror stream cannot have AllowAtomicPublish.</summary>
    public const int MirrorWithAtomicPublish = 10198;
}

/// <summary>
/// A single staged message within an in-flight atomic batch.
/// </summary>
public sealed class StagedBatchMessage
{
    public required string Subject { get; init; }
    public required ReadOnlyMemory<byte> Payload { get; init; }
    public string? MsgId { get; init; }
    public ulong ExpectedLastSeq { get; init; }
    public ulong ExpectedLastSubjectSeq { get; init; }
    public string? ExpectedLastSubjectSeqSubject { get; init; }
}

/// <summary>
/// State for one in-flight atomic batch (identified by Nats-Batch-Id).
/// Go reference: server/jetstream_batching.go batchGroup struct.
/// </summary>
internal sealed class InFlightBatch
{
    private readonly List<StagedBatchMessage> _messages = [];
    private readonly HashSet<string> _stagedMsgIds = new(StringComparer.Ordinal);

    public DateTimeOffset CreatedAt { get; } = DateTimeOffset.UtcNow;
    public int Count => _messages.Count;
    public IReadOnlyList<StagedBatchMessage> Messages => _messages;

    public void Add(StagedBatchMessage msg)
    {
        _messages.Add(msg);
        if (msg.MsgId != null)
            _stagedMsgIds.Add(msg.MsgId);
    }

    public bool ContainsMsgId(string msgId) => _stagedMsgIds.Contains(msgId);
}

/// <summary>
/// Engine that manages all in-flight atomic batches for a single stream.
/// Go reference: server/jetstream_batching.go streamBatches struct.
///
/// The protocol:
///   1. Publisher sends messages with Nats-Batch-Id + Nats-Batch-Sequence headers.
///      Messages with seq &gt; 1 are staged; they are not committed to the store yet.
///   2. The final message has Nats-Batch-Commit: 1 (or "eob"). This triggers atomic
///      commit of all staged messages.
///   3. On error (gap, timeout, disable, delete, leader-change), the batch is rolled back.
/// </summary>
public sealed class AtomicBatchPublishEngine
{
    // Go: streamDefaultMaxBatchInflightPerStream = 50
    public const int DefaultMaxInflightPerStream = 50;
    // Go: streamDefaultMaxBatchSize = 1000
    public const int DefaultMaxBatchSize = 1000;
    // Go: streamDefaultMaxBatchTimeout = 10s
    public static readonly TimeSpan DefaultBatchTimeout = TimeSpan.FromSeconds(10);
    // Go: max batch ID length = 64
    public const int MaxBatchIdLength = 64;

    private readonly ConcurrentDictionary<string, InFlightBatch> _batches = new(StringComparer.Ordinal);
    private readonly int _maxInflightPerStream;
    private readonly int _maxBatchSize;
    private readonly TimeSpan _batchTimeout;

    public AtomicBatchPublishEngine(
        int maxInflightPerStream = DefaultMaxInflightPerStream,
        int maxBatchSize = DefaultMaxBatchSize,
        TimeSpan? batchTimeout = null)
    {
        _maxInflightPerStream = maxInflightPerStream;
        _maxBatchSize = maxBatchSize;
        _batchTimeout = batchTimeout ?? DefaultBatchTimeout;
    }

    /// <summary>
    /// Returns the number of in-flight batches.
    /// </summary>
    public int InflightCount => _batches.Count;

    /// <summary>
    /// Validates and stages/commits a batch message.
    /// Returns a result indicating: stage (empty ack), commit (full ack), or error.
    /// </summary>
    public AtomicBatchResult Process(
        BatchPublishRequest req,
        PublishPreconditions preconditions,
        int streamDuplicateWindowMs,
        Func<StagedBatchMessage, PubAck?> commitSingle)
    {
        // Validate batch ID length (max 64 chars).
        if (req.BatchId.Length > MaxBatchIdLength)
            return AtomicBatchResult.Error(AtomicBatchPublishErrorCodes.InvalidBatchId,
                "atomic publish batch ID is invalid");

        // Validate commit value: only "1" and "eob" are valid; empty/"0"/"anything-else" is invalid.
        if (req.IsCommit && req.CommitValue is not ("1" or "eob"))
            return AtomicBatchResult.Error(AtomicBatchPublishErrorCodes.InvalidCommit,
                "atomic publish batch commit is invalid");

        EvictExpiredBatches();

        // Check inflight limit.
        if (!_batches.ContainsKey(req.BatchId) && _batches.Count >= _maxInflightPerStream)
            return AtomicBatchResult.Error(AtomicBatchPublishErrorCodes.IncompleteBatch,
                "atomic publish batch is incomplete");

        // Sequence 1 starts a new batch.
        if (req.BatchSeq == 1)
        {
            // Start new batch or replace any previous one with the same ID.
            var newBatch = new InFlightBatch();

            // Check for duplicate msg ID in existing store (preconditions).
            if (req.MsgId != null && preconditions.IsDuplicate(req.MsgId, streamDuplicateWindowMs, out _))
                return AtomicBatchResult.Error(AtomicBatchPublishErrorCodes.ContainsDuplicate,
                    "atomic publish batch contains duplicate message id");

            // Stage this message.
            var staged = new StagedBatchMessage
            {
                Subject = req.Subject,
                Payload = req.Payload,
                MsgId = req.MsgId,
                ExpectedLastSeq = req.ExpectedLastSeq,
                ExpectedLastSubjectSeq = req.ExpectedLastSubjectSeq,
                ExpectedLastSubjectSeqSubject = req.ExpectedLastSubjectSeqSubject,
            };
            newBatch.Add(staged);

            // Single-message batch that commits immediately.
            if (req.IsCommit)
            {
                // Commit just this one message.
                var ack = commitSingle(staged);
                if (ack == null || ack.ErrorCode != null)
                {
                    return AtomicBatchResult.Error(
                        ack?.ErrorCode ?? AtomicBatchPublishErrorCodes.IncompleteBatch,
                        "atomic publish batch is incomplete");
                }

                if (req.MsgId != null)
                {
                    preconditions.Record(req.MsgId, ack.Seq);
                    preconditions.TrimOlderThan(streamDuplicateWindowMs);
                }

                return AtomicBatchResult.Committed(new PubAck
                {
                    Stream = ack.Stream,
                    Seq = ack.Seq,
                    BatchId = req.BatchId,
                    BatchSize = 1,
                });
            }

            // Check batch size limit for single-message staging.
            if (1 > _maxBatchSize)
                return AtomicBatchResult.Error(AtomicBatchPublishErrorCodes.TooLargeBatch,
                    $"atomic publish batch is too large: {_maxBatchSize}");

            _batches[req.BatchId] = newBatch;
            return AtomicBatchResult.Staged();
        }

        // Non-first sequence: find the in-flight batch.
        if (!_batches.TryGetValue(req.BatchId, out var batch))
            return AtomicBatchResult.Error(AtomicBatchPublishErrorCodes.IncompleteBatch,
                "atomic publish batch is incomplete");

        // Sequence must be consecutive (no gaps).
        if (req.BatchSeq != (ulong)batch.Count + 1)
        {
            _batches.TryRemove(req.BatchId, out _);
            return AtomicBatchResult.Error(AtomicBatchPublishErrorCodes.IncompleteBatch,
                "atomic publish batch is incomplete");
        }

        // Check for duplicate msg ID in store or already staged in this batch.
        if (req.MsgId != null)
        {
            if (preconditions.IsDuplicate(req.MsgId, streamDuplicateWindowMs, out _))
            {
                _batches.TryRemove(req.BatchId, out _);
                return AtomicBatchResult.Error(AtomicBatchPublishErrorCodes.ContainsDuplicate,
                    "atomic publish batch contains duplicate message id");
            }

            if (batch.ContainsMsgId(req.MsgId))
            {
                _batches.TryRemove(req.BatchId, out _);
                return AtomicBatchResult.Error(AtomicBatchPublishErrorCodes.ContainsDuplicate,
                    "atomic publish batch contains duplicate message id");
            }
        }

        var stagedMsg = new StagedBatchMessage
        {
            Subject = req.Subject,
            Payload = req.Payload,
            MsgId = req.MsgId,
            ExpectedLastSeq = req.ExpectedLastSeq,
            ExpectedLastSubjectSeq = req.ExpectedLastSubjectSeq,
            ExpectedLastSubjectSeqSubject = req.ExpectedLastSubjectSeqSubject,
        };
        batch.Add(stagedMsg);

        // Check batch size limit.
        if (batch.Count > _maxBatchSize)
        {
            _batches.TryRemove(req.BatchId, out _);
            return AtomicBatchResult.Error(AtomicBatchPublishErrorCodes.TooLargeBatch,
                $"atomic publish batch is too large: {_maxBatchSize}");
        }

        if (!req.IsCommit)
        {
            // Not committing yet — return empty ack (flow control).
            return AtomicBatchResult.Staged();
        }

        // Validate commit value.
        if (req.CommitValue is not ("1" or "eob"))
        {
            _batches.TryRemove(req.BatchId, out _);
            return AtomicBatchResult.Error(AtomicBatchPublishErrorCodes.InvalidCommit,
                "atomic publish batch commit is invalid");
        }

        // Determine effective batch for eob: exclude the commit-marker message itself.
        var effectiveBatch = req.CommitValue == "eob"
            ? batch.Messages.Take(batch.Messages.Count - 1).ToList()
            : batch.Messages;

        _batches.TryRemove(req.BatchId, out _);

        // Atomically commit all staged messages.
        ulong lastSeq = 0;
        foreach (var msg in effectiveBatch)
        {
            var ack = commitSingle(msg);
            if (ack == null || ack.ErrorCode != null)
                return AtomicBatchResult.Error(
                    ack?.ErrorCode ?? AtomicBatchPublishErrorCodes.IncompleteBatch,
                    "atomic publish batch is incomplete");

            if (msg.MsgId != null)
            {
                preconditions.Record(msg.MsgId, ack.Seq);
                preconditions.TrimOlderThan(streamDuplicateWindowMs);
            }

            lastSeq = ack.Seq;
        }

        return AtomicBatchResult.Committed(new PubAck
        {
            Stream = effectiveBatch.Count > 0 ? (commitSingle == null ? "" : effectiveBatch[^1].Subject) : "",
            Seq = lastSeq,
            BatchId = req.BatchId,
            BatchSize = effectiveBatch.Count,
        });
    }

    /// <summary>
    /// Removes all batches from the engine (called on stream disable/delete/leader-change).
    /// </summary>
    public void Clear() => _batches.Clear();

    /// <summary>
    /// Returns whether a batch with the given ID is currently in-flight.
    /// </summary>
    public bool HasBatch(string batchId) => _batches.ContainsKey(batchId);

    private void EvictExpiredBatches()
    {
        var cutoff = DateTimeOffset.UtcNow - _batchTimeout;
        foreach (var (key, batch) in _batches)
        {
            if (batch.CreatedAt < cutoff)
                _batches.TryRemove(key, out _);
        }
    }
}

/// <summary>
/// Request to process a batch publish message.
/// </summary>
public sealed class BatchPublishRequest
{
    public required string BatchId { get; init; }
    public required ulong BatchSeq { get; init; }
    public required string Subject { get; init; }
    public required ReadOnlyMemory<byte> Payload { get; init; }

    /// <summary>
    /// True if Nats-Batch-Commit header is present with a non-empty value.
    /// </summary>
    public bool IsCommit { get; init; }

    /// <summary>
    /// Raw value of the Nats-Batch-Commit header ("1", "eob", or other).
    /// Only meaningful when IsCommit = true.
    /// </summary>
    public string? CommitValue { get; init; }

    public string? MsgId { get; init; }
    public ulong ExpectedLastSeq { get; init; }
    public ulong ExpectedLastSubjectSeq { get; init; }
    public string? ExpectedLastSubjectSeqSubject { get; init; }
}

/// <summary>
/// Result of processing a batch publish message.
/// </summary>
public sealed class AtomicBatchResult
{
    public enum ResultKind { Staged, Committed, Error }

    public ResultKind Kind { get; private init; }
    public PubAck? CommitAck { get; private init; }
    public int ErrorCode { get; private init; }
    public string ErrorDescription { get; private init; } = string.Empty;

    public static AtomicBatchResult Staged() => new() { Kind = ResultKind.Staged };

    public static AtomicBatchResult Committed(PubAck ack) =>
        new() { Kind = ResultKind.Committed, CommitAck = ack };

    public static AtomicBatchResult Error(int code, string description) =>
        new() { Kind = ResultKind.Error, ErrorCode = code, ErrorDescription = description };
}
