using System.Buffers;
using System.IO.Pipelines;
using System.Net;
using System.Net.Sockets;
using System.Security.Cryptography;
using System.Text;
using System.Text.Json;
using System.Threading.Channels;
using Microsoft.Extensions.Logging;
using NATS.Server.Auth;
using NATS.Server.Auth.Jwt;
using NATS.Server.IO;
using NATS.Server.JetStream.Publish;
using NATS.Server.Protocol;
using NATS.Server.Subscriptions;
using NATS.Server.Tls;
using NATS.Server.WebSocket;

namespace NATS.Server;

public interface IMessageRouter
{
    void ProcessMessage(string subject, string? replyTo, ReadOnlyMemory<byte> headers,
        ReadOnlyMemory<byte> payload, INatsClient sender);
    void RemoveClient(INatsClient client);
    void PublishConnectEvent(INatsClient client);
    void PublishDisconnectEvent(INatsClient client);
}

public interface ISubListAccess
{
    SubList SubList { get; }
}

/// <summary>
/// Outbound data wrapper that tracks an optional pooled buffer for return after writing.
/// Go reference: client.go — queueOutbound appends slice refs to pooled buffers; writeLoop returns them.
/// </summary>
internal readonly struct OutboundData
{
    public readonly ReadOnlyMemory<byte> Data;
    public readonly byte[]? PoolBuffer;

    public OutboundData(ReadOnlyMemory<byte> data, byte[]? poolBuffer = null)
    {
        Data = data;
        PoolBuffer = poolBuffer;
    }

    public int Length => Data.Length;
}

public sealed class NatsClient : INatsClient, IDisposable
{
    private static readonly ClientCommandMatrix CommandMatrix = new();
    private readonly Socket _socket;
    private readonly Stream _stream;
    private readonly NatsOptions _options;
    private readonly ServerInfo _serverInfo;
    private readonly AuthService _authService;
    private readonly byte[]? _nonce;
    private readonly NatsParser _parser;
    private readonly AdaptiveReadBuffer _adaptiveReadBuffer = new();
    private readonly OutboundBufferPool _outboundBufferPool = new();
    private readonly Channel<OutboundData> _outbound = Channel.CreateBounded<OutboundData>(
        new BoundedChannelOptions(8192) { SingleReader = true, FullMode = BoundedChannelFullMode.Wait });
    private long _pendingBytes;

    /// <summary>
    /// True when the connection uses a plain TCP socket (no TLS/WebSocket) and we can use
    /// Socket.SendAsync scatter-gather directly, bypassing NetworkStream overhead.
    /// Go reference: client.go flushOutbound — net.Buffers.WriteTo for writev().
    /// </summary>
    private readonly bool _isPlainSocket;

    // Per-client direct write buffer for the hot message delivery path.
    // Go reference: client.go outbound.nb — pooled buffers filled by queueOutbound under lock.
    // Publishers format MSG directly into this buffer, avoiding Channel + pool overhead.
    // Double-buffer: _directBuf is filled by producers, _writeBuf is owned by the write loop.
    // On flush, they are swapped under lock — zero allocation.
    private byte[] _directBuf = new byte[65536];
    private byte[] _writeBuf = new byte[65536];
    private int _directBufUsed;
    private SpinLock _directBufLock = new(enableThreadOwnerTracking: false);
    private readonly Channel<byte> _flushSignal = Channel.CreateBounded<byte>(
        new BoundedChannelOptions(1) { SingleReader = true, FullMode = BoundedChannelFullMode.DropWrite });
    private CancellationTokenSource? _clientCts;
    private readonly Dictionary<string, Subscription> _subs = new();
    private readonly ILogger _logger;
    private ClientPermissions? _permissions;
    private readonly ServerStats _serverStats;

    public ulong Id { get; }
    public ClientKind Kind { get; }
    public ClientOptions? ClientOpts { get; private set; }
    public MessageTraceContext TraceContext { get; private set; } = MessageTraceContext.Empty;
    public IMessageRouter? Router { get; set; }
    public Account? Account { get; private set; }
    public ClientPermissions? Permissions => _permissions;

    /// <summary>
    /// MQTT client-id for monitoring (/connz mqtt_client field).
    /// Set when this NatsClient proxies an MQTT connection via MqttNatsClientAdapter.
    /// </summary>
    public string? MqttClientId { get; set; }

    private readonly ClientFlagHolder _flags = new();
    public bool ConnectReceived => _flags.HasFlag(ClientFlags.ConnectReceived);
    public ClientClosedReason CloseReason { get; private set; }

    public void SetTraceMode(bool enabled)
    {
        if (enabled)
        {
            _flags.SetFlag(ClientFlags.TraceMode);
            _parser.Logger = _logger;
        }
        else
        {
            _flags.ClearFlag(ClientFlags.TraceMode);
            _parser.Logger = _options.Trace ? _logger : null;
        }
    }

    public DateTime StartTime { get; }
    private long _lastActivityTicks;
    public DateTime LastActivity => new(Interlocked.Read(ref _lastActivityTicks), DateTimeKind.Utc);
    public string? RemoteIp { get; }
    public int RemotePort { get; }

    // Stats
    public long InMsgs;
    public long OutMsgs;
    public long InBytes;
    public long OutBytes;

    // Close reason tracking
    private int _skipFlushOnClose;
    public bool ShouldSkipFlush => Volatile.Read(ref _skipFlushOnClose) != 0;

    // PING keepalive state
    private int _pingsOut;
    private long _lastIn;

    // RTT tracking
    private long _rttStartTicks;
    private long _rtt;
    public TimeSpan Rtt => new(Interlocked.Read(ref _rtt));

    public bool IsMqtt { get; set; }
    public bool IsWebSocket { get; set; }
    public WsUpgradeResult? WsInfo { get; set; }

    public TlsConnectionState? TlsState { get; set; }
    public bool InfoAlreadySent { get; set; }

    public IReadOnlyDictionary<string, Subscription> Subscriptions => _subs;
    public PubAck? LastJetStreamPubAck { get; private set; }

    public NatsClient(ulong id, Stream stream, Socket socket, NatsOptions options, ServerInfo serverInfo,
        AuthService authService, byte[]? nonce, ILogger logger, ServerStats serverStats,
        ClientKind kind = ClientKind.Client)
    {
        Id = id;
        Kind = kind;
        _socket = socket;
        _stream = stream;
        _options = options;
        _serverInfo = serverInfo;
        _authService = authService;
        _nonce = nonce;
        _logger = logger;
        _serverStats = serverStats;
        _parser = new NatsParser(options.MaxPayload, options.Trace ? logger : null);
        _isPlainSocket = stream is System.Net.Sockets.NetworkStream;
        StartTime = DateTime.UtcNow;
        _lastActivityTicks = StartTime.Ticks;
        if (socket.RemoteEndPoint is IPEndPoint ep)
        {
            RemoteIp = ep.Address.ToString();
            RemotePort = ep.Port;
        }
    }

    public byte[]? GetNonce() => _nonce?.ToArray();

    public string GetName() => ClientOpts?.Name ?? string.Empty;

    public ClientConnectionType ClientType()
    {
        if (Kind != ClientKind.Client)
            return ClientConnectionType.NonClient;
        if (IsMqtt)
            return ClientConnectionType.Mqtt;
        if (IsWebSocket)
            return ClientConnectionType.WebSocket;
        return ClientConnectionType.Nats;
    }

    public override string ToString()
    {
        var endpoint = RemoteIp is null ? "unknown" : $"{RemoteIp}:{RemotePort}";
        return $"{Kind} cid={Id} endpoint={endpoint}";
    }

    public bool QueueOutbound(ReadOnlyMemory<byte> data) => QueueOutboundCore(new OutboundData(data));

    /// <summary>
    /// Queues a pooled buffer for outbound writing. The write loop will return the buffer
    /// to the pool after writing. Avoids the .ToArray() copy on the hot path.
    /// Go reference: client.go queueOutbound — queues slice refs into pooled buffers.
    /// </summary>
    private bool QueueOutboundPooled(byte[] buffer, int length)
        => QueueOutboundCore(new OutboundData(new ReadOnlyMemory<byte>(buffer, 0, length), buffer));

    private bool QueueOutboundCore(OutboundData data)
    {
        if (_flags.HasFlag(ClientFlags.CloseConnection))
            return false;

        var pending = Interlocked.Add(ref _pendingBytes, data.Length);
        if (pending > _options.MaxPending)
        {
            Interlocked.Add(ref _pendingBytes, -data.Length);
            if (!_flags.HasFlag(ClientFlags.CloseConnection))
            {
                _flags.SetFlag(ClientFlags.CloseConnection);
                _flags.SetFlag(ClientFlags.IsSlowConsumer);
                Interlocked.Increment(ref _serverStats.SlowConsumers);
                Interlocked.Increment(ref _serverStats.SlowConsumerClients);
                _ = CloseWithReasonAsync(ClientClosedReason.SlowConsumerPendingBytes, NatsProtocol.ErrSlowConsumer);
            }
            return false;
        }

        if (!_outbound.Writer.TryWrite(data))
        {
            Interlocked.Add(ref _pendingBytes, -data.Length);
            if (!_flags.HasFlag(ClientFlags.CloseConnection))
            {
                _flags.SetFlag(ClientFlags.CloseConnection);
                _flags.SetFlag(ClientFlags.IsSlowConsumer);
                Interlocked.Increment(ref _serverStats.SlowConsumers);
                Interlocked.Increment(ref _serverStats.SlowConsumerClients);
                _ = CloseWithReasonAsync(ClientClosedReason.SlowConsumerPendingBytes, NatsProtocol.ErrSlowConsumer);
            }
            return false;
        }

        SignalFlushPending();
        _flushSignal.Writer.TryWrite(0);
        return true;
    }

    public long PendingBytes => Interlocked.Read(ref _pendingBytes);

    /// <summary>
    /// Maximum number of pending flush signals before forcing a flush.
    /// Go reference: server/client.go (maxFlushPending, pcd)
    /// </summary>
    public const int MaxFlushPending = 10;

    /// <summary>
    /// Current pending flush signal count. When the write loop drains queued data
    /// and _flushSignalsPending is below MaxFlushPending, it can briefly coalesce
    /// additional writes before flushing to reduce syscalls.
    /// </summary>
    private int _flushSignalsPending;

    /// <summary>
    /// Records that a flush signal has been posted. Called after each QueueOutbound write.
    /// Go reference: server/client.go pcd (post-channel-data) flush signaling.
    /// </summary>
    public void SignalFlushPending()
    {
        Interlocked.Increment(ref _flushSignalsPending);
    }

    /// <summary>
    /// Resets the flush signal counter after a flush completes.
    /// </summary>
    public void ResetFlushPending()
    {
        Interlocked.Exchange(ref _flushSignalsPending, 0);
    }

    /// <summary>
    /// Current number of pending flush signals.
    /// </summary>
    public int FlushSignalsPending => Volatile.Read(ref _flushSignalsPending);

    /// <summary>
    /// Whether more writes should be coalesced before flushing.
    /// Returns true when pending flush signals are below MaxFlushPending,
    /// indicating the write loop may briefly wait for more data.
    /// Go reference: server/client.go — fsp (flush signal pending) check.
    /// </summary>
    public bool ShouldCoalesceFlush => FlushSignalsPending < MaxFlushPending;

    public async Task RunAsync(CancellationToken ct)
    {
        _clientCts = CancellationTokenSource.CreateLinkedTokenSource(ct);
        Interlocked.Exchange(ref _lastIn, Environment.TickCount64);
        var pipe = new Pipe();
        try
        {
            // Send INFO (skip if already sent during TLS negotiation)
            if (!InfoAlreadySent)
                SendInfo();

            // Start auth timeout if auth is required
            Task? authTimeoutTask = null;
            if (_authService.IsAuthRequired)
            {
                authTimeoutTask = Task.Run(async () =>
                {
                    try
                    {
                        await Task.Delay(_options.AuthTimeout, _clientCts!.Token);
                        if (!ConnectReceived)
                        {
                            _logger.LogDebug("Client {ClientId} auth timeout", Id);
                            await SendErrAndCloseAsync(NatsProtocol.ErrAuthTimeout, ClientClosedReason.AuthenticationTimeout);
                        }
                    }
                    catch (OperationCanceledException)
                    {
                        // Normal -- client connected or was cancelled
                    }
                }, _clientCts.Token);
            }

            // Start read pump, command processing, write loop, and ping timer in parallel
            var fillTask = FillPipeAsync(pipe.Writer, _clientCts.Token);
            var processTask = ProcessCommandsAsync(pipe.Reader, _clientCts.Token);
            var pingTask = RunPingTimerAsync(_clientCts.Token);
            var writeTask = RunWriteLoopAsync(_clientCts.Token);

            if (authTimeoutTask != null)
                await Task.WhenAny(fillTask, processTask, pingTask, writeTask, authTimeoutTask);
            else
                await Task.WhenAny(fillTask, processTask, pingTask, writeTask);
        }
        catch (OperationCanceledException)
        {
            _logger.LogDebug("Client {ClientId} operation cancelled", Id);
            MarkClosed(ClientClosedReason.ServerShutdown);
        }
        catch (IOException)
        {
            MarkClosed(ClientClosedReason.ReadError);
        }
        catch (SocketException)
        {
            MarkClosed(ClientClosedReason.ReadError);
        }
        catch (Exception ex)
        {
            _logger.LogDebug(ex, "Client {ClientId} connection error", Id);
            MarkClosed(ClientClosedReason.ReadError);
        }
        finally
        {
            MarkClosed(ClientClosedReason.ClientClosed);
            _outbound.Writer.TryComplete();
            _flushSignal.Writer.TryComplete();
            try { _socket.Shutdown(SocketShutdown.Both); }
            catch (SocketException) { }
            catch (ObjectDisposedException) { }
            Router?.RemoveClient(this);
        }
    }

    private async Task FillPipeAsync(PipeWriter writer, CancellationToken ct)
    {
        try
        {
            while (!ct.IsCancellationRequested)
            {
                var memory = writer.GetMemory(_adaptiveReadBuffer.CurrentSize);
                int bytesRead = await _stream.ReadAsync(memory, ct);
                if (bytesRead == 0)
                    break;

                _adaptiveReadBuffer.RecordRead(bytesRead);
                writer.Advance(bytesRead);
                var result = await writer.FlushAsync(ct);
                if (result.IsCompleted)
                    break;
            }
        }
        finally
        {
            await writer.CompleteAsync();
        }
    }

    private async Task ProcessCommandsAsync(PipeReader reader, CancellationToken ct)
    {
        try
        {
            while (!ct.IsCancellationRequested)
            {
                var result = await reader.ReadAsync(ct);
                var buffer = result.Buffer;

                long localInMsgs = 0;
                long localInBytes = 0;

                while (_parser.TryParseView(ref buffer, out var cmdView))
                {
                    Interlocked.Exchange(ref _lastIn, Environment.TickCount64);

                    // Handle Pub/HPub inline to allow ref parameter passing for stat batching.
                    // DispatchCommandAsync is async and cannot accept ref parameters.
                    if (cmdView.Type is CommandType.Pub or CommandType.HPub
                        && (!_authService.IsAuthRequired || ConnectReceived))
                    {
                        Interlocked.Exchange(ref _lastActivityTicks, DateTime.UtcNow.Ticks);
                        ProcessPub(cmdView, ref localInMsgs, ref localInBytes);
                        if (ClientOpts?.Verbose == true)
                            WriteProtocol(NatsProtocol.OkBytes);
                    }
                    else
                    {
                        await DispatchCommandAsync(cmdView.Materialize(), ct);
                    }
                }

                if (localInMsgs > 0)
                {
                    Interlocked.Add(ref InMsgs, localInMsgs);
                    Interlocked.Add(ref _serverStats.InMsgs, localInMsgs);
                }

                if (localInBytes > 0)
                {
                    Interlocked.Add(ref InBytes, localInBytes);
                    Interlocked.Add(ref _serverStats.InBytes, localInBytes);
                }

                reader.AdvanceTo(buffer.Start, buffer.End);

                if (result.IsCompleted)
                    break;
            }
        }
        finally
        {
            await reader.CompleteAsync();
        }
    }

    private async ValueTask DispatchCommandAsync(ParsedCommand cmd, CancellationToken ct)
    {
        Interlocked.Exchange(ref _lastActivityTicks, DateTime.UtcNow.Ticks);

        if (!CommandMatrix.IsAllowed(Kind, cmd.Operation))
        {
            _logger.LogDebug("Command {Command} is not allowed for client kind {ClientKind}", cmd.Operation, Kind);
            await SendErrAndCloseAsync("Parser Error");
            return;
        }

        // If auth is required and CONNECT hasn't been received yet,
        // only allow CONNECT and PING commands
        if (_authService.IsAuthRequired && !ConnectReceived)
        {
            switch (cmd.Type)
            {
                case CommandType.Connect:
                    await ProcessConnectAsync(cmd);
                    return;
                case CommandType.Ping:
                    WriteProtocol(NatsProtocol.PongBytes);
                    return;
                default:
                    // Ignore all other commands until authenticated
                    return;
            }
        }

        switch (cmd.Type)
        {
            case CommandType.Connect:
                await ProcessConnectAsync(cmd);
                if (ClientOpts?.Verbose == true)
                    WriteProtocol(NatsProtocol.OkBytes);
                break;

            case CommandType.Ping:
                WriteProtocol(NatsProtocol.PongBytes);
                if (ClientOpts?.Verbose == true)
                    WriteProtocol(NatsProtocol.OkBytes);
                break;

            case CommandType.Pong:
                Interlocked.Exchange(ref _pingsOut, 0);
                var rttStart = Interlocked.Read(ref _rttStartTicks);
                if (rttStart > 0)
                {
                    var elapsed = DateTime.UtcNow.Ticks - rttStart;
                    if (elapsed <= 0) elapsed = 1; // min 1 tick for Windows granularity
                    Interlocked.Exchange(ref _rtt, elapsed);
                }
                _flags.SetFlag(ClientFlags.FirstPongSent);
                break;

            case CommandType.Sub:
                ProcessSub(cmd);
                if (ClientOpts?.Verbose == true)
                    WriteProtocol(NatsProtocol.OkBytes);
                break;

            case CommandType.Unsub:
                ProcessUnsub(cmd);
                if (ClientOpts?.Verbose == true)
                    WriteProtocol(NatsProtocol.OkBytes);
                break;

            case CommandType.Pub:
            case CommandType.HPub:
                // Pub/HPub is handled inline in ProcessCommandsAsync for stat batching
                break;
        }
    }

    private async ValueTask ProcessConnectAsync(ParsedCommand cmd)
    {
        ClientOpts = JsonSerializer.Deserialize<ClientOptions>(cmd.Payload.Span)
                     ?? new ClientOptions();
        TraceContext = MessageTraceContext.CreateFromConnect(ClientOpts);

        // Authenticate if auth is required
        AuthResult? authResult = null;
        if (_authService.IsAuthRequired)
        {
            var context = new ClientAuthContext
            {
                Opts = ClientOpts,
                Nonce = _nonce ?? [],
                ClientCertificate = TlsState?.PeerCert,
                ConnectionType = JwtConnectionTypes.Standard,
            };

            authResult = _authService.Authenticate(context);
            if (authResult == null)
            {
                _logger.LogWarning("Client {ClientId} authentication failed", Id);
                await SendErrAndCloseAsync(NatsProtocol.ErrAuthorizationViolation, ClientClosedReason.AuthenticationViolation);
                return;
            }

            // Build permissions from auth result
            _permissions = ClientPermissions.Build(authResult.Permissions);

            // Resolve account
            if (Router is NatsServer server)
            {
                var accountName = authResult.AccountName ?? Account.GlobalAccountName;
                Account = server.GetOrCreateAccount(accountName);
                if (authResult.MaxJetStreamStreams > 0)
                    Account.MaxJetStreamStreams = authResult.MaxJetStreamStreams;
                if (!string.IsNullOrWhiteSpace(authResult.JetStreamTier))
                    Account.JetStreamTier = authResult.JetStreamTier;
                if (!Account.AddClient(Id))
                {
                    Account = null;
                    await SendErrAndCloseAsync("maximum connections for account exceeded",
                        ClientClosedReason.AuthenticationViolation);
                    return;
                }
            }

            _logger.LogDebug("Client {ClientId} authenticated as {Identity}", Id, authResult.Identity);

            // Clear nonce after use -- defense-in-depth against memory dumps
            if (_nonce != null)
                CryptographicOperations.ZeroMemory(_nonce);
        }

        // If no account was assigned by auth, assign global account
        if (Account == null && Router is NatsServer server2)
        {
            Account = server2.GetOrCreateAccount(Account.GlobalAccountName);
            if (!Account.AddClient(Id))
            {
                Account = null;
                await SendErrAndCloseAsync("maximum connections for account exceeded",
                    ClientClosedReason.AuthenticationViolation);
                return;
            }
        }

        // Validate no_responders requires headers
        if (ClientOpts.NoResponders && !ClientOpts.Headers)
        {
            _logger.LogDebug("Client {ClientId} no_responders requires headers", Id);
            await CloseWithReasonAsync(ClientClosedReason.NoRespondersRequiresHeaders,
                NatsProtocol.ErrNoRespondersRequiresHeaders);
            return;
        }

        _flags.SetFlag(ClientFlags.ConnectReceived);
        _flags.SetFlag(ClientFlags.ConnectProcessFinished);
        _logger.LogDebug("CONNECT received from client {ClientId}, name={ClientName}", Id, ClientOpts?.Name);

        // Publish connect advisory to the system event bus
        Router?.PublishConnectEvent(this);

        // Start auth expiry timer if needed
        if (_authService.IsAuthRequired && authResult?.Expiry is { } expiry)
        {
            var remaining = expiry - DateTimeOffset.UtcNow;
            if (remaining > TimeSpan.Zero)
            {
                _ = Task.Run(async () =>
                {
                    try
                    {
                        await Task.Delay(remaining, _clientCts!.Token);
                        _logger.LogDebug("Client {ClientId} authentication expired", Id);
                        await SendErrAndCloseAsync("Authentication Expired",
                            ClientClosedReason.AuthenticationExpired);
                    }
                    catch (OperationCanceledException) { }
                }, _clientCts!.Token);
            }
            else
            {
                await SendErrAndCloseAsync("Authentication Expired",
                    ClientClosedReason.AuthenticationExpired);
                return;
            }
        }
    }

    private void ProcessSub(ParsedCommand cmd)
    {
        // Permission check for subscribe
        if (_permissions != null && !_permissions.IsSubscribeAllowed(cmd.Subject!, cmd.Queue))
        {
            _logger.LogDebug("Client {ClientId} subscribe permission denied for {Subject}", Id, cmd.Subject);
            SendErr(NatsProtocol.ErrPermissionsSubscribe);
            return;
        }

        // Per-connection subscription limit
        if (_options.MaxSubs > 0 && _subs.Count >= _options.MaxSubs)
        {
            _logger.LogDebug("Client {ClientId} max subscriptions exceeded", Id);
            _ = SendErrAndCloseAsync(NatsProtocol.ErrMaxSubscriptionsExceeded,
                ClientClosedReason.MaxSubscriptionsExceeded);
            return;
        }

        // Per-account subscription limit
        if (Account != null && !Account.IncrementSubscriptions())
        {
            _logger.LogDebug("Client {ClientId} account subscription limit exceeded", Id);
            _ = SendErrAndCloseAsync(NatsProtocol.ErrMaxSubscriptionsExceeded,
                ClientClosedReason.MaxSubscriptionsExceeded);
            return;
        }

        var sub = new Subscription
        {
            Subject = cmd.Subject!,
            Queue = cmd.Queue,
            Sid = cmd.Sid!,
        };

        _subs[cmd.Sid!] = sub;
        sub.Client = this;

        _logger.LogDebug("SUB {Subject} {Sid} from client {ClientId}", cmd.Subject, cmd.Sid, Id);

        Account?.SubList.Insert(sub);
        if (Router is NatsServer server)
            server.OnLocalSubscription(Account?.Name ?? Account.GlobalAccountName, sub.Subject, sub.Queue);
    }

    private void ProcessUnsub(ParsedCommand cmd)
    {
        _logger.LogDebug("UNSUB {Sid} from client {ClientId}", cmd.Sid, Id);

        if (!_subs.TryGetValue(cmd.Sid!, out var sub))
            return;

        if (cmd.MaxMessages > 0)
        {
            sub.MaxMessages = cmd.MaxMessages;
            // Will be cleaned up when MessageCount reaches MaxMessages
            return;
        }

        _subs.Remove(cmd.Sid!);
        Account?.DecrementSubscriptions();

        Account?.SubList.Remove(sub);
        if (Router is NatsServer server)
            server.OnLocalUnsubscription(Account?.Name ?? Account.GlobalAccountName, sub.Subject, sub.Queue);
    }

    private void ProcessPub(ParsedCommandView cmd, ref long localInMsgs, ref long localInBytes)
    {
        var payloadMemory = cmd.GetPayloadMemory();
        localInMsgs++;
        localInBytes += payloadMemory.Length;

        // Max payload validation (always, hard close)
        if (payloadMemory.Length > _options.MaxPayload)
        {
            _logger.LogWarning("Client {ClientId} exceeded max payload: {Size} > {MaxPayload}",
                Id, payloadMemory.Length, _options.MaxPayload);
            _ = SendErrAndCloseAsync(NatsProtocol.ErrMaxPayloadViolation, ClientClosedReason.MaxPayloadExceeded);
            return;
        }

        var subject = Encoding.ASCII.GetString(cmd.Subject.Span);

        // Pedantic mode: validate publish subject
        if (ClientOpts?.Pedantic == true && !SubjectMatch.IsValidPublishSubject(subject))
        {
            _logger.LogDebug("Client {ClientId} invalid publish subject: {Subject}", Id, subject);
            SendErr(NatsProtocol.ErrInvalidPublishSubject);
            return;
        }

        // Permission check for publish
        if (_permissions != null && !_permissions.IsPublishAllowed(subject))
        {
            _logger.LogDebug("Client {ClientId} publish permission denied for {Subject}", Id, subject);
            SendErr(NatsProtocol.ErrPermissionsPublish);
            return;
        }

        ReadOnlyMemory<byte> headers = default;
        ReadOnlyMemory<byte> payload = payloadMemory;

        if (cmd.Type == CommandType.HPub && cmd.HeaderSize > 0)
        {
            headers = payloadMemory[..cmd.HeaderSize];
            payload = payloadMemory[cmd.HeaderSize..];
        }

        var replyTo = cmd.ReplyTo.IsEmpty ? null : Encoding.ASCII.GetString(cmd.ReplyTo.Span);
        Router?.ProcessMessage(subject, replyTo, headers, payload, this);
    }

    public void RecordJetStreamPubAck(PubAck ack)
    {
        LastJetStreamPubAck = ack;
    }

    private void SendInfo()
    {
        // Use the cached INFO bytes from the server when there is no per-connection
        // nonce (i.e. NKey auth is not active for this connection).  When a nonce is
        // present the _serverInfo was already cloned with the nonce embedded, so we
        // must serialise it individually.
        if (_nonce == null && Router is NatsServer server)
        {
            QueueOutbound(server.CachedInfoLine);
        }
        else
        {
            var infoJson = JsonSerializer.Serialize(_serverInfo);
            var infoLine = Encoding.ASCII.GetBytes($"INFO {infoJson}\r\n");
            QueueOutbound(infoLine);
        }
    }

    public void SendMessage(string subject, string sid, string? replyTo,
        ReadOnlyMemory<byte> headers, ReadOnlyMemory<byte> payload)
    {
        SendMessageNoFlush(subject, sid, replyTo, headers, payload);
        _flushSignal.Writer.TryWrite(0);
    }

    /// <summary>
    /// Queues a message into the client's outbound buffer without signaling the write loop.
    /// Callers must call <see cref="SignalFlush"/> after all messages in a batch are queued.
    /// Go reference: client.go addToPCD — deferred flush via pcd map.
    /// </summary>
    public void SendMessageNoFlush(string subject, string sid, string? replyTo,
        ReadOnlyMemory<byte> headers, ReadOnlyMemory<byte> payload)
    {
        // Batch per-client stats (single thread writes these during delivery).
        // Server-wide stats use Interlocked since multiple threads update them.
        OutMsgs++;
        OutBytes += payload.Length + headers.Length;
        Interlocked.Increment(ref _serverStats.OutMsgs);
        Interlocked.Add(ref _serverStats.OutBytes, payload.Length + headers.Length);

        // Format MSG header on the stack (no heap allocation).
        // Go reference: client.go msgHeader — formats into per-client 1KB scratch buffer (msgb).
        Span<byte> headerBuf = stackalloc byte[512];
        int pos = 0;

        // Write prefix
        if (headers.Length > 0)
        {
            "HMSG "u8.CopyTo(headerBuf);
            pos = 5;
        }
        else
        {
            "MSG "u8.CopyTo(headerBuf);
            pos = 4;
        }

        // Subject
        pos += Encoding.ASCII.GetBytes(subject, headerBuf[pos..]);
        headerBuf[pos++] = (byte)' ';

        // SID
        pos += Encoding.ASCII.GetBytes(sid, headerBuf[pos..]);
        headerBuf[pos++] = (byte)' ';

        // Reply-to
        if (replyTo != null)
        {
            pos += Encoding.ASCII.GetBytes(replyTo, headerBuf[pos..]);
            headerBuf[pos++] = (byte)' ';
        }

        // Sizes
        if (headers.Length > 0)
        {
            int totalSize = headers.Length + payload.Length;
            headers.Length.TryFormat(headerBuf[pos..], out int written);
            pos += written;
            headerBuf[pos++] = (byte)' ';
            totalSize.TryFormat(headerBuf[pos..], out written);
            pos += written;
        }
        else
        {
            payload.Length.TryFormat(headerBuf[pos..], out int written);
            pos += written;
        }

        // CRLF
        headerBuf[pos++] = (byte)'\r';
        headerBuf[pos++] = (byte)'\n';

        // Write header + body + CRLF directly into the per-client buffer under lock.
        // Go reference: client.go queueOutbound — appends slice refs under client.mu.
        var totalLen = pos + headers.Length + payload.Length + 2;
        var lockTaken = false;
        _directBufLock.Enter(ref lockTaken);
        try
        {
            // Grow buffer if needed
            var needed = _directBufUsed + totalLen;
            if (needed > _directBuf.Length)
            {
                var newSize = Math.Max(_directBuf.Length * 2, needed);
                var newBuf = new byte[newSize];
                _directBuf.AsSpan(0, _directBufUsed).CopyTo(newBuf);
                _directBuf = newBuf;
            }

            var dst = _directBuf.AsSpan(_directBufUsed);

            // Header
            headerBuf[..pos].CopyTo(dst);
            _directBufUsed += pos;
            dst = _directBuf.AsSpan(_directBufUsed);

            // Headers (HMSG)
            if (headers.Length > 0)
            {
                headers.Span.CopyTo(dst);
                _directBufUsed += headers.Length;
                dst = _directBuf.AsSpan(_directBufUsed);
            }

            // Payload
            if (payload.Length > 0)
            {
                payload.Span.CopyTo(dst);
                _directBufUsed += payload.Length;
            }

            // Trailing CRLF
            _directBuf[_directBufUsed++] = (byte)'\r';
            _directBuf[_directBufUsed++] = (byte)'\n';
        }
        finally
        {
            if (lockTaken) _directBufLock.Exit();
        }

        var pending = Interlocked.Add(ref _pendingBytes, totalLen);
        if (pending > _options.MaxPending)
        {
            if (!_flags.HasFlag(ClientFlags.CloseConnection))
            {
                _flags.SetFlag(ClientFlags.CloseConnection);
                _flags.SetFlag(ClientFlags.IsSlowConsumer);
                Interlocked.Increment(ref _serverStats.SlowConsumers);
                Interlocked.Increment(ref _serverStats.SlowConsumerClients);
                _ = CloseWithReasonAsync(ClientClosedReason.SlowConsumerPendingBytes, NatsProtocol.ErrSlowConsumer);
            }
        }
    }

    /// <summary>
    /// Signals the write loop to flush buffered data. Call once after batching
    /// multiple <see cref="SendMessageNoFlush"/> calls to the same client.
    /// Go reference: client.go:1324 flushClients — one signal per unique client.
    /// </summary>
    public void SignalFlush() => _flushSignal.Writer.TryWrite(0);

    private void WriteProtocol(byte[] data)
    {
        QueueOutbound(data);
    }

    public void SendErr(string message)
    {
        var errLine = Encoding.ASCII.GetBytes($"-ERR '{message}'\r\n");
        QueueOutbound(errLine);
    }

    private async Task RunWriteLoopAsync(CancellationToken ct)
    {
        _flags.SetFlag(ClientFlags.WriteLoopStarted);
        var outboundReader = _outbound.Reader;
        var flushReader = _flushSignal.Reader;

        // Scatter-gather list for Socket.SendAsync (plaintext path).
        // Go reference: client.go flushOutbound — net.Buffers.WriteTo → writev().
        var segments = new List<ArraySegment<byte>>(64);
        var returnList = new List<byte[]>(32);

        // Coalesce buffer for TLS/WebSocket path (SslStream doesn't support scatter-gather).
        const int coalesceBufSize = 65536;
        byte[]? coalesceBuf = _isPlainSocket ? null : new byte[coalesceBufSize];

        try
        {
            // Wait on _flushSignal — signaled by both SendMessage (direct buffer)
            // and QueueOutboundCore (protocol channel). Single wait point, two data sources.
            while (await flushReader.WaitToReadAsync(ct))
            {
                // Drain all pending signals
                while (flushReader.TryRead(out _)) { }

                long batchBytes = 0;

                // 1. Drain direct buffer (message data from SendMessage).
                // Swap _directBuf and _writeBuf under SpinLock — zero allocation.
                int directLen = 0;
                var lockTaken = false;
                _directBufLock.Enter(ref lockTaken);
                try
                {
                    if (_directBufUsed > 0)
                    {
                        (_directBuf, _writeBuf) = (_writeBuf, _directBuf);
                        directLen = _directBufUsed;
                        _directBufUsed = 0;
                    }
                }
                finally
                {
                    if (lockTaken) _directBufLock.Exit();
                }

                if (_isPlainSocket)
                {
                    // Plaintext path: collect ArraySegments for scatter-gather write.
                    if (directLen > 0)
                    {
                        segments.Add(new ArraySegment<byte>(_writeBuf, 0, directLen));
                        batchBytes += directLen;
                    }

                    // 2. Drain outbound channel (protocol data: INFO, PING, ERR, etc.)
                    while (outboundReader.TryRead(out var item))
                    {
                        batchBytes += item.Data.Length;
                        if (System.Runtime.InteropServices.MemoryMarshal.TryGetArray(item.Data, out var seg))
                            segments.Add(seg);
                        else
                            segments.Add(new ArraySegment<byte>(item.Data.ToArray()));
                        if (item.PoolBuffer != null)
                            returnList.Add(item.PoolBuffer);
                    }

                    if (segments.Count > 0)
                    {
                        await _socket.SendAsync(segments, SocketFlags.None);
                        segments.Clear();
                    }
                }
                else
                {
                    // TLS/WebSocket path: coalesce into single buffer.
                    var coalescePos = 0;

                    // Direct buffer data first
                    if (directLen > 0)
                    {
                        if (directLen <= coalesceBufSize)
                        {
                            _writeBuf.AsSpan(0, directLen).CopyTo(coalesceBuf);
                            coalescePos = directLen;
                        }
                        else
                        {
                            await _stream.WriteAsync(new ReadOnlyMemory<byte>(_writeBuf, 0, directLen), ct);
                        }

                        batchBytes += directLen;
                    }

                    // Outbound channel data
                    while (outboundReader.TryRead(out var item))
                    {
                        var data = item.Data;
                        batchBytes += data.Length;

                        if (coalescePos + data.Length <= coalesceBufSize)
                        {
                            data.Span.CopyTo(coalesceBuf.AsSpan(coalescePos));
                            coalescePos += data.Length;
                        }
                        else
                        {
                            if (coalescePos > 0)
                            {
                                await _stream.WriteAsync(coalesceBuf.AsMemory(0, coalescePos), ct);
                                coalescePos = 0;
                            }

                            await _stream.WriteAsync(data, ct);
                        }

                        if (item.PoolBuffer != null)
                            returnList.Add(item.PoolBuffer);
                    }

                    if (coalescePos > 0)
                        await _stream.WriteAsync(coalesceBuf.AsMemory(0, coalescePos), ct);

                    // SslStream requires explicit flush to push encrypted data.
                    await _stream.FlushAsync(ct);
                }

                // Return pooled buffers after writing
                foreach (var buf in returnList)
                    _outboundBufferPool.ReturnBuffer(buf);
                returnList.Clear();

                ResetFlushPending();
                Interlocked.Add(ref _pendingBytes, -batchBytes);
            }
        }
        catch (OperationCanceledException)
        {
            // Normal shutdown
        }
        catch (IOException)
        {
            await CloseWithReasonAsync(ClientClosedReason.WriteError);
        }
        catch (SocketException)
        {
            await CloseWithReasonAsync(ClientClosedReason.WriteError);
        }
    }

    public async Task SendErrAndCloseAsync(string message, ClientClosedReason reason = ClientClosedReason.ProtocolViolation)
    {
        await CloseWithReasonAsync(reason, message);
    }

    private async Task CloseWithReasonAsync(ClientClosedReason reason, string? errMessage = null)
    {
        CloseReason = reason;
        if (errMessage != null)
            SendErr(errMessage);
        _flags.SetFlag(ClientFlags.CloseConnection);

        // Complete both channels so the write loop drains remaining data and exits
        _outbound.Writer.TryComplete();
        _flushSignal.Writer.TryComplete();

        // Give the write loop a short window to flush the final batch before canceling
        await Task.Delay(50);

        if (_clientCts is { } cts)
            await cts.CancelAsync();
        else
            _socket.Close();
    }

    private async Task RunPingTimerAsync(CancellationToken ct)
    {
        using var timer = new PeriodicTimer(_options.PingInterval);
        try
        {
            while (await timer.WaitForNextTickAsync(ct))
            {
                // Delay first PING until client has responded with PONG or 2 seconds elapsed.
                // Go: opts.go — DisableShortFirstPing bypasses this grace period.
                if (!_options.DisableShortFirstPing
                    && !_flags.HasFlag(ClientFlags.FirstPongSent)
                    && (DateTime.UtcNow - StartTime).TotalSeconds < 2)
                {
                    continue;
                }

                var elapsed = Environment.TickCount64 - Interlocked.Read(ref _lastIn);
                if (elapsed < (long)_options.PingInterval.TotalMilliseconds)
                {
                    // Client was recently active, skip ping
                    Interlocked.Exchange(ref _pingsOut, 0);
                    continue;
                }

                if (Volatile.Read(ref _pingsOut) + 1 > _options.MaxPingsOut)
                {
                    _logger.LogDebug("Client {ClientId} stale connection -- closing", Id);
                    Interlocked.Increment(ref _serverStats.StaleConnections);
                    Interlocked.Increment(ref _serverStats.StaleConnectionClients);
                    await SendErrAndCloseAsync(NatsProtocol.ErrStaleConnection, ClientClosedReason.StaleConnection);
                    return;
                }

                var currentPingsOut = Interlocked.Increment(ref _pingsOut);
                _logger.LogDebug("Client {ClientId} sending PING ({PingsOut}/{MaxPingsOut})",
                    Id, currentPingsOut, _options.MaxPingsOut);
                Interlocked.Exchange(ref _rttStartTicks, DateTime.UtcNow.Ticks);
                WriteProtocol(NatsProtocol.PingBytes);
            }
        }
        catch (OperationCanceledException)
        {
            // Normal shutdown
        }
    }

    /// <summary>
    /// Marks this connection as closed with the given reason.
    /// Sets skip-flush flag for error-related reasons.
    /// Only the first call sets the reason (subsequent calls are no-ops).
    /// </summary>
    public void MarkClosed(ClientClosedReason reason)
    {
        if (CloseReason != ClientClosedReason.None)
            return;

        CloseReason = reason;

        switch (reason)
        {
            case ClientClosedReason.ReadError:
            case ClientClosedReason.WriteError:
            case ClientClosedReason.SlowConsumerPendingBytes:
            case ClientClosedReason.SlowConsumerWriteDeadline:
            case ClientClosedReason.TlsHandshakeError:
                Volatile.Write(ref _skipFlushOnClose, 1);
                break;
        }

        _logger.LogDebug("Client {ClientId} connection closed: {CloseReason}", Id, reason);
    }

    /// <summary>
    /// Flushes pending data (unless skip-flush is set) and closes the connection.
    /// </summary>
    public async Task FlushAndCloseAsync(bool minimalFlush = false)
    {
        if (!ShouldSkipFlush)
        {
            try
            {
                using var flushCts = new CancellationTokenSource(minimalFlush
                    ? TimeSpan.FromMilliseconds(100)
                    : TimeSpan.FromSeconds(1));
                await _stream.FlushAsync(flushCts.Token);
            }
            catch (Exception)
            {
                // Best effort flush — don't let it prevent close
            }
        }

        try { _socket.Shutdown(SocketShutdown.Both); }
        catch (SocketException) { }
        catch (ObjectDisposedException) { }
    }

    public void RemoveSubscription(string sid)
    {
        if (_subs.Remove(sid))
            Account?.DecrementSubscriptions();
    }

    public void RemoveAllSubscriptions(SubList subList)
    {
        foreach (var sub in _subs.Values)
            subList.Remove(sub);
        _subs.Clear();
    }

    public void Dispose()
    {
        _permissions?.Dispose();
        _outbound.Writer.TryComplete();
        _clientCts?.Dispose();
        _stream.Dispose();
        _socket.Dispose();
    }

    /// <summary>
    /// Policy for handling write timeouts based on client kind.
    /// Go reference: server/client.go — CLIENT connections close on timeout,
    /// ROUTER/GATEWAY/LEAF connections attempt TCP-level flush recovery.
    /// </summary>
    public enum WriteTimeoutPolicy
    {
        /// <summary>Close the connection on write timeout (used for CLIENT kind).</summary>
        Close,

        /// <summary>Attempt TCP-level flush and continue (used for ROUTER, GATEWAY, LEAF).</summary>
        TcpFlush,
    }

    /// <summary>
    /// Returns the write timeout policy for the given client kind.
    /// Go reference: server/client.go — routes/gateways/leafnodes get TcpFlush,
    /// regular clients get Close.
    /// </summary>
    public static WriteTimeoutPolicy GetWriteTimeoutPolicy(ClientKind kind) => kind switch
    {
        ClientKind.Client => WriteTimeoutPolicy.Close,
        ClientKind.Router => WriteTimeoutPolicy.TcpFlush,
        ClientKind.Gateway => WriteTimeoutPolicy.TcpFlush,
        ClientKind.Leaf => WriteTimeoutPolicy.TcpFlush,
        _ => WriteTimeoutPolicy.Close,
    };

    /// <summary>
    /// Result of a flush operation, tracking partial write progress.
    /// Go reference: server/client.go — partial write handling for routes and gateways.
    /// </summary>
    public readonly record struct FlushResult(long BytesAttempted, long BytesWritten)
    {
        /// <summary>Whether the flush was only partially completed.</summary>
        public bool IsPartial => BytesWritten < BytesAttempted;

        /// <summary>Number of bytes remaining to be written.</summary>
        public long BytesRemaining => BytesAttempted - BytesWritten;
    }

    /// <summary>
    /// Blocks producers when the client's outbound buffer is near capacity.
    /// Go reference: server/client.go (stc channel, stalledRoute handling).
    /// When pending bytes exceed 75% of maxPending, producers must wait until
    /// the write loop drains enough data.
    /// </summary>
    public sealed class StallGate
    {
        private readonly long _threshold;
        private volatile SemaphoreSlim? _semaphore;
        private readonly Lock _gate = new();

        /// <summary>
        /// Creates a stall gate with the given maxPending capacity.
        /// The stall threshold is set at 75% of maxPending.
        /// Go reference: server/client.go stc channel creation.
        /// </summary>
        public StallGate(long maxPending)
        {
            _threshold = maxPending * 3 / 4;
        }

        /// <summary>Whether producers are currently being stalled.</summary>
        public bool IsStalled
        {
            get { lock (_gate) return _semaphore is not null; }
        }

        /// <summary>
        /// Updates pending byte count and activates/deactivates the stall gate.
        /// Go reference: server/client.go stalledRoute check.
        /// </summary>
        public void UpdatePending(long pending)
        {
            lock (_gate)
            {
                if (pending >= _threshold && _semaphore is null)
                {
                    _semaphore = new SemaphoreSlim(0, 1);
                }
                else if (pending < _threshold && _semaphore is not null)
                {
                    Release();
                }
            }
        }

        /// <summary>
        /// Waits for the stall gate to release. Returns true if released,
        /// false if timed out (indicating the client should be closed as slow consumer).
        /// Go reference: server/client.go stc channel receive with timeout.
        /// </summary>
        public async Task<bool> WaitAsync(TimeSpan timeout)
        {
            SemaphoreSlim? sem;
            lock (_gate) sem = _semaphore;
            if (sem is null) return true;
            return await sem.WaitAsync(timeout);
        }

        /// <summary>
        /// Releases any blocked producers. Called when the write loop has drained
        /// enough data to bring pending bytes below the threshold.
        /// </summary>
        public void Release()
        {
            lock (_gate)
            {
                if (_semaphore is not null)
                {
                    _semaphore.Release();
                    _semaphore = null;
                }
            }
        }
    }
}
