// 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.

using System.Buffers.Binary;
using System.Text;
using ZB.MOM.NatsNet.Server.WebSocket;

namespace ZB.MOM.NatsNet.Server;

public sealed partial class ClientConnection
{
    internal (List<OutboundChunk> chunks, long attempted) WsCollapsePtoNB() => CollapsePtoNB();

    internal List<byte[]> WsRead(WsReadInfo readInfo, Stream reader, byte[] buffer)
    {
        var bufs = new List<byte[]>();
        var pos = 0;
        var max = buffer.Length;
        var maxPayload = Volatile.Read(ref _mpay);

        while (pos != max)
        {
            if (readInfo.FrameStart)
            {
                var b0 = buffer[pos];
                var frameType = (WsOpCode)(b0 & 0xF);
                var final = (b0 & WsConstants.FinalBit) != 0;
                var compressed = (b0 & WsConstants.Rsv1Bit) != 0;
                pos++;

                var (firstLen, newPos) = WebSocketHelpers.WsGet(reader, buffer, pos, 1);
                pos = newPos;
                var b1 = firstLen[0];

                if (readInfo.Mask && (b1 & WsConstants.MaskBit) == 0)
                    throw WsHandleProtocolError("mask bit missing");

                readInfo.Rem = b1 & 0x7F;

                switch (frameType)
                {
                    case WsOpCode.Ping:
                    case WsOpCode.Pong:
                    case WsOpCode.Close:
                        if (readInfo.Rem > WsConstants.MaxControlPayloadSize)
                            throw WsHandleProtocolError($"control frame length bigger than maximum allowed of {WsConstants.MaxControlPayloadSize} bytes");
                        if (!final)
                            throw WsHandleProtocolError("control frame does not have final bit set");
                        break;
                    case WsOpCode.Text:
                    case WsOpCode.Binary:
                        if (!readInfo.FinalFrameReceived)
                            throw WsHandleProtocolError("new message started before final frame for previous message was received");
                        readInfo.FinalFrameReceived = final;
                        readInfo.FrameCompressed = compressed;
                        break;
                    case WsOpCode.Continuation:
                        if (readInfo.FinalFrameReceived || compressed)
                            throw WsHandleProtocolError("invalid continuation frame");
                        readInfo.FinalFrameReceived = final;
                        break;
                    default:
                        throw WsHandleProtocolError($"unknown opcode {(int)frameType}");
                }

                if (readInfo.Rem == 126)
                {
                    var (extended, p) = WebSocketHelpers.WsGet(reader, buffer, pos, 2);
                    pos = p;
                    readInfo.Rem = BinaryPrimitives.ReadUInt16BigEndian(extended);
                }
                else if (readInfo.Rem == 127)
                {
                    var (extended, p) = WebSocketHelpers.WsGet(reader, buffer, pos, 8);
                    pos = p;
                    readInfo.Rem = checked((int)BinaryPrimitives.ReadUInt64BigEndian(extended));
                }

                if (readInfo.Mask)
                {
                    var (maskKey, p) = WebSocketHelpers.WsGet(reader, buffer, pos, 4);
                    pos = p;
                    Array.Copy(maskKey, 0, readInfo.MaskKey, 0, 4);
                    readInfo.MaskKeyPosition = 0;
                }

                if (WebSocketHelpers.WsIsControlFrame(frameType))
                {
                    pos = WsHandleControlFrame(readInfo, frameType, reader, buffer, pos);
                    continue;
                }

                readInfo.FrameStart = false;
            }

            if (pos >= max)
                continue;

            var n = readInfo.Rem;
            if (pos + n > max)
                n = max - pos;

            var payload = buffer.AsSpan(pos, n).ToArray();
            pos += n;
            readInfo.Rem -= n;

            if (readInfo.Mask)
                readInfo.Unmask(payload);

            var addToBufs = true;
            if (readInfo.FrameCompressed)
            {
                addToBufs = false;
                readInfo.CompressedBuffers.Add(payload);
                if (readInfo.FinalFrameReceived && readInfo.Rem == 0)
                {
                    payload = readInfo.Decompress(maxPayload);
                    readInfo.FrameCompressed = false;
                    addToBufs = true;
                }
            }

            if (addToBufs)
                bufs.Add(payload);

            if (readInfo.Rem == 0)
                readInfo.FrameStart = true;
        }

        return bufs;
    }

    internal int WsHandleControlFrame(WsReadInfo readInfo, WsOpCode frameType, Stream networkConnection, byte[] buffer, int pos)
    {
        byte[] payload = [];
        if (readInfo.Rem > 0)
        {
            (payload, pos) = WebSocketHelpers.WsGet(networkConnection, buffer, pos, readInfo.Rem);
            if (readInfo.Mask)
                readInfo.Unmask(payload);
            readInfo.Rem = 0;
        }

        switch (frameType)
        {
            case WsOpCode.Close:
            {
                var status = WsConstants.CloseNoStatusReceived;
                string body = string.Empty;

                var payloadLength = payload.Length;
                var hasStatus = payloadLength >= WsConstants.CloseStatusSize;
                var hasBody = payloadLength > WsConstants.CloseStatusSize;
                if (hasStatus)
                {
                    status = BinaryPrimitives.ReadUInt16BigEndian(payload.AsSpan(0, 2));
                    if (hasBody)
                    {
                        body = Encoding.UTF8.GetString(payload.AsSpan(WsConstants.CloseStatusSize));
                        if (!Encoding.UTF8.GetBytes(body).AsSpan().SequenceEqual(payload.AsSpan(WsConstants.CloseStatusSize)))
                        {
                            status = WsConstants.CloseInvalidPayloadData;
                            body = "invalid utf8 body in close frame";
                        }
                    }
                }

                byte[]? closeMessage = null;
                if (status != WsConstants.CloseNoStatusReceived)
                    closeMessage = WebSocketHelpers.WsCreateCloseMessage(status, body);

                WsEnqueueControlMessage(WsOpCode.Close, closeMessage ?? []);
                throw new EndOfStreamException();
            }
            case WsOpCode.Ping:
                WsEnqueueControlMessage(WsOpCode.Pong, payload);
                break;
            case WsOpCode.Pong:
                break;
        }

        return pos;
    }

    internal void WsEnqueueControlMessage(WsOpCode controlMessage, byte[] payload)
    {
        lock (_mu)
            WsEnqueueControlMessageLocked(controlMessage, payload);
    }

    internal void WsEnqueueControlMessageLocked(WsOpCode controlMessage, byte[] payload)
    {
        if (Ws == null)
            Ws = new WebsocketConnection();

        var useMasking = Ws.MaskWrite;
        var headerSize = 2 + (useMasking ? 4 : 0);
        var control = NbPool.Get(headerSize + payload.Length);

        var (n, key) = WebSocketHelpers.WsFillFrameHeader(control, useMasking, first: true, final: true, compressed: false, controlMessage, payload.Length);
        var totalLen = n;
        if (payload.Length > 0)
        {
            Array.Copy(payload, 0, control, n, payload.Length);
            if (useMasking && key != null)
                WebSocketHelpers.WsMaskBuf(key, control.AsSpan(n, payload.Length));
            totalLen += payload.Length;
        }

        var frame = control[..totalLen];
        OutPb += totalLen;
        if (controlMessage == WsOpCode.Close)
        {
            Ws.CloseSent = true;
            Ws.CloseMessage = frame;
        }
        else
        {
            Ws.Frames.Add(frame);
            Ws.FrameSize += frame.Length;
        }

        FlushSignal();
    }

    internal void WsEnqueueCloseMessage(ClosedState reason)
    {
        var status = reason switch
        {
            ClosedState.ClientClosed => WsConstants.CloseNormalClosure,
            ClosedState.AuthenticationTimeout or ClosedState.AuthenticationViolation or ClosedState.SlowConsumerPendingBytes or
            ClosedState.SlowConsumerWriteDeadline or ClosedState.MaxAccountConnectionsExceeded or ClosedState.MaxConnectionsExceeded or
            ClosedState.MaxControlLineExceeded or ClosedState.MaxSubscriptionsExceeded or ClosedState.MissingAccount or
            ClosedState.AuthenticationExpired or ClosedState.Revocation => WsConstants.ClosePolicyViolation,
            ClosedState.TlsHandshakeError => WsConstants.CloseTlsHandshake,
            ClosedState.ParseError or ClosedState.ProtocolViolation or ClosedState.BadClientProtocolVersion => WsConstants.CloseProtocolError,
            ClosedState.MaxPayloadExceeded => WsConstants.CloseMessageTooBig,
            ClosedState.WriteError or ClosedState.ReadError or ClosedState.StaleConnection or ClosedState.ServerShutdown => WsConstants.CloseGoingAway,
            _ => WsConstants.CloseInternalError,
        };

        var body = WebSocketHelpers.WsCreateCloseMessage(status, reason.ToString());
        WsEnqueueControlMessageLocked(WsOpCode.Close, body);
    }

    internal Exception WsHandleProtocolError(string message)
    {
        var payload = WebSocketHelpers.WsCreateCloseMessage(WsConstants.CloseProtocolError, message);
        WsEnqueueControlMessage(WsOpCode.Close, payload);
        return new InvalidDataException(message);
    }
}
