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

using System.Text;

namespace ZB.MOM.NatsNet.Server.Mqtt;

/// <summary>
/// MQTT packet parsing and processing handlers.
/// Mirrors the mqttParseConnect / mqttProcessConnect / mqttEnqueueConnAck
/// functions in server/mqtt.go.
/// </summary>
internal static class MqttPacketHandlers
{
    /// <summary>
    /// Parses an MQTT CONNECT packet from the reader.
    /// Returns (returnCode, connectProto, error).
    /// returnCode == 0 means success; non-zero is a CONNACK return code.
    /// Mirrors Go <c>mqttParseConnect()</c>.
    /// </summary>
    public static (byte rc, MqttConnectProto? cp, Exception? err) ParseConnect(MqttReader r)
    {
        // --- Protocol Name ---
        string protoName;
        try { protoName = r.ReadString("protocol name"); }
        catch (Exception ex) { return (0, null, ex); }

        if (protoName != "MQTT")
            return (MqttConnAckRc.UnacceptableProtocol, null,
                new InvalidOperationException($"unsupported MQTT protocol: \"{protoName}\""));

        // --- Protocol Level ---
        byte level;
        try { level = r.ReadByte("protocol level"); }
        catch (Exception ex) { return (0, null, ex); }

        if (level != 0x04)
            return (MqttConnAckRc.UnacceptableProtocol, null,
                new InvalidOperationException($"unsupported MQTT protocol level: {level}"));

        // --- Connect Flags ---
        byte flags;
        try { flags = r.ReadByte("connect flags"); }
        catch (Exception ex) { return (0, null, ex); }

        if ((flags & MqttConnectFlag.Reserved) != 0)
            return (0, null, new InvalidOperationException("CONNECT flags reserved bit must be 0"));

        bool cleanSession = (flags & MqttConnectFlag.CleanSession) != 0;
        bool willFlag = (flags & MqttConnectFlag.WillFlag) != 0;
        byte willQos = (byte)((flags & MqttConnectFlag.WillQoS) >> 3);
        bool willRetain = (flags & MqttConnectFlag.WillRetain) != 0;
        bool hasPassword = (flags & MqttConnectFlag.PasswordFlag) != 0;
        bool hasUsername = (flags & MqttConnectFlag.UsernameFlag) != 0;

        // Validate Will flags.
        if (!willFlag)
        {
            if (willQos != 0)
                return (0, null, new InvalidOperationException("Will QoS must be 0 when Will Flag is 0"));
            if (willRetain)
                return (0, null, new InvalidOperationException("Will Retain must be 0 when Will Flag is 0"));
        }
        else
        {
            if (willQos > 2)
                return (0, null, new InvalidOperationException($"invalid Will QoS: {willQos}"));
        }

        // Username/password consistency.
        if (hasPassword && !hasUsername)
            return (0, null, new InvalidOperationException("password flag without username flag"));

        // --- Keep Alive ---
        ushort keepAlive;
        try { keepAlive = r.ReadUInt16("keep alive"); }
        catch (Exception ex) { return (0, null, ex); }

        // --- Client ID ---
        string clientId;
        try { clientId = r.ReadString("client id"); }
        catch (Exception ex) { return (0, null, ex); }

        if (string.IsNullOrEmpty(clientId))
        {
            if (!cleanSession)
                return (MqttConnAckRc.IdentifierRejected, null,
                    new InvalidOperationException("empty client ID requires clean session flag"));
            // Generate a unique client ID.
            clientId = Guid.NewGuid().ToString("N");
        }

        // --- Will Topic & Message ---
        MqttWill? will = null;
        if (willFlag)
        {
            string willTopic;
            try { willTopic = r.ReadString("will topic"); }
            catch (Exception ex) { return (0, null, ex); }

            if (string.IsNullOrEmpty(willTopic))
                return (0, null, new InvalidOperationException("empty will topic"));

            byte[] willMsg;
            try { willMsg = r.ReadBytes("will message", copy: true); }
            catch (Exception ex) { return (0, null, ex); }

            // Convert MQTT topic to NATS subject.
            var topicBytes = Encoding.UTF8.GetBytes(willTopic);
            byte[] subjectBytes;
            try { subjectBytes = MqttSubjectConverter.MqttTopicToNatsPubSubject(topicBytes); }
            catch (Exception ex) { return (0, null, ex); }

            will = new MqttWill
            {
                Topic = willTopic,
                Subject = Encoding.UTF8.GetString(subjectBytes),
                Msg = willMsg.Length > 0 ? willMsg : null,
                Qos = willQos,
                Retain = willRetain,
            };
        }

        // --- Username ---
        string username = string.Empty;
        if (hasUsername)
        {
            try { username = r.ReadString("username"); }
            catch (Exception ex) { return (0, null, ex); }

            if (string.IsNullOrEmpty(username))
                return (0, null, new InvalidOperationException("empty username"));
        }

        // --- Password ---
        byte[]? password = null;
        if (hasPassword)
        {
            try { password = r.ReadBytes("password", copy: true); }
            catch (Exception ex) { return (0, null, ex); }
        }

        var cp = new MqttConnectProto
        {
            ClientId = clientId,
            Will = will,
            Username = username,
            Password = password,
            CleanSession = cleanSession,
            KeepAlive = keepAlive,
        };

        return (MqttConnAckRc.Accepted, cp, null);
    }

    /// <summary>
    /// Processes a parsed CONNECT packet: sets client state, sends CONNACK.
    /// Minimal implementation — full session management deferred to Task 6.
    /// Mirrors Go <c>mqttProcessConnect()</c>.
    /// </summary>
    public static Exception? ProcessConnect(ClientConnection c, MqttConnectProto cp)
    {
        var mqtt = c.Mqtt!;

        // Store client identity.
        mqtt.ClientId = cp.ClientId;
        mqtt.CleanSession = cp.CleanSession;
        mqtt.KeepAlive = cp.KeepAlive;
        mqtt.Will = cp.Will;

        // Store auth credentials on client options.
        if (!string.IsNullOrEmpty(cp.Username))
            c.Opts.Username = cp.Username;
        if (cp.Password != null)
            c.Opts.Password = Encoding.UTF8.GetString(cp.Password);

        // Mark as connected.
        c.Flags |= ClientFlags.ConnectReceived;

        // Set keep-alive read deadline.
        if (cp.KeepAlive > 0)
        {
            // MQTT spec: server MUST disconnect if no packet within 1.5x keep-alive.
            var deadline = TimeSpan.FromSeconds(cp.KeepAlive * 1.5);
            mqtt.KeepAlive = cp.KeepAlive;
            // TODO: set read deadline on connection stream (Task 7)
        }

        // Send CONNACK (accepted, no session present for now).
        EnqueueConnAck(c, MqttConnAckRc.Accepted, sessionPresent: false);

        return null;
    }

    /// <summary>
    /// Enqueues a CONNACK packet to the client.
    /// Mirrors Go <c>mqttEnqueueConnAck()</c>.
    /// </summary>
    public static void EnqueueConnAck(ClientConnection c, byte rc, bool sessionPresent)
    {
        byte sp = 0;
        if (rc == MqttConnAckRc.Accepted && sessionPresent)
            sp = 1;

        ReadOnlySpan<byte> connack = [MqttPacket.ConnectAck, 0x02, sp, rc];
        lock (c)
        {
            c.EnqueueProto(connack);
        }
    }

    /// <summary>
    /// Handles DISCONNECT: clears the will message and closes the connection.
    /// Mirrors Go DISCONNECT case in mqttParse().
    /// </summary>
    public static void HandleDisconnect(ClientConnection c)
    {
        // Per MQTT spec 3.1.2-8: discard the will message on clean disconnect.
        lock (c)
        {
            if (c.Mqtt != null)
                c.Mqtt.Will = null;
        }

        // Close the connection cleanly.
        c.CloseConnection(ClosedState.ClientClosed);
    }

    // =========================================================================
    // Will message delivery
    // =========================================================================

    /// <summary>
    /// Delivers the will message if one is set and the connection was not cleanly
    /// disconnected (i.e., DISCONNECT was not received, so will is still non-null).
    /// Called from CloseConnection(). Mirrors Go <c>mqttHandleWill()</c>.
    /// </summary>
    public static void DeliverWill(ClientConnection c)
    {
        MqttWill? will;
        lock (c)
        {
            if (c.Mqtt == null) return;
            will = c.Mqtt.Will;
            c.Mqtt.Will = null; // Prevent duplicate delivery.
        }

        if (will == null)
            return;

        // Deliver the will message via internal NATS routing.
        var payload = will.Msg ?? [];
        c.ProcessInboundClientMsg(payload);
    }

    // =========================================================================
    // PUBLISH parsing + processing
    // =========================================================================

    /// <summary>
    /// Parses an MQTT PUBLISH packet from the reader.
    /// <paramref name="pl"/> is the remaining length, <paramref name="flags"/> is the lower
    /// nibble of the first byte (DUP/QoS/RETAIN).
    /// Mirrors Go <c>mqttParsePub()</c>.
    /// </summary>
    public static (MqttPublishInfo? pp, Exception? err) ParsePub(
        MqttReader r, int pl, byte flags, bool rejectQoS2)
    {
        var start = r.Position;

        // Extract QoS from flags bits 1-2.
        var qos = (byte)((flags & MqttPubFlag.QoS) >> 1);
        if (qos > 2)
            return (null, new InvalidOperationException($"QoS value must be 0, 1, or 2, got {qos}"));
        if (qos == 2 && rejectQoS2)
            return (null, new InvalidOperationException("QoS-2 PUBLISH rejected by server policy"));

        // Topic name (2-byte length prefix + UTF-8 bytes).
        byte[] topicBytes;
        try { topicBytes = r.ReadBytes("topic", copy: false); }
        catch (Exception ex) { return (null, ex); }

        if (topicBytes.Length == 0)
            return (null, new InvalidOperationException("empty topic in PUBLISH"));

        // Convert MQTT topic to NATS subject.
        byte[] subjectBytes;
        try { subjectBytes = MqttSubjectConverter.MqttTopicToNatsPubSubject(topicBytes); }
        catch (Exception ex) { return (null, ex); }

        // Packet identifier (QoS > 0 only).
        ushort pi = 0;
        if (qos > 0)
        {
            try { pi = r.ReadUInt16("packet identifier"); }
            catch (Exception ex) { return (null, ex); }
            if (pi == 0)
                return (null, new InvalidOperationException("packet identifier must not be 0 for QoS > 0"));
        }

        // Payload = remaining bytes in the packet.
        var consumed = r.Position - start;
        var payloadLen = pl - consumed;
        byte[]? payload = null;
        if (payloadLen > 0)
        {
            try { payload = r.ReadSlice(payloadLen, "publish payload"); }
            catch (Exception ex) { return (null, ex); }
        }
        else if (payloadLen < 0)
        {
            return (null, new InvalidOperationException("PUBLISH packet payload length underflow"));
        }

        var pp = new MqttPublishInfo
        {
            Topic = Encoding.UTF8.GetString(topicBytes),
            Subject = Encoding.UTF8.GetString(subjectBytes),
            Msg = payload,
            Qos = qos,
            Retain = (flags & MqttPubFlag.Retain) != 0,
            Dup = (flags & MqttPubFlag.Dup) != 0,
            Pi = pi,
        };

        return (pp, null);
    }

    /// <summary>
    /// Processes an inbound MQTT PUBLISH packet.
    /// QoS 0: routes immediately. QoS 1: deliver + PUBACK. QoS 2: store + PUBREC.
    /// Mirrors Go <c>mqttProcessPub()</c>.
    /// </summary>
    public static Exception? ProcessPub(ClientConnection c, MqttPublishInfo pp)
    {
        var payload = pp.Msg ?? [];

        switch (pp.Qos)
        {
            case 0:
                c.ProcessInboundClientMsg(payload);
                return null;

            case 1:
                // QoS 1: deliver immediately, then send PUBACK.
                c.ProcessInboundClientMsg(payload);
                EnqueuePubResponse(c, MqttPacket.PubAck, pp.Pi);
                return null;

            case 2:
                // QoS 2: store message pending PUBREL, then send PUBREC.
                // Full JetStream persistence deferred to Task 6 — uses in-memory store.
                lock (c)
                {
                    c.Mqtt!.QoS2Pending[pp.Pi] = pp;
                }
                EnqueuePubResponse(c, MqttPacket.PubRec, pp.Pi);
                return null;

            default:
                return new InvalidOperationException($"invalid QoS: {pp.Qos}");
        }
    }

    // =========================================================================
    // PUBACK / PUBREC / PUBREL / PUBCOMP handling
    // =========================================================================

    /// <summary>
    /// Parses a packet-identifier-only packet (PUBACK, PUBREC, PUBREL, PUBCOMP).
    /// Returns (pi, error). Mirrors Go <c>mqttParsePIPacket()</c>.
    /// </summary>
    public static (ushort pi, Exception? err) ParsePiPacket(MqttReader r)
    {
        ushort pi;
        try { pi = r.ReadUInt16("packet identifier"); }
        catch (Exception ex) { return (0, ex); }
        if (pi == 0)
            return (0, new InvalidOperationException("packet identifier must not be 0"));
        return (pi, null);
    }

    /// <summary>
    /// Processes an inbound PUBREL packet (QoS 2 phase 2).
    /// Retrieves the stored QoS 2 message, delivers it, and sends PUBCOMP.
    /// Mirrors Go <c>mqttProcessPubRel()</c>.
    /// </summary>
    public static Exception? ProcessPubRel(ClientConnection c, ushort pi)
    {
        // Always send PUBCOMP, even if message not found (idempotency).
        MqttPublishInfo? pp;
        lock (c)
        {
            c.Mqtt!.QoS2Pending.Remove(pi, out pp);
        }

        // Send PUBCOMP.
        EnqueuePubResponse(c, MqttPacket.PubComp, pi);

        // Deliver the stored message if found.
        if (pp != null)
        {
            var payload = pp.Msg ?? [];
            c.ProcessInboundClientMsg(payload);
        }

        return null;
    }

    /// <summary>
    /// Processes an inbound PUBACK (client acknowledges QoS 1 message from server).
    /// Removes from pending tracking. Mirrors Go <c>mqttProcessPubAck()</c>.
    /// </summary>
    public static Exception? ProcessPubAck(ClientConnection c, ushort pi)
    {
        lock (c)
        {
            c.Mqtt!.Pending.Remove(pi);
        }
        return null;
    }

    /// <summary>
    /// Processes an inbound PUBREC (client acknowledges QoS 2 message from server, phase 1).
    /// Transitions to PUBREL phase. Mirrors Go <c>mqttProcessPubRec()</c>.
    /// </summary>
    public static Exception? ProcessPubRec(ClientConnection c, ushort pi)
    {
        // In the full implementation, this would store a PUBREL in JetStream.
        // For now, send PUBREL immediately and remove from pending.
        lock (c)
        {
            c.Mqtt!.Pending.Remove(pi);
        }
        EnqueuePubResponse(c, MqttPacket.PubRel, pi);
        return null;
    }

    /// <summary>
    /// Processes an inbound PUBCOMP (client acknowledges PUBREL, QoS 2 complete).
    /// Mirrors Go <c>mqttProcessPubComp()</c>.
    /// </summary>
    public static Exception? ProcessPubComp(ClientConnection c, ushort pi)
    {
        // Final cleanup — remove any tracking for this PI.
        lock (c)
        {
            c.Mqtt!.Pending.Remove(pi);
        }
        return null;
    }

    /// <summary>
    /// Enqueues a PUBACK/PUBREC/PUBREL/PUBCOMP response packet.
    /// All four share the same 4-byte format: [type] [0x02] [PI high] [PI low].
    /// PUBREL has bit 1 set in byte 0 per MQTT spec [3.6.1-1].
    /// Mirrors Go <c>mqttEnqueuePubResponse()</c>.
    /// </summary>
    public static void EnqueuePubResponse(ClientConnection c, byte packetType, ushort pi)
    {
        var b0 = packetType;
        // PUBREL requires fixed header bits 0010 per MQTT spec.
        if (packetType == MqttPacket.PubRel)
            b0 |= 0x02;

        ReadOnlySpan<byte> packet = [b0, 0x02, (byte)(pi >> 8), (byte)(pi & 0xFF)];
        lock (c)
        {
            c.EnqueueProto(packet);
        }
    }

    // =========================================================================
    // SUBSCRIBE / UNSUBSCRIBE parsing + processing
    // =========================================================================

    /// <summary>
    /// Parses SUBSCRIBE or UNSUBSCRIBE packet filters from the reader.
    /// Returns (packetId, filters, error).
    /// Mirrors Go <c>mqttParseSubsOrUnsubs()</c>.
    /// </summary>
    public static (ushort pi, List<MqttFilter>? filters, Exception? err) ParseSubsOrUnsubs(
        MqttReader r, byte firstByte, int pl, bool isSub)
    {
        var kind = isSub ? "SUBSCRIBE" : "UNSUBSCRIBE";

        // Validate reserved flags (must be 0x02 per spec).
        var expectedFlags = isSub ? MqttConst.SubscribeFlags : MqttConst.UnsubscribeFlags;
        var actualFlags = (byte)(firstByte & MqttPacket.FlagMask);
        if (actualFlags != expectedFlags)
            return (0, null, new InvalidOperationException(
                $"{kind} reserved flags must be 0x{expectedFlags:X2}, got 0x{actualFlags:X2}"));

        var end = r.Position + pl;

        // Packet identifier.
        ushort pi;
        try { pi = r.ReadUInt16("packet identifier"); }
        catch (Exception ex) { return (0, null, ex); }
        if (pi == 0)
            return (0, null, new InvalidOperationException("packet identifier must not be 0"));

        var filters = new List<MqttFilter>();
        while (r.Position < end)
        {
            // Topic filter (2-byte length + UTF-8 bytes).
            byte[] topicBytes;
            try { topicBytes = r.ReadBytes("topic filter", copy: false); }
            catch (Exception ex) { return (0, null, ex); }

            if (topicBytes.Length == 0)
                return (0, null, new InvalidOperationException("empty topic filter"));

            // Convert MQTT filter to NATS subject.
            string natsSubject;
            try
            {
                var subjBytes = MqttSubjectConverter.MqttFilterToNatsSubject(topicBytes);
                natsSubject = Encoding.UTF8.GetString(subjBytes);
            }
            catch (Exception ex) { return (0, null, ex); }

            byte qos = 0;
            if (isSub)
            {
                try { qos = r.ReadByte("QoS"); }
                catch (Exception ex) { return (0, null, ex); }
                if (qos > 2)
                    return (0, null, new InvalidOperationException($"invalid QoS value: {qos}"));
            }

            filters.Add(new MqttFilter
            {
                Filter = natsSubject,
                Qos = qos,
                Ttopic = topicBytes,
            });
        }

        if (filters.Count == 0)
            return (0, null, new InvalidOperationException(
                $"{kind} must contain at least one topic filter"));

        return (pi, filters, null);
    }

    /// <summary>
    /// Processes parsed SUBSCRIBE filters: creates NATS subscriptions and sends SUBACK.
    /// Mirrors Go <c>mqttProcessSubs()</c> (minimal — session/JetStream deferred to Task 5-6).
    /// </summary>
    public static Exception? ProcessSubs(ClientConnection c, ushort pi, List<MqttFilter> filters)
    {
        var mqtt = c.Mqtt!;

        foreach (var f in filters)
        {
            // Downgrade QoS 2 if configured.
            if (f.Qos == 2 && mqtt.DowngradeQoS2Sub)
                f.Qos = 1;

            // Create NATS subscription using the filter as SID.
            var sid = Encoding.UTF8.GetBytes(f.Filter);
            var subject = Encoding.UTF8.GetBytes(f.Filter);
            var (_, err) = c.ProcessSub(subject, queue: null, sid: sid, noForward: false);
            if (err != null)
            {
                // Mark as failure in SUBACK.
                f.Qos = MqttConst.SubAckFailure;
            }
        }

        EnqueueSubAck(c, pi, filters);
        return null;
    }

    /// <summary>
    /// Processes parsed UNSUBSCRIBE filters: removes subscriptions and sends UNSUBACK.
    /// Mirrors Go <c>mqttProcessUnsubs()</c> (minimal — session cleanup deferred to Task 6).
    /// </summary>
    public static Exception? ProcessUnsubs(ClientConnection c, ushort pi, List<MqttFilter> filters)
    {
        foreach (var f in filters)
        {
            var sid = Encoding.UTF8.GetBytes(f.Filter);
            c.RemoveSubBySid(sid);
        }

        EnqueueUnsubAck(c, pi);
        return null;
    }

    // =========================================================================
    // SUBACK / UNSUBACK encoding
    // =========================================================================

    /// <summary>
    /// Enqueues a SUBACK packet. Mirrors Go <c>mqttEnqueueSubAck()</c>.
    /// Format: [0x90] [var-int length] [PI high] [PI low] [QoS per filter...]
    /// </summary>
    public static void EnqueueSubAck(ClientConnection c, ushort pi, List<MqttFilter> filters)
    {
        var payloadLen = 2 + filters.Count;
        var packet = new byte[1 + VarIntSize(payloadLen) + payloadLen];
        var pos = 0;
        packet[pos++] = MqttPacket.SubAck;
        pos += WriteVarInt(packet.AsSpan(pos), payloadLen);
        packet[pos++] = (byte)(pi >> 8);
        packet[pos++] = (byte)(pi & 0xFF);
        foreach (var f in filters)
            packet[pos++] = f.Qos;

        lock (c)
        {
            c.EnqueueProto(packet.AsSpan(0, pos));
        }
    }

    /// <summary>
    /// Enqueues an UNSUBACK packet. Mirrors Go <c>mqttEnqueueUnsubAck()</c>.
    /// Format: [0xB0] [0x02] [PI high] [PI low]
    /// </summary>
    public static void EnqueueUnsubAck(ClientConnection c, ushort pi)
    {
        ReadOnlySpan<byte> unsuback =
            [MqttPacket.UnsubAck, 0x02, (byte)(pi >> 8), (byte)(pi & 0xFF)];
        lock (c)
        {
            c.EnqueueProto(unsuback);
        }
    }

    // =========================================================================
    // Helpers
    // =========================================================================

    private static int VarIntSize(int value)
    {
        if (value < 0x80) return 1;
        if (value < 0x4000) return 2;
        if (value < 0x200000) return 3;
        return 4;
    }

    private static int WriteVarInt(Span<byte> buf, int value)
    {
        var pos = 0;
        do
        {
            var b = (byte)(value & 0x7F);
            value >>= 7;
            if (value > 0) b |= 0x80;
            buf[pos++] = b;
        } while (value > 0);
        return pos;
    }
}
