using System.IO.Compression;

namespace NATS.Server.WebSocket;

/// <summary>
/// Negotiated permessage-deflate parameters per RFC 7692 Section 7.1.
/// Captures the results of extension parameter negotiation during the
/// WebSocket upgrade handshake.
/// </summary>
public readonly record struct WsDeflateParams(
    bool ServerNoContextTakeover,
    bool ClientNoContextTakeover,
    int ServerMaxWindowBits,
    int ClientMaxWindowBits)
{
    /// <summary>
    /// Default parameters matching NATS Go server behavior:
    /// both sides use no_context_takeover, default 15-bit windows.
    /// </summary>
    public static readonly WsDeflateParams Default = new(
        ServerNoContextTakeover: true,
        ClientNoContextTakeover: true,
        ServerMaxWindowBits: 15,
        ClientMaxWindowBits: 15);

    /// <summary>
    /// Builds the Sec-WebSocket-Extensions response header value from negotiated parameters.
    /// Only includes parameters that differ from the default RFC values.
    /// Reference: RFC 7692 Section 7.1.
    /// </summary>
    public string ToResponseHeaderValue()
    {
        var parts = new List<string> { WsConstants.PmcExtension };

        if (ServerNoContextTakeover)
            parts.Add(WsConstants.PmcSrvNoCtx);
        if (ClientNoContextTakeover)
            parts.Add(WsConstants.PmcCliNoCtx);
        if (ServerMaxWindowBits is > 0 and < 15)
            parts.Add($"server_max_window_bits={ServerMaxWindowBits}");
        if (ClientMaxWindowBits is > 0 and < 15)
            parts.Add($"client_max_window_bits={ClientMaxWindowBits}");

        return string.Join("; ", parts);
    }
}

/// <summary>
/// Parses and negotiates permessage-deflate extension parameters from the
/// Sec-WebSocket-Extensions header per RFC 7692 Section 7.
/// Reference: golang/nats-server/server/websocket.go — wsPMCExtensionSupport.
/// </summary>
public static class WsDeflateNegotiator
{
    /// <summary>
    /// Parses the Sec-WebSocket-Extensions header value and negotiates
    /// permessage-deflate parameters. Returns null if no valid
    /// permessage-deflate offer is found.
    /// </summary>
    public static WsDeflateParams? Negotiate(string? extensionHeader)
    {
        if (string.IsNullOrEmpty(extensionHeader))
            return null;

        // The header may contain multiple extensions separated by commas
        var extensions = extensionHeader.Split(',');
        foreach (var extension in extensions)
        {
            var trimmed = extension.Trim();
            var parts = trimmed.Split(';');

            // First part must be the extension name
            if (parts.Length == 0)
                continue;

            if (!string.Equals(parts[0].Trim(), WsConstants.PmcExtension, StringComparison.OrdinalIgnoreCase))
                continue;

            // Found permessage-deflate — parse parameters
            // Note: serverNoCtx and clientNoCtx are parsed but always overridden
            // with true below (NATS enforces no_context_takeover for both sides).
            int serverMaxWindowBits = 15;
            int clientMaxWindowBits = 15;

            for (int i = 1; i < parts.Length; i++)
            {
                var param = parts[i].Trim();

                if (string.Equals(param, WsConstants.PmcSrvNoCtx, StringComparison.OrdinalIgnoreCase))
                {
                    // Parsed but overridden: NATS always enforces no_context_takeover.
                }
                else if (string.Equals(param, WsConstants.PmcCliNoCtx, StringComparison.OrdinalIgnoreCase))
                {
                    // Parsed but overridden: NATS always enforces no_context_takeover.
                }
                else if (param.StartsWith("server_max_window_bits", StringComparison.OrdinalIgnoreCase))
                {
                    serverMaxWindowBits = ParseWindowBits(param, 15);
                }
                else if (param.StartsWith("client_max_window_bits", StringComparison.OrdinalIgnoreCase))
                {
                    // client_max_window_bits with no value means the client supports it
                    // and the server may choose a value. Per RFC 7692 Section 7.1.2.2,
                    // an offer with just "client_max_window_bits" (no value) indicates
                    // the client can accept any value 8-15.
                    clientMaxWindowBits = ParseWindowBits(param, 15);
                }
            }

            // NATS server always enforces no_context_takeover for both sides
            // (matching Go behavior) to avoid holding compressor state per connection.
            return new WsDeflateParams(
                ServerNoContextTakeover: true,
                ClientNoContextTakeover: true,
                ServerMaxWindowBits: ClampWindowBits(serverMaxWindowBits),
                ClientMaxWindowBits: ClampWindowBits(clientMaxWindowBits));
        }

        return null;
    }

    private static int ParseWindowBits(string param, int defaultValue)
    {
        var eqIdx = param.IndexOf('=');
        if (eqIdx < 0)
            return defaultValue;

        var valueStr = param[(eqIdx + 1)..].Trim();
        if (int.TryParse(valueStr, out var bits))
            return bits;

        return defaultValue;
    }

    private static int ClampWindowBits(int bits)
    {
        // RFC 7692: valid range is 8-15
        if (bits < 8) return 8;
        if (bits > 15) return 15;
        return bits;
    }
}

/// <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();
    }
}
