// Copyright 2020-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

using System.Buffers.Binary;
using System.Security.Cryptography;

namespace ZB.MOM.NatsNet.Server.WebSocket;

internal static class WebSocketHelpers
{
    public static (byte[] bytes, int newPos) WsGet(Stream reader, byte[] buf, int pos, int needed)
    {
        var available = buf.Length - pos;
        if (available >= needed)
            return (buf[pos..(pos + needed)], pos + needed);

        var b = new byte[needed];
        var start = 0;
        if (available > 0)
        {
            System.Buffer.BlockCopy(buf, pos, b, 0, available);
            start = available;
        }

        while (start < needed)
        {
            var n = reader.Read(b, start, needed - start);
            if (n <= 0)
                throw new EndOfStreamException();
            start += n;
        }

        return (b, pos + available);
    }

    public static bool WsIsControlFrame(WsOpCode frameType) => frameType >= WsOpCode.Close;

    public static (byte[] header, byte[]? key) WsCreateFrameHeader(bool useMasking, bool compressed, WsOpCode frameType, int length)
    {
        var frameHeader = NbPool.Get(WsConstants.MaxFrameHeaderSize);
        var (n, key) = WsFillFrameHeader(frameHeader, useMasking, first: true, final: true, compressed, frameType, length);
        return (frameHeader[..n], key);
    }

    public static (int n, byte[]? key) WsFillFrameHeader(
        byte[] frameHeader,
        bool useMasking,
        bool first,
        bool final,
        bool compressed,
        WsOpCode frameType,
        int length)
    {
        byte b0 = 0;
        if (first)
            b0 = (byte)frameType;
        if (final)
            b0 |= WsConstants.FinalBit;
        if (compressed)
            b0 |= WsConstants.Rsv1Bit;

        byte b1 = 0;
        if (useMasking)
            b1 |= WsConstants.MaskBit;

        int n;
        if (length <= 125)
        {
            n = 2;
            frameHeader[0] = b0;
            frameHeader[1] = (byte)(b1 | (byte)length);
        }
        else if (length < 65536)
        {
            n = 4;
            frameHeader[0] = b0;
            frameHeader[1] = (byte)(b1 | 126);
            BinaryPrimitives.WriteUInt16BigEndian(frameHeader.AsSpan(2), (ushort)length);
        }
        else
        {
            n = 10;
            frameHeader[0] = b0;
            frameHeader[1] = (byte)(b1 | 127);
            BinaryPrimitives.WriteUInt64BigEndian(frameHeader.AsSpan(2), (ulong)length);
        }

        byte[]? key = null;
        if (useMasking)
        {
            key = new byte[4];
            RandomNumberGenerator.Fill(key);
            Array.Copy(key, 0, frameHeader, n, 4);
            n += 4;
        }

        return (n, key);
    }

    public static void WsMaskBuf(ReadOnlySpan<byte> key, Span<byte> buffer)
    {
        for (var i = 0; i < buffer.Length; i++)
            buffer[i] ^= key[i & 0x3];
    }

    public static void WsMaskBufs(ReadOnlySpan<byte> key, IReadOnlyList<byte[]> buffers)
    {
        var pos = 0;
        for (var i = 0; i < buffers.Count; i++)
        {
            var buffer = buffers[i];
            for (var j = 0; j < buffer.Length; j++)
            {
                buffer[j] ^= key[pos & 0x3];
                pos++;
            }
        }
    }

    public static byte[] WsCreateCloseMessage(int status, string body)
    {
        if (body.Length > WsConstants.MaxControlPayloadSize - 2)
            body = string.Concat(body.AsSpan(0, WsConstants.MaxControlPayloadSize - 5), "...");

        var payload = System.Text.Encoding.UTF8.GetBytes(body);
        var buffer = new byte[2 + payload.Length];
        BinaryPrimitives.WriteUInt16BigEndian(buffer.AsSpan(0, 2), (ushort)status);
        payload.CopyTo(buffer.AsSpan(2));
        return buffer;
    }
}
