namespace NATS.Server.WebSocket;

/// <summary>
/// Stream wrapper that transparently frames/deframes WebSocket around raw TCP I/O.
/// NatsClient uses this as its _stream -- FillPipeAsync and RunWriteLoopAsync work unchanged.
/// Ported from golang/nats-server/server/websocket.go wsUpgrade/wrapWebsocket pattern.
/// </summary>
public sealed class WsConnection : Stream
{
    private readonly Stream _inner;
    private readonly bool _compress;
    private readonly bool _maskRead;
    private readonly bool _maskWrite;
    private readonly bool _browser;
    private readonly bool _noCompFrag;
    private WsReadInfo _readInfo;
    // Read-side state: accessed only from the single FillPipeAsync reader task (no synchronization needed)
    private readonly Queue<byte[]> _readQueue = new();
    private int _readOffset;
    private readonly object _writeLock = new();
    private readonly List<ControlFrameAction> _pendingControlWrites = [];

    public bool CloseReceived => _readInfo.CloseReceived;
    public int CloseStatus => _readInfo.CloseStatus;

    public WsConnection(Stream inner, bool compress, bool maskRead, bool maskWrite, bool browser, bool noCompFrag)
    {
        _inner = inner;
        _compress = compress;
        _maskRead = maskRead;
        _maskWrite = maskWrite;
        _browser = browser;
        _noCompFrag = noCompFrag;
        _readInfo = new WsReadInfo(expectMask: maskRead);
    }

    public override async ValueTask<int> ReadAsync(Memory<byte> buffer, CancellationToken ct = default)
    {
        // Drain any buffered decoded payloads first
        if (_readQueue.Count > 0)
            return DrainReadQueue(buffer.Span);

        while (true)
        {
            // Read raw bytes from inner stream
            var rawBuf = new byte[Math.Max(buffer.Length, 4096)];
            int bytesRead = await _inner.ReadAsync(rawBuf.AsMemory(), ct);
            if (bytesRead == 0) return 0;

            // Decode frames
            var payloads = WsReadInfo.ReadFrames(_readInfo, new MemoryStream(rawBuf, 0, bytesRead), bytesRead, maxPayload: 1024 * 1024);

            // Collect control frame responses
            if (_readInfo.PendingControlFrames.Count > 0)
            {
                lock (_writeLock)
                    _pendingControlWrites.AddRange(_readInfo.PendingControlFrames);
                _readInfo.PendingControlFrames.Clear();
                // Write pending control frames
                await FlushControlFramesAsync(ct);
            }

            if (_readInfo.CloseReceived)
                return 0;

            foreach (var payload in payloads)
                _readQueue.Enqueue(payload);

            // If no payloads were decoded (e.g. only frame headers were read),
            // continue reading instead of returning 0 which signals end-of-stream
            if (_readQueue.Count > 0)
                return DrainReadQueue(buffer.Span);
        }
    }

    public override async ValueTask WriteAsync(ReadOnlyMemory<byte> buffer, CancellationToken ct = default)
    {
        var data = buffer.Span;

        if (_compress && data.Length > WsConstants.CompressThreshold)
        {
            var compressed = WsCompression.Compress(data);
            await WriteFramedAsync(compressed, compressed: true, ct);
        }
        else
        {
            await WriteFramedAsync(data.ToArray(), compressed: false, ct);
        }
    }

    private async ValueTask WriteFramedAsync(byte[] payload, bool compressed, CancellationToken ct)
    {
        if (_browser && payload.Length > WsConstants.FrameSizeForBrowsers && !(_noCompFrag && compressed))
        {
            // Fragment for browsers
            int offset = 0;
            bool first = true;
            while (offset < payload.Length)
            {
                int chunkLen = Math.Min(WsConstants.FrameSizeForBrowsers, payload.Length - offset);
                bool final = offset + chunkLen >= payload.Length;
                var fh = new byte[WsConstants.MaxFrameHeaderSize];
                var (n, key) = WsFrameWriter.FillFrameHeader(fh, _maskWrite,
                    first: first, final: final, compressed: first && compressed,
                    opcode: WsConstants.BinaryMessage, payloadLength: chunkLen);

                var chunk = payload.AsSpan(offset, chunkLen).ToArray();
                if (_maskWrite && key != null)
                    WsFrameWriter.MaskBuf(key, chunk);

                await _inner.WriteAsync(fh.AsMemory(0, n), ct);
                await _inner.WriteAsync(chunk.AsMemory(), ct);
                offset += chunkLen;
                first = false;
            }
        }
        else
        {
            var (header, key) = WsFrameWriter.CreateFrameHeader(_maskWrite, compressed, WsConstants.BinaryMessage, payload.Length);
            if (_maskWrite && key != null)
                WsFrameWriter.MaskBuf(key, payload);
            await _inner.WriteAsync(header.AsMemory(), ct);
            await _inner.WriteAsync(payload.AsMemory(), ct);
        }
    }

    private async Task FlushControlFramesAsync(CancellationToken ct)
    {
        List<ControlFrameAction> toWrite;
        lock (_writeLock)
        {
            if (_pendingControlWrites.Count == 0) return;
            toWrite = [.. _pendingControlWrites];
            _pendingControlWrites.Clear();
        }

        foreach (var action in toWrite)
        {
            var frame = WsFrameWriter.BuildControlFrame(action.Opcode, action.Payload, _maskWrite);
            await _inner.WriteAsync(frame, ct);
        }
        await _inner.FlushAsync(ct);
    }

    /// <summary>
    /// Sends a WebSocket close frame.
    /// </summary>
    public async Task SendCloseAsync(ClientClosedReason reason, CancellationToken ct = default)
    {
        var status = WsFrameWriter.MapCloseStatus(reason);
        var closePayload = WsFrameWriter.CreateCloseMessage(status, reason.ToReasonString());
        var frame = WsFrameWriter.BuildControlFrame(WsConstants.CloseMessage, closePayload, _maskWrite);
        await _inner.WriteAsync(frame, ct);
        await _inner.FlushAsync(ct);
    }

    private int DrainReadQueue(Span<byte> buffer)
    {
        int written = 0;
        while (_readQueue.Count > 0 && written < buffer.Length)
        {
            var current = _readQueue.Peek();
            int available = current.Length - _readOffset;
            int toCopy = Math.Min(available, buffer.Length - written);
            current.AsSpan(_readOffset, toCopy).CopyTo(buffer[written..]);
            written += toCopy;
            _readOffset += toCopy;
            if (_readOffset >= current.Length)
            {
                _readQueue.Dequeue();
                _readOffset = 0;
            }
        }
        return written;
    }

    // Stream abstract members
    public override bool CanRead => true;
    public override bool CanWrite => true;
    public override bool CanSeek => false;
    public override long Length => throw new NotSupportedException();
    public override long Position { get => throw new NotSupportedException(); set => throw new NotSupportedException(); }
    public override void Flush() => _inner.Flush();
    public override Task FlushAsync(CancellationToken ct) => _inner.FlushAsync(ct);
    public override int Read(byte[] buffer, int offset, int count) => throw new NotSupportedException("Use ReadAsync");
    public override void Write(byte[] buffer, int offset, int count) => throw new NotSupportedException("Use WriteAsync");
    public override long Seek(long offset, SeekOrigin origin) => throw new NotSupportedException();
    public override void SetLength(long value) => throw new NotSupportedException();

    protected override void Dispose(bool disposing)
    {
        if (disposing)
            _inner.Dispose();
        base.Dispose(disposing);
    }

    public override async ValueTask DisposeAsync()
    {
        await _inner.DisposeAsync();
        GC.SuppressFinalize(this);
    }
}
