using System.IO.Compression;

namespace NATS.Server.WebSocket;

/// <summary>
/// permessage-deflate compression/decompression for WebSocket frames (RFC 7692).
/// Ported from golang/nats-server/server/websocket.go lines 403-440 and 1391-1466.
/// </summary>
public static class WsCompression
{
    /// <summary>
    /// Compresses data using deflate. Removes trailing 4 bytes (sync marker)
    /// per RFC 7692 Section 7.2.1.
    /// </summary>
    /// <remarks>
    /// We call Flush() but intentionally do not Dispose() the DeflateStream before
    /// reading output, because Dispose writes a final deflate block (0x03 0x00) that
    /// would be corrupted by the 4-byte tail strip. Flush() alone writes a sync flush
    /// ending with 0x00 0x00 0xff 0xff, matching Go's flate.Writer.Flush() behavior.
    /// </remarks>
    public static byte[] Compress(ReadOnlySpan<byte> data)
    {
        var output = new MemoryStream();
        var deflate = new DeflateStream(output, CompressionLevel.Fastest, leaveOpen: true);
        try
        {
            deflate.Write(data);
            deflate.Flush();

            var compressed = output.ToArray();

            // Remove trailing 4-byte sync marker (0x00 0x00 0xff 0xff) per RFC 7692
            if (compressed.Length >= 4)
                return compressed[..^4];

            return compressed;
        }
        finally
        {
            deflate.Dispose();
            output.Dispose();
        }
    }

    /// <summary>
    /// Decompresses collected compressed buffers.
    /// Appends trailer bytes before decompressing per RFC 7692 Section 7.2.2.
    /// Ported from golang/nats-server/server/websocket.go lines 403-440.
    /// The Go code appends compressLastBlock (9 bytes) which includes the sync
    /// marker plus a final empty stored block to signal end-of-stream to the
    /// flate reader.
    /// </summary>
    public static byte[] Decompress(List<byte[]> compressedBuffers, int maxPayload)
    {
        if (maxPayload <= 0)
            maxPayload = 1024 * 1024; // Default 1MB

        // Concatenate all compressed buffers + trailer.
        // Per RFC 7692 Section 7.2.2, append the sync flush marker (0x00 0x00 0xff 0xff)
        // that was stripped during compression. The Go reference appends compressLastBlock
        // (9 bytes) for Go's flate reader; .NET's DeflateStream only needs the 4-byte trailer.
        int totalLen = 0;
        foreach (var buf in compressedBuffers)
            totalLen += buf.Length;
        totalLen += WsConstants.DecompressTrailer.Length;

        var combined = new byte[totalLen];
        int offset = 0;
        foreach (var buf in compressedBuffers)
        {
            buf.CopyTo(combined, offset);
            offset += buf.Length;
        }

        WsConstants.DecompressTrailer.CopyTo(combined, offset);

        using var input = new MemoryStream(combined);
        using var deflate = new DeflateStream(input, CompressionMode.Decompress);
        using var output = new MemoryStream();

        var readBuf = new byte[4096];
        int totalRead = 0;
        int n;
        while ((n = deflate.Read(readBuf, 0, readBuf.Length)) > 0)
        {
            totalRead += n;
            if (totalRead > maxPayload)
                throw new InvalidOperationException("decompressed data exceeds maximum payload size");
            output.Write(readBuf, 0, n);
        }

        return output.ToArray();
    }
}