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

using System.Text;

namespace ZB.MOM.NatsNet.Server;

/// <summary>
/// Wire-level NATS message header constants.
/// </summary>
public static class NatsHeaderConstants
{
    /// <summary>NATS header status line: <c>"NATS/1.0\r\n"</c>. Mirrors Go <c>hdrLine</c>.</summary>
    public const string HdrLine      = "NATS/1.0\r\n";

    /// <summary>Empty header block with blank line terminator. Mirrors Go <c>emptyHdrLine</c>.</summary>
    public const string EmptyHdrLine = "NATS/1.0\r\n\r\n";

    // JetStream expected-sequence headers (defined in server/stream.go, used by header utilities).
    public const string JsExpectedStream        = "Nats-Expected-Stream";
    public const string JsExpectedLastSeq       = "Nats-Expected-Last-Sequence";
    public const string JsExpectedLastSubjSeq   = "Nats-Expected-Last-Subject-Sequence";
    public const string JsExpectedLastSubjSeqSubj = "Nats-Expected-Last-Subject-Sequence-Subject";
    public const string JsExpectedLastMsgId     = "Nats-Expected-Last-Msg-Id";

    // Other commonly used headers.
    public const string JsMsgId   = "Nats-Msg-Id";
    public const string JsMsgRollup = "Nats-Rollup";
    public const string JsMsgSize = "Nats-Msg-Size";
    public const string JsResponseType = "Nats-Response-Type";
    public const string JsMessageTtl = "Nats-TTL";
    public const string JsMarkerReason = "Nats-Marker-Reason";
    public const string JsMessageIncr = "Nats-Incr";
    public const string JsBatchId = "Nats-Batch-Id";
    public const string JsBatchSeq = "Nats-Batch-Sequence";
    public const string JsBatchCommit = "Nats-Batch-Commit";

    // Scheduling headers.
    public const string JsSchedulePattern = "Nats-Schedule";
    public const string JsScheduleTtl = "Nats-Schedule-TTL";
    public const string JsScheduleTarget = "Nats-Schedule-Target";
    public const string JsScheduleSource = "Nats-Schedule-Source";
    public const string JsScheduler = "Nats-Scheduler";
    public const string JsScheduleNext = "Nats-Schedule-Next";
    public const string JsScheduleNextPurge = "purge";

    // Rollup values.
    public const string JsMsgRollupSubject = "sub";
    public const string JsMsgRollupAll = "all";

    // Marker reasons.
    public const string JsMarkerReasonMaxAge = "MaxAge";
    public const string JsMarkerReasonPurge = "Purge";
    public const string JsMarkerReasonRemove = "Remove";
}

/// <summary>
/// Low-level NATS message header manipulation utilities.
/// Mirrors the package-level functions in server/client.go:
/// <c>genHeader</c>, <c>removeHeaderIfPresent</c>, <c>removeHeaderIfPrefixPresent</c>,
/// <c>getHeader</c>, <c>sliceHeader</c>, <c>getHeaderKeyIndex</c>, <c>setHeader</c>.
/// </summary>
public static class NatsMessageHeaders
{
    private static readonly byte[] CrLfBytes = "\r\n"u8.ToArray();

    // -------------------------------------------------------------------------
    // genHeader (feature 506)
    // -------------------------------------------------------------------------

    /// <summary>
    /// Generates a header buffer by appending <c>key: value\r\n</c> to an existing header,
    /// or starting a fresh <c>NATS/1.0\r\n</c> block if <paramref name="hdr"/> is empty/null.
    /// Mirrors Go <c>genHeader</c>.
    /// </summary>
    /// <param name="hdr">Existing header bytes, or <c>null</c> to start fresh.</param>
    /// <param name="key">Header key.</param>
    /// <param name="value">Header value.</param>
    public static byte[] GenHeader(byte[]? hdr, string key, string value)
    {
        var sb = new StringBuilder();

        // Strip trailing CRLF from existing header to reopen for appending,
        // or start fresh with the header status line.
        const int LenCrLf = 2;
        if (hdr is { Length: > LenCrLf })
        {
            // Write all but the trailing "\r\n"
            sb.Append(Encoding.ASCII.GetString(hdr, 0, hdr.Length - LenCrLf));
        }
        else
        {
            sb.Append(NatsHeaderConstants.HdrLine);
        }

        // Append "key: value\r\n\r\n" (HTTP header format).
        sb.Append(key);
        sb.Append(": ");
        sb.Append(value);
        sb.Append("\r\n\r\n");

        return Encoding.ASCII.GetBytes(sb.ToString());
    }

    // -------------------------------------------------------------------------
    // removeHeaderIfPresent (feature 504)
    // -------------------------------------------------------------------------

    /// <summary>
    /// Removes the first occurrence of <paramref name="key"/> header from <paramref name="hdr"/>.
    /// Returns <c>null</c> if the result would be an empty header block.
    /// Mirrors Go <c>removeHeaderIfPresent</c>.
    /// </summary>
    public static byte[]? RemoveHeaderIfPresent(byte[] hdr, string key)
    {
        int start = GetHeaderKeyIndex(key, hdr);
        // Key must exist and be preceded by '\n' (not at position 0).
        if (start < 1 || hdr[start - 1] != '\n')
            return hdr;

        int index = start + key.Length;
        if (index >= hdr.Length || hdr[index] != ':')
            return hdr;

        // Find CRLF following this header line.
        int crlfIdx = IndexOfCrLf(hdr, start);
        if (crlfIdx < 0)
            return hdr;

        // Remove from 'start' through end of CRLF.
        int removeEnd = start + crlfIdx + 2; // +2 for "\r\n"
        var result = new byte[hdr.Length - (removeEnd - start)];
        Buffer.BlockCopy(hdr, 0,         result, 0,              start);
        Buffer.BlockCopy(hdr, removeEnd, result, start, hdr.Length - removeEnd);

        // If nothing meaningful remains, return null.
        if (result.Length <= NatsHeaderConstants.EmptyHdrLine.Length)
            return null;

        return result;
    }

    // -------------------------------------------------------------------------
    // removeHeaderIfPrefixPresent (feature 505)
    // -------------------------------------------------------------------------

    /// <summary>
    /// Removes all headers whose names start with <paramref name="prefix"/>.
    /// Returns <c>null</c> if the result would be an empty header block.
    /// Mirrors Go <c>removeHeaderIfPrefixPresent</c>.
    /// </summary>
    public static byte[]? RemoveHeaderIfPrefixPresent(byte[] hdr, string prefix)
    {
        var prefixBytes = Encoding.ASCII.GetBytes(prefix);
        var working = hdr.ToList(); // work on a list for easy splicing
        int index = 0;

        while (index < working.Count)
        {
            // Look for prefix starting at current index.
            int found = IndexOf(working, prefixBytes, index);
            if (found < 0)
                break;

            // Must be preceded by '\n'.
            if (found < 1 || working[found - 1] != '\n')
                break;

            // Find CRLF after this prefix's key:value line.
            int crlfIdx = IndexOfCrLf(working, found + prefix.Length);
            if (crlfIdx < 0)
                break;

            int removeEnd = found + prefix.Length + crlfIdx + 2;
            working.RemoveRange(found, removeEnd - found);

            // Don't advance index — there may be more headers at same position.
            if (working.Count <= NatsHeaderConstants.EmptyHdrLine.Length)
                return null;
        }

        return working.ToArray();
    }

    // -------------------------------------------------------------------------
    // getHeaderKeyIndex (feature 510)
    // -------------------------------------------------------------------------

    /// <summary>
    /// Returns the byte offset of <paramref name="key"/> in <paramref name="hdr"/>,
    /// or <c>-1</c> if not found.
    /// The key must be preceded by <c>\r\n</c> and followed by <c>:</c>.
    /// Mirrors Go <c>getHeaderKeyIndex</c>.
    /// </summary>
    public static int GetHeaderKeyIndex(string key, byte[] hdr)
    {
        if (hdr.Length == 0) return -1;

        var bkey   = Encoding.ASCII.GetBytes(key);
        int keyLen = bkey.Length;
        int hdrLen = hdr.Length;
        int offset = 0;

        while (true)
        {
            int index = IndexOf(hdr, bkey, offset);
            // Need index >= 2 (room for preceding \r\n) and enough space for trailing colon.
            if (index < 2) return -1;

            // Preceded by \r\n ?
            if (hdr[index - 1] != '\n' || hdr[index - 2] != '\r')
            {
                offset = index + keyLen;
                continue;
            }

            // Immediately followed by ':' ?
            if (index + keyLen >= hdrLen)
                return -1;

            if (hdr[index + keyLen] != ':')
            {
                offset = index + keyLen;
                continue;
            }

            return index;
        }
    }

    // -------------------------------------------------------------------------
    // sliceHeader (feature 509)
    // -------------------------------------------------------------------------

    /// <summary>
    /// Returns a slice of <paramref name="hdr"/> containing the value of <paramref name="key"/>,
    /// or <c>null</c> if not found.
    /// The returned slice shares memory with <paramref name="hdr"/>.
    /// Mirrors Go <c>sliceHeader</c>.
    /// </summary>
    public static ReadOnlyMemory<byte>? SliceHeader(string key, byte[] hdr)
    {
        if (hdr.Length == 0) return null;

        int index = GetHeaderKeyIndex(key, hdr);
        if (index == -1) return null;

        // Skip over key + ':' separator.
        index += key.Length + 1;
        int hdrLen = hdr.Length;

        // Skip leading whitespace.
        while (index < hdrLen && hdr[index] == ' ')
            index++;

        int start = index;
        // Collect until CRLF.
        while (index < hdrLen)
        {
            if (hdr[index] == '\r' && index + 1 < hdrLen && hdr[index + 1] == '\n')
                break;
            index++;
        }

        // Return a slice with capped length == value length (no extra capacity).
        return new ReadOnlyMemory<byte>(hdr, start, index - start);
    }

    // -------------------------------------------------------------------------
    // getHeader (feature 508)
    // -------------------------------------------------------------------------

    /// <summary>
    /// Returns a copy of the value for the header named <paramref name="key"/>,
    /// or <c>null</c> if not found.
    /// Mirrors Go <c>getHeader</c>.
    /// </summary>
    public static byte[]? GetHeader(string key, byte[] hdr)
    {
        var slice = SliceHeader(key, hdr);
        if (slice is null) return null;

        // Return a fresh copy.
        return slice.Value.ToArray();
    }

    // -------------------------------------------------------------------------
    // setHeader (feature 511)
    // -------------------------------------------------------------------------

    /// <summary>
    /// Replaces the value of the first existing <paramref name="key"/> header in
    /// <paramref name="hdr"/>, or appends a new header if the key is absent.
    /// Returns a new buffer when the new value is larger; modifies in-place otherwise.
    /// Mirrors Go <c>setHeader</c>.
    /// </summary>
    public static byte[] SetHeader(string key, string val, byte[] hdr)
    {
        int start = GetHeaderKeyIndex(key, hdr);
        if (start >= 0)
        {
            int valStart = start + key.Length + 1; // skip past ':'
            int hdrLen   = hdr.Length;

            // Preserve a single leading space if present.
            if (valStart < hdrLen && hdr[valStart] == ' ')
                valStart++;

            // Find the CR before the CRLF.
            int crIdx = IndexOf(hdr, [(byte)'\r'], valStart);
            if (crIdx < 0) return hdr; // malformed

            int valEnd    = crIdx;
            int oldValLen = valEnd - valStart;
            var valBytes  = Encoding.ASCII.GetBytes(val);

            int extra = valBytes.Length - oldValLen;
            if (extra > 0)
            {
                // New value is larger — must allocate a new buffer.
                var newHdr = new byte[hdrLen + extra];
                Buffer.BlockCopy(hdr, 0,        newHdr, 0,              valStart);
                Buffer.BlockCopy(valBytes, 0,   newHdr, valStart,       valBytes.Length);
                Buffer.BlockCopy(hdr, valEnd,   newHdr, valStart + valBytes.Length, hdrLen - valEnd);
                return newHdr;
            }

            // Write in place (new value fits).
            int n = valBytes.Length;
            Buffer.BlockCopy(valBytes, 0, hdr, valStart, n);
            // Shift remainder left.
            Buffer.BlockCopy(hdr, valEnd, hdr, valStart + n, hdrLen - valEnd);
            return hdr[..(valStart + n + hdrLen - valEnd)];
        }

        // Key not present — append.
        bool hasTrailingCrLf = hdr.Length >= 2
            && hdr[^2] == '\r'
            && hdr[^1] == '\n';

        byte[] suffix;
        if (hasTrailingCrLf)
        {
            // Strip trailing CRLF, append "key: val\r\n\r\n".
            suffix = Encoding.ASCII.GetBytes($"{key}: {val}\r\n");
            var result = new byte[hdr.Length - 2 + suffix.Length + 2];
            Buffer.BlockCopy(hdr, 0, result, 0, hdr.Length - 2);
            Buffer.BlockCopy(suffix, 0, result, hdr.Length - 2, suffix.Length);
            result[^2] = (byte)'\r';
            result[^1] = (byte)'\n';
            return result;
        }

        suffix = Encoding.ASCII.GetBytes($"{key}: {val}\r\n");
        var newBuf = new byte[hdr.Length + suffix.Length];
        Buffer.BlockCopy(hdr, 0, newBuf, 0, hdr.Length);
        Buffer.BlockCopy(suffix, 0, newBuf, hdr.Length, suffix.Length);
        return newBuf;
    }

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

    private static int IndexOf(byte[] haystack, byte[] needle, int offset)
    {
        var span = haystack.AsSpan(offset);
        int idx  = span.IndexOf(needle);
        return idx < 0 ? -1 : offset + idx;
    }

    private static int IndexOf(List<byte> haystack, byte[] needle, int offset)
    {
        for (int i = offset; i <= haystack.Count - needle.Length; i++)
        {
            bool match = true;
            for (int j = 0; j < needle.Length; j++)
            {
                if (haystack[i + j] != needle[j]) { match = false; break; }
            }
            if (match) return i;
        }
        return -1;
    }

    /// <summary>Returns the offset of the first \r\n in <paramref name="hdr"/> at or after <paramref name="offset"/>.</summary>
    private static int IndexOfCrLf(byte[] hdr, int offset)
    {
        var span = hdr.AsSpan(offset);
        int idx  = span.IndexOf(CrLfBytes);
        return idx; // relative to offset
    }

    private static int IndexOfCrLf(List<byte> hdr, int offset)
    {
        for (int i = offset; i < hdr.Count - 1; i++)
        {
            if (hdr[i] == '\r' && hdr[i + 1] == '\n')
                return i - offset;
        }
        return -1;
    }
}
