using System.Buffers;
using System.IO.Pipelines;
using System.Net.Sockets;
using System.Security.Cryptography.X509Certificates;
using System.Text;
using NATS.Server.Auth;
using static NATS.Server.Mqtt.MqttBinaryDecoder;

namespace NATS.Server.Mqtt;

public sealed class MqttConnection : IAsyncDisposable
{
    private readonly TcpClient? _tcpClient;
    private readonly Stream _stream;
    private readonly MqttListener _listener;
    private readonly MqttProtocolParser _parser = new();
    private readonly SemaphoreSlim _writeGate = new(1, 1);
    private readonly bool _useBinaryProtocol;
    private readonly X509Certificate2? _clientCert;
    private string _clientId = string.Empty;
    private bool _cleanSession = true;
    private TimeSpan _idleTimeout = Timeout.InfiniteTimeSpan;
    private bool _connected;
    private bool _willCleared;
    private MqttConnectInfo _connectInfo;
    private readonly Dictionary<string, string> _topicToSid = new(StringComparer.Ordinal);
    private int _nextSid;

    /// <summary>Auth result after successful CONNECT (populated for AuthService path).</summary>
    public AuthResult? AuthResult { get; private set; }

    /// <summary>
    /// The adapter that bridges this MQTT connection to the NATS SubList routing.
    /// Created on CONNECT when running with a NatsServer router; null in test-only mode.
    /// </summary>
    public MqttNatsClientAdapter? Adapter { get; private set; }

    public string ClientId => _clientId;

    /// <summary>
    /// Creates a connection from a TcpClient (standard accept path).
    /// </summary>
    public MqttConnection(TcpClient client, MqttListener listener, bool useBinaryProtocol = true)
    {
        _tcpClient = client;
        _stream = client.GetStream();
        _listener = listener;
        _useBinaryProtocol = useBinaryProtocol;
    }

    /// <summary>
    /// Creates a connection from an arbitrary Stream (for TLS wrapping or testing).
    /// </summary>
    public MqttConnection(Stream stream, MqttListener listener, bool useBinaryProtocol = true)
    {
        _stream = stream;
        _listener = listener;
        _useBinaryProtocol = useBinaryProtocol;
    }

    /// <summary>
    /// Creates a connection from a Stream with a TLS client certificate.
    /// Used by the accept loop after TLS handshake completes.
    /// </summary>
    public MqttConnection(Stream stream, MqttListener listener, bool useBinaryProtocol, X509Certificate2? clientCert)
    {
        _stream = stream;
        _listener = listener;
        _useBinaryProtocol = useBinaryProtocol;
        _clientCert = clientCert;
    }

    public async Task RunAsync(CancellationToken ct)
    {
        if (_useBinaryProtocol)
            await RunBinaryAsync(ct);
        else
            await RunTextAsync(ct);
    }

    private async Task RunBinaryAsync(CancellationToken ct)
    {
        var pipeReader = PipeReader.Create(_stream, new StreamPipeReaderOptions(leaveOpen: true));

        try
        {
            while (!ct.IsCancellationRequested)
            {
                ReadResult readResult;
                try
                {
                    // Apply idle timeout for keepalive
                    if (_idleTimeout != Timeout.InfiniteTimeSpan && _idleTimeout > TimeSpan.Zero)
                    {
                        using var timeoutCts = CancellationTokenSource.CreateLinkedTokenSource(ct);
                        timeoutCts.CancelAfter(_idleTimeout);
                        readResult = await pipeReader.ReadAsync(timeoutCts.Token);
                    }
                    else
                    {
                        readResult = await pipeReader.ReadAsync(ct);
                    }
                }
                catch (OperationCanceledException) when (!ct.IsCancellationRequested)
                {
                    // Keepalive timeout
                    break;
                }
                catch
                {
                    break;
                }

                var buffer = readResult.Buffer;

                if (buffer.IsEmpty && readResult.IsCompleted)
                    break;

                while (MqttPacketReader.TryRead(buffer, out var packet, out var consumed))
                {
                    buffer = buffer.Slice(consumed);

                    try
                    {
                        var shouldContinue = await ProcessBinaryPacketAsync(packet!, ct);
                        if (!shouldContinue)
                        {
                            pipeReader.AdvanceTo(consumed);
                            return;
                        }
                    }
                    catch (FormatException)
                    {
                        // Protocol violation — disconnect
                        pipeReader.AdvanceTo(consumed);
                        return;
                    }
                }

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

                if (readResult.IsCompleted)
                    break;
            }
        }
        finally
        {
            await pipeReader.CompleteAsync();

            // Publish will message if not cleanly disconnected
            if (_connected && !_willCleared && _connectInfo.WillTopic != null)
            {
                RoutePublish(_connectInfo.WillTopic, _connectInfo.WillMessage ?? []);
            }
        }
    }

    /// <summary>
    /// Processes a single binary MQTT control packet.
    /// Returns false if the connection should be closed.
    /// </summary>
    private async Task<bool> ProcessBinaryPacketAsync(MqttControlPacket packet, CancellationToken ct)
    {
        // MQTT 3.1.1: First packet MUST be CONNECT
        if (!_connected && packet.Type != MqttControlPacketType.Connect)
            return false;

        switch (packet.Type)
        {
            case MqttControlPacketType.Connect:
                return await HandleConnectAsync(packet, ct);

            case MqttControlPacketType.Publish:
                return await HandlePublishAsync(packet, ct);

            case MqttControlPacketType.PubAck:
                HandlePubAck(packet);
                return true;

            case MqttControlPacketType.PubRec:
                await HandlePubRecAsync(packet, ct);
                return true;

            case MqttControlPacketType.PubRel:
                // Fixed-header flags must be 0x02 for PUBREL
                if (packet.Flags != 0x02)
                    return false;
                await HandlePubRelAsync(packet, ct);
                return true;

            case MqttControlPacketType.PubComp:
                HandlePubComp(packet);
                return true;

            case MqttControlPacketType.Subscribe:
                // Fixed-header flags must be 0x02 for SUBSCRIBE
                if (packet.Flags != MqttProtocolConstants.SubscribeFlags)
                    return false;
                await HandleSubscribeAsync(packet, ct);
                return true;

            case MqttControlPacketType.Unsubscribe:
                // Fixed-header flags must be 0x02 for UNSUBSCRIBE
                if (packet.Flags != 0x02)
                    return false;
                await HandleUnsubscribeAsync(packet, ct);
                return true;

            case MqttControlPacketType.PingReq:
                await WriteBinaryAsync(MqttPacketWriter.WritePingResp(), ct);
                return true;

            case MqttControlPacketType.Disconnect:
                // Clean disconnect — clear will message
                _willCleared = true;
                return false;

            default:
                // Unknown packet type — disconnect
                return false;
        }
    }

    private async Task<bool> HandleConnectAsync(MqttControlPacket packet, CancellationToken ct)
    {
        if (_connected)
            return false; // Second CONNECT is a protocol violation

        var connectInfo = MqttBinaryDecoder.ParseConnect(packet.Payload.Span);
        _connectInfo = connectInfo;

        // MQTT 3.1.1: Reserved bit (bit 0 of connect flags) must be 0
        // This is implicitly validated because we parse individual flag bits

        // Protocol level must be 4 for MQTT 3.1.1
        if (connectInfo.ProtocolLevel != 4)
        {
            await WriteBinaryAsync(
                MqttPacketWriter.WriteConnAck(0x00, MqttProtocolConstants.ConnAckUnacceptableProtocolVersion), ct);
            return false;
        }

        // Will QoS range check (0-2)
        if (connectInfo.WillQoS > 2)
        {
            await WriteBinaryAsync(
                MqttPacketWriter.WriteConnAck(0x00, MqttProtocolConstants.ConnAckIdentifierRejected), ct);
            return false;
        }

        // Empty client-id handling
        if (string.IsNullOrEmpty(connectInfo.ClientId))
        {
            if (connectInfo.CleanSession)
            {
                // Generate a unique client ID
                _clientId = $"auto-{Guid.NewGuid():N}";
            }
            else
            {
                // Empty client-id with persistent session is not allowed
                await WriteBinaryAsync(
                    MqttPacketWriter.WriteConnAck(0x00, MqttProtocolConstants.ConnAckIdentifierRejected), ct);
                return false;
            }
        }
        else
        {
            _clientId = connectInfo.ClientId;
        }

        // Auth check via AuthService (passes TLS client cert for cert-mapping auth)
        var authResult = _listener.AuthenticateMqtt(connectInfo.Username, connectInfo.Password, _clientCert);
        if (authResult == null)
        {
            await WriteBinaryAsync(
                MqttPacketWriter.WriteConnAck(0x00, MqttProtocolConstants.ConnAckNotAuthorized), ct);
            return false;
        }

        AuthResult = authResult;

        // Create MqttNatsClientAdapter for cross-protocol routing (when running with NatsServer)
        if (_listener.AllocateClientId != null)
        {
            var adapterId = _listener.AllocateClientId();
            Adapter = new MqttNatsClientAdapter(this, adapterId);
            Adapter.Account = _listener.ResolveAccount?.Invoke(authResult.AccountName);
            _listener.RegisterMqttAdapter(Adapter);
        }

        // Duplicate client-id takeover
        _listener.TakeoverExistingConnection(_clientId, this);

        _cleanSession = connectInfo.CleanSession;
        _idleTimeout = _listener.ResolveKeepAliveTimeout(connectInfo.KeepAlive);

        // Session-present bit: 1 if resuming existing session, 0 otherwise
        var pending = _listener.OpenSession(_clientId, _cleanSession);
        byte sessionPresent = (byte)(!_cleanSession && pending.Count > 0 ? 0x01 : 0x00);

        _connected = true;
        await WriteBinaryAsync(
            MqttPacketWriter.WriteConnAck(sessionPresent, MqttProtocolConstants.ConnAckAccepted), ct);

        // Redeliver pending QoS 1 messages
        foreach (var redelivery in pending)
        {
            var payloadBytes = Encoding.UTF8.GetBytes(redelivery.Payload);
            await WriteBinaryAsync(
                MqttPacketWriter.WritePublish(redelivery.Topic, payloadBytes, qos: 1, dup: true,
                    packetId: (ushort)redelivery.PacketId), ct);
        }

        return true;
    }

    private async Task<bool> HandlePublishAsync(MqttControlPacket packet, CancellationToken ct)
    {
        var publishInfo = MqttBinaryDecoder.ParsePublish(packet.Payload.Span, packet.Flags);

        // Non-zero packet identifier required for QoS > 0
        if (publishInfo.QoS > 0 && publishInfo.PacketId == 0)
            return false;

        switch (publishInfo.QoS)
        {
            case 0:
                RoutePublish(publishInfo.Topic, publishInfo.Payload);
                break;

            case 1:
                _listener.RecordPendingPublish(_clientId, publishInfo.PacketId, publishInfo.Topic,
                    Encoding.UTF8.GetString(publishInfo.Payload.Span));
                await WriteBinaryAsync(MqttPacketWriter.WritePubAck(publishInfo.PacketId), ct);
                RoutePublish(publishInfo.Topic, publishInfo.Payload);
                break;

            case 2:
                // QoS 2 step 1: store and send PUBREC (delivery deferred to PUBREL)
                _listener.RecordPendingPublish(_clientId, publishInfo.PacketId, publishInfo.Topic,
                    Encoding.UTF8.GetString(publishInfo.Payload.Span));
                await WriteBinaryAsync(MqttPacketWriter.WritePubRec(publishInfo.PacketId), ct);
                break;
        }

        // Handle retained messages
        if (publishInfo.Retain)
        {
            _listener.SetRetainedMessage(publishInfo.Topic,
                publishInfo.Payload.Length == 0 ? null : Encoding.UTF8.GetString(publishInfo.Payload.Span));
        }

        return true;
    }

    /// <summary>
    /// Routes a published message through the NATS SubList (cross-protocol) if a router
    /// and adapter are available, otherwise falls back to MQTT-only fan-out.
    /// </summary>
    private void RoutePublish(string mqttTopic, ReadOnlyMemory<byte> payload)
    {
        if (_listener.Router != null && Adapter != null)
        {
            var natsSubject = MqttTopicMapper.MqttToNats(mqttTopic);
            _listener.Router.ProcessMessage(natsSubject, null, default, payload, Adapter);
        }
        else
        {
            // Test-only fallback: MQTT-only fan-out
            _ = _listener.PublishAsync(mqttTopic, Encoding.UTF8.GetString(payload.Span), this, CancellationToken.None);
        }
    }

    private void HandlePubAck(MqttControlPacket packet)
    {
        if (packet.Payload.Length < 2) return;
        var packetId = (ushort)((packet.Payload.Span[0] << 8) | packet.Payload.Span[1]);
        _listener.AckPendingPublish(_clientId, packetId);
    }

    private async Task HandlePubRecAsync(MqttControlPacket packet, CancellationToken ct)
    {
        if (packet.Payload.Length < 2) return;
        var packetId = (ushort)((packet.Payload.Span[0] << 8) | packet.Payload.Span[1]);
        await WriteBinaryAsync(MqttPacketWriter.WritePubRel(packetId), ct);
    }

    private async Task HandlePubRelAsync(MqttControlPacket packet, CancellationToken ct)
    {
        if (packet.Payload.Length < 2) return;
        var packetId = (ushort)((packet.Payload.Span[0] << 8) | packet.Payload.Span[1]);

        // QoS 2 step 2: deliver the stored message, then ack and send PUBCOMP
        var pending = _listener.GetPendingPublish(_clientId, packetId);
        if (pending != null)
        {
            RoutePublish(pending.Topic, Encoding.UTF8.GetBytes(pending.Payload));
        }

        _listener.AckPendingPublish(_clientId, packetId);
        await WriteBinaryAsync(MqttPacketWriter.WritePubComp(packetId), ct);
    }

    private void HandlePubComp(MqttControlPacket packet)
    {
        if (packet.Payload.Length < 2) return;
        var packetId = (ushort)((packet.Payload.Span[0] << 8) | packet.Payload.Span[1]);
        _listener.AckPendingPublish(_clientId, packetId);
    }

    private async Task HandleSubscribeAsync(MqttControlPacket packet, CancellationToken ct)
    {
        var subscribeInfo = MqttBinaryDecoder.ParseSubscribe(packet.Payload.Span, packet.Flags);

        // Grant QoS (cap at 2)
        var grantedQoS = new byte[subscribeInfo.Filters.Count];
        for (var i = 0; i < subscribeInfo.Filters.Count; i++)
        {
            var (topicFilter, requestedQoS) = subscribeInfo.Filters[i];

            if (Adapter != null)
            {
                // Route through SubList for cross-protocol delivery
                var natsSubject = MqttTopicMapper.MqttToNats(topicFilter);
                var sid = $"$MQTT_{Interlocked.Increment(ref _nextSid)}";
                Adapter.AddSubscription(natsSubject, sid);
                _topicToSid[topicFilter] = sid;
            }
            else
            {
                // Test-only fallback: MQTT-only subscription
                _listener.RegisterSubscription(this, topicFilter);
            }

            grantedQoS[i] = Math.Min(requestedQoS, (byte)2);

            // Deliver retained messages for this topic filter
            var retained = _listener.GetRetainedMessage(topicFilter);
            if (retained != null)
            {
                var retainedPayload = Encoding.UTF8.GetBytes(retained);
                await WriteBinaryAsync(
                    MqttPacketWriter.WritePublish(topicFilter, retainedPayload, qos: 0, retain: true, packetId: 0), ct);
            }
        }

        await WriteBinaryAsync(MqttPacketWriter.WriteSubAck(subscribeInfo.PacketId, grantedQoS), ct);
    }

    private async Task HandleUnsubscribeAsync(MqttControlPacket packet, CancellationToken ct)
    {
        var unsubInfo = MqttBinaryDecoder.ParseUnsubscribe(packet.Payload.Span, packet.Flags);

        foreach (var filter in unsubInfo.Filters)
        {
            if (Adapter != null && _topicToSid.Remove(filter, out var sid))
            {
                Adapter.RemoveSubscription(sid);
            }
            else
            {
                _listener.UnregisterSubscription(this, filter);
            }
        }

        await WriteBinaryAsync(MqttPacketWriter.WriteUnsubAck(unsubInfo.PacketId), ct);
    }

    /// <summary>
    /// Sends a binary MQTT PUBLISH packet to this connection (for message delivery).
    /// </summary>
    public async Task SendBinaryPublishAsync(string topic, ReadOnlyMemory<byte> payload, byte qos,
        bool retain, ushort packetId, CancellationToken ct)
    {
        var packet = MqttPacketWriter.WritePublish(topic, payload.Span, qos, retain, packetId: packetId);
        await WriteBinaryAsync(packet, ct);
    }

    /// <summary>
    /// Sends a message to the connection. Used by the listener for fan-out delivery.
    /// In binary mode, sends a PUBLISH packet; in text mode, sends a text line.
    /// </summary>
    public Task SendMessageAsync(string topic, string payload, CancellationToken ct)
    {
        if (_useBinaryProtocol)
        {
            var payloadBytes = Encoding.UTF8.GetBytes(payload);
            return SendBinaryPublishAsync(topic, payloadBytes, qos: 0, retain: false, packetId: 0, ct);
        }

        return WriteLineAsync($"MSG {topic} {payload}", ct);
    }

    public async ValueTask DisposeAsync()
    {
        // Clean up adapter subscriptions and unregister from listener
        if (Adapter != null)
        {
            Adapter.RemoveAllSubscriptions();
            _listener.UnregisterMqttAdapter(Adapter);
        }

        _listener.Unregister(this);
        _writeGate.Dispose();
        await _stream.DisposeAsync();
        _tcpClient?.Dispose();
    }

    /// <summary>
    /// Forces this connection to close (used for duplicate client-id takeover).
    /// </summary>
    internal void ForceClose()
    {
        try { _stream.Close(); }
        catch { /* best effort */ }
        try { _tcpClient?.Close(); }
        catch { /* best effort */ }
    }

    private async Task WriteBinaryAsync(byte[] data, CancellationToken ct)
    {
        await _writeGate.WaitAsync(ct);
        try
        {
            await _stream.WriteAsync(data, ct);
            await _stream.FlushAsync(ct);
        }
        finally
        {
            _writeGate.Release();
        }
    }

    // --- Text protocol methods (for backward compatibility during test migration) ---
    // TODO: Remove after test migration — deadline: Phase 3 completion

    private async Task RunTextAsync(CancellationToken ct)
    {
        while (!ct.IsCancellationRequested)
        {
            string line;
            try
            {
                line = await ReadLineAsync(ct, _idleTimeout);
            }
            catch
            {
                break;
            }

            var packet = _parser.ParseLine(line);
            switch (packet.Type)
            {
                case MqttPacketType.Connect:
                    if (!_listener.TryAuthenticate(packet.Username, packet.Password))
                    {
                        await WriteLineAsync("ERR mqtt auth failed", ct);
                        return;
                    }

                    _clientId = packet.ClientId;
                    _cleanSession = packet.CleanSession;
                    _idleTimeout = _listener.ResolveKeepAliveTimeout(packet.KeepAliveSeconds);
                    var pending = _listener.OpenSession(_clientId, _cleanSession);
                    await WriteLineAsync("CONNACK", ct);
                    foreach (var redelivery in pending)
                        await WriteLineAsync($"REDLIVER {redelivery.PacketId} {redelivery.Topic} {redelivery.Payload}", ct);
                    break;
                case MqttPacketType.Subscribe:
                    _listener.RegisterSubscription(this, packet.Topic);
                    await WriteLineAsync($"SUBACK {packet.Topic}", ct);
                    break;
                case MqttPacketType.Publish:
                    await _listener.PublishAsync(packet.Topic, packet.Payload, this, ct);
                    break;
                case MqttPacketType.PublishQos1:
                    _listener.RecordPendingPublish(_clientId, packet.PacketId, packet.Topic, packet.Payload);
                    await WriteLineAsync($"PUBACK {packet.PacketId}", ct);
                    await _listener.PublishAsync(packet.Topic, packet.Payload, this, ct);
                    break;
                case MqttPacketType.Ack:
                    _listener.AckPendingPublish(_clientId, packet.PacketId);
                    break;
            }
        }
    }

    private async Task WriteLineAsync(string line, CancellationToken ct)
    {
        await _writeGate.WaitAsync(ct);
        try
        {
            var bytes = Encoding.UTF8.GetBytes(line + "\n");
            await _stream.WriteAsync(bytes, ct);
            await _stream.FlushAsync(ct);
        }
        finally
        {
            _writeGate.Release();
        }
    }

    private async Task<string> ReadLineAsync(CancellationToken ct, TimeSpan idleTimeout)
    {
        CancellationToken token = ct;
        CancellationTokenSource? timeoutCts = null;
        if (idleTimeout != Timeout.InfiniteTimeSpan && idleTimeout > TimeSpan.Zero)
        {
            timeoutCts = CancellationTokenSource.CreateLinkedTokenSource(ct);
            timeoutCts.CancelAfter(idleTimeout);
            token = timeoutCts.Token;
        }

        var bytes = new List<byte>(64);
        var single = new byte[1];
        try
        {
            while (true)
            {
                var read = await _stream.ReadAsync(single.AsMemory(0, 1), token);
                if (read == 0)
                    throw new IOException("mqtt closed");
                if (single[0] == (byte)'\n')
                    break;
                if (single[0] != (byte)'\r')
                    bytes.Add(single[0]);
            }
        }
        catch (OperationCanceledException) when (timeoutCts != null && !ct.IsCancellationRequested)
        {
            throw new IOException("mqtt keepalive timeout");
        }
        finally
        {
            timeoutCts?.Dispose();
        }

        return Encoding.UTF8.GetString([.. bytes]);
    }
}
