// 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
//
// 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/websocket.go in the NATS server Go source.

using System.IO.Compression;
using ZB.MOM.NatsNet.Server.Internal;

namespace ZB.MOM.NatsNet.Server.WebSocket;

/// <summary>
/// Per-connection WebSocket read state.
/// Mirrors Go <c>wsReadInfo</c> struct in server/websocket.go.
/// </summary>
internal sealed class WsReadInfo
{
    public int Rem { get; set; }
    public bool FrameStart { get; set; }
    public bool FinalFrameReceived { get; set; }
    public bool FrameCompressed { get; set; }
    public bool Mask { get; set; }
    public byte MaskKeyPosition { get; set; }
    public byte[] MaskKey { get; } = new byte[4];
    public List<byte[]> CompressedBuffers { get; } = [];
    public int CompressedOffset { get; set; }

    public void Init()
    {
        FrameStart = true;
        FinalFrameReceived = true;
    }

    public int Read(byte[] destination, int offset, int count)
    {
        if (count == 0)
            return 0;
        if (CompressedBuffers.Count == 0)
            return 0;

        var copied = 0;
        var remaining = count;
        while (CompressedBuffers.Count > 0 && remaining > 0)
        {
            var buffer = CompressedBuffers[0];
            var available = buffer.Length - CompressedOffset;
            if (available <= 0)
            {
                NextCBuf();
                continue;
            }

            var n = Math.Min(available, remaining);
            Array.Copy(buffer, CompressedOffset, destination, offset + copied, n);
            copied += n;
            remaining -= n;
            CompressedOffset += n;
            NextCBuf();
        }

        return copied;
    }

    public byte[]? NextCBuf()
    {
        if (CompressedBuffers.Count == 0)
            return null;
        if (CompressedOffset != CompressedBuffers[0].Length)
            return CompressedBuffers[0];

        CompressedOffset = 0;
        if (CompressedBuffers.Count == 1)
        {
            CompressedBuffers.Clear();
            return null;
        }

        CompressedBuffers.RemoveAt(0);
        return CompressedBuffers[0];
    }

    public byte ReadByte()
    {
        if (CompressedBuffers.Count == 0)
            throw new EndOfStreamException();

        var b = CompressedBuffers[0][CompressedOffset];
        CompressedOffset++;
        NextCBuf();
        return b;
    }

    public byte[] Decompress(int maxPayload)
    {
        if (maxPayload <= 0)
            maxPayload = ServerConstants.MaxPayloadSize;

        CompressedOffset = 0;
        var input = new MemoryStream();
        foreach (var buffer in CompressedBuffers)
            input.Write(buffer, 0, buffer.Length);
        input.Write(WsConstants.CompressLastBlock, 0, WsConstants.CompressLastBlock.Length);
        input.Position = 0;

        using var deflate = new DeflateStream(input, System.IO.Compression.CompressionMode.Decompress, leaveOpen: true);
        using var output = new MemoryStream();
        var tmp = new byte[4096];
        while (true)
        {
            var n = deflate.Read(tmp, 0, tmp.Length);
            if (n == 0)
                break;
            if (output.Length + n > maxPayload)
                throw ServerErrors.ErrMaxPayload;
            output.Write(tmp, 0, n);
        }

        CompressedBuffers.Clear();
        return output.ToArray();
    }

    public void Unmask(byte[] buf) => Unmask(buf.AsSpan());

    public void Unmask(Span<byte> buf)
    {
        var p = (int)MaskKeyPosition;
        if (buf.Length < 16)
        {
            for (var i = 0; i < buf.Length; i++)
            {
                buf[i] ^= MaskKey[p & 3];
                p++;
            }
            MaskKeyPosition = (byte)(p & 3);
            return;
        }

        var key8 = new byte[8];
        for (var i = 0; i < 8; i++)
            key8[i] = MaskKey[(p + i) & 3];
        var mask64 = BitConverter.ToUInt64(key8, 0);

        var n8 = (buf.Length / 8) * 8;
        for (var i = 0; i < n8; i += 8)
        {
            var value = BitConverter.ToUInt64(buf[i..(i + 8)]) ^ mask64;
            var bytes = BitConverter.GetBytes(value);
            bytes.CopyTo(buf[i..(i + 8)]);
        }

        for (var i = n8; i < buf.Length; i++)
        {
            buf[i] ^= MaskKey[p & 3];
            p++;
        }
        MaskKeyPosition = (byte)(p & 3);
    }
}

/// <summary>
/// Client-level WebSocket runtime state.
/// Mirrors Go <c>websocket</c> struct in websocket.go.
/// </summary>
internal sealed class WebsocketConnection
{
    public List<byte[]> Frames { get; } = [];
    public long FrameSize { get; set; }
    public byte[]? CloseMessage { get; set; }
    public bool Compress { get; set; }
    public bool CloseSent { get; set; }
    public bool Browser { get; set; }
    public bool NoCompressedFragment { get; set; }
    public bool MaskRead { get; set; } = true;
    public bool MaskWrite { get; set; }
    public string ClientIP { get; set; } = string.Empty;
}

/// <summary>
/// Server-level WebSocket state, shared across all WebSocket connections.
/// Mirrors Go <c>srvWebsocket</c> struct in server/websocket.go.
/// </summary>
internal sealed class SrvWebsocket
{
    public Lock Mu { get; } = new();
    public System.Net.Sockets.TcpListener? Listener { get; set; }
    public Exception? ListenerErr { get; set; }
    public Dictionary<string, AllowedOrigin> AllowedOrigins { get; } = new(StringComparer.OrdinalIgnoreCase);
    public bool SameOrigin { get; set; }
    public List<string> ConnectUrls { get; } = [];
    public RefCountedUrlSet ConnectUrlsMap { get; } = new();
    public bool AuthOverride { get; set; }
    public string RawHeaders { get; set; } = string.Empty;
    public System.Net.Security.SslServerAuthenticationOptions? TlsConfig { get; set; }
    public bool Compression { get; set; }
    public string Host { get; set; } = string.Empty;
    public int Port { get; set; }

    public Exception? checkOrigin(string requestHost, string? origin)
    {
        if (!SameOrigin && AllowedOrigins.Count == 0)
            return null;
        if (string.IsNullOrWhiteSpace(origin))
            return new InvalidOperationException("origin header missing");
        if (!Uri.TryCreate(origin, UriKind.Absolute, out var uri))
            return new InvalidOperationException("invalid origin");
        if (SameOrigin && !string.Equals(uri.Host, requestHost, StringComparison.OrdinalIgnoreCase))
            return new InvalidOperationException("origin not same as host");
        if (AllowedOrigins.Count == 0)
            return null;
        if (!AllowedOrigins.TryGetValue(uri.Host, out var allowed))
            return new InvalidOperationException("origin host not allowed");

        var port = uri.IsDefaultPort
            ? uri.Scheme.Equals("https", StringComparison.OrdinalIgnoreCase) ? "443" : "80"
            : uri.Port.ToString();
        if (!string.Equals(allowed.Scheme, uri.Scheme, StringComparison.OrdinalIgnoreCase))
            return new InvalidOperationException("origin scheme not allowed");
        if (!string.Equals(allowed.Port, port, StringComparison.Ordinal))
            return new InvalidOperationException("origin port not allowed");
        return null;
    }
}

internal readonly record struct AllowedOrigin(string Scheme, string Port);
