saltylab-ios/SulTee/SulTee/MQTTClient.swift

import Foundation
import Network

/// Minimal MQTT 3.1.1 client — CONNECT + PUBLISH (QoS 0) + SUBSCRIBE (QoS 0) + PINGREQ.
/// Supports both publish and subscribe; no QoS 1/2 needed for this use-case.
final class MQTTClient {

    struct Config {
        var host: String
        var port: UInt16
        var username: String
        var password: String
        var clientID: String
        var keepAlive: UInt16 = 60
    }

    enum State { case disconnected, connecting, connected }

    private(set) var state: State = .disconnected

    /// Called on the main queue for every received PUBLISH message: (topic, payload).
    var onMessage: ((String, String) -> Void)?

    private var config: Config
    private var connection: NWConnection?
    private var pingTimer: DispatchSourceTimer?
    private var shouldRun = false
    private var subscriptions: [String] = []   // persisted across reconnects
    private var nextPacketID: UInt16 = 1
    private let queue = DispatchQueue(label: "mqtt.client", qos: .utility)

    init(config: Config) {
        self.config = config
    }

    // MARK: - Public

    func connect() {
        shouldRun = true
        guard state == .disconnected else { return }
        openConnection()
    }

    func disconnect() {
        shouldRun = false
        pingTimer?.cancel()
        pingTimer = nil
        connection?.cancel()
        connection = nil
        state = .disconnected
    }

    /// Publish a UTF-8 string payload to `topic` at QoS 0.
    func publish(topic: String, payload: String) {
        guard state == .connected else { return }
        connection?.send(content: buildPublish(topic: topic, payload: payload),
                        completion: .idempotent)
    }

    /// Subscribe to `topic` at QoS 0. Stored and re-sent automatically on reconnect.
    func subscribe(topic: String) {
        if !subscriptions.contains(topic) { subscriptions.append(topic) }
        guard state == .connected else { return }
        sendSubscribe(topic: topic)
    }

    // MARK: - Connection lifecycle

    private func openConnection() {
        state = .connecting
        let conn = NWConnection(host: NWEndpoint.Host(config.host),
                               port: NWEndpoint.Port(rawValue: config.port)!,
                               using: .tcp)
        conn.stateUpdateHandler = { [weak self] s in self?.handleStateChange(s) }
        conn.start(queue: queue)
        connection = conn
        scheduleRead()
    }

    private func handleStateChange(_ newState: NWConnection.State) {
        switch newState {
        case .ready:
            sendConnect()
            schedulePing()
        case .failed, .cancelled:
            state = .disconnected
            pingTimer?.cancel()
            pingTimer = nil
            reconnectIfNeeded()
        default:
            break
        }
    }

    private func reconnectIfNeeded() {
        guard shouldRun else { return }
        queue.asyncAfter(deadline: .now() + 3) { [weak self] in self?.openConnection() }
    }

    // MARK: - Read loop

    private func scheduleRead() {
        connection?.receive(minimumIncompleteLength: 2, maximumLength: 4096) { [weak self] data, _, _, error in
            guard let self else { return }
            if let data, !data.isEmpty { self.handleIncoming(data) }
            if error == nil { self.scheduleRead() }
        }
    }

    /// Parse one or more MQTT packets from `data`.
    private func handleIncoming(_ data: Data) {
        var i = data.startIndex

        while i < data.endIndex {
            let firstByte = data[i]
            i = data.index(after: i)
            let packetType = firstByte & 0xF0

            // Decode variable-length remaining-length field
            var multiplier = 1
            var remaining = 0
            var lenByte: UInt8 = 0
            repeat {
                guard i < data.endIndex else { return }
                lenByte = data[i]
                i = data.index(after: i)
                remaining += Int(lenByte & 0x7F) * multiplier
                multiplier *= 128
            } while lenByte & 0x80 != 0

            guard let payloadEnd = data.index(i, offsetBy: remaining, limitedBy: data.endIndex) else { break }

            switch packetType {
            case 0x20:  // CONNACK
                state = .connected
                for topic in subscriptions { sendSubscribe(topic: topic) }

            case 0x30:  // PUBLISH (QoS 0 — no packet identifier)
                var j = i
                if data.distance(from: j, to: payloadEnd) >= 2 {
                    let topicLen = Int(data[j]) << 8 | Int(data[data.index(after: j)])
                    j = data.index(j, offsetBy: 2)
                    if let topicEnd = data.index(j, offsetBy: topicLen, limitedBy: payloadEnd) {
                        let topic = String(bytes: data[j..<topicEnd], encoding: .utf8) ?? ""
                        let payload = String(bytes: data[topicEnd..<payloadEnd], encoding: .utf8) ?? ""
                        let t = topic, p = payload
                        DispatchQueue.main.async { self.onMessage?(t, p) }
                    }
                }

            case 0x90:  // SUBACK — ignore
                break
            case 0xD0:  // PINGRESP — ignore
                break
            default:
                break
            }

            i = payloadEnd
        }
    }

    // MARK: - Keep-alive ping

    private func schedulePing() {
        let t = DispatchSource.makeTimerSource(queue: queue)
        t.schedule(deadline: .now() + Double(config.keepAlive / 2),
                   repeating: Double(config.keepAlive / 2))
        t.setEventHandler { [weak self] in
            self?.connection?.send(content: Data([0xC0, 0x00]), completion: .idempotent)
        }
        t.resume()
        pingTimer = t
    }

    // MARK: - MQTT packet builders

    private func sendConnect() {
        connection?.send(content: buildConnect(), completion: .idempotent)
    }

    private func buildConnect() -> Data {
        var payload = Data()
        payload += mqttString("MQTT")
        payload.append(0x04)                    // protocol level 3.1.1
        payload.append(0xC2)                    // flags: username + password + clean session
        payload += uint16BE(config.keepAlive)
        payload += mqttString(config.clientID)
        payload += mqttString(config.username)
        payload += mqttString(config.password)
        return mqttPacket(type: 0x10, payload: payload)
    }

    private func buildPublish(topic: String, payload: String) -> Data {
        var body = Data()
        body += mqttString(topic)
        body += payload.data(using: .utf8) ?? Data()
        return mqttPacket(type: 0x30, payload: body)
    }

    private func sendSubscribe(topic: String) {
        var payload = Data()
        payload += uint16BE(nextPacketID)
        nextPacketID &+= 1
        payload += mqttString(topic)
        payload.append(0x00)  // QoS 0
        connection?.send(content: mqttPacket(type: 0x82, payload: payload),
                        completion: .idempotent)
    }

    // MARK: - Encoding helpers

    private func mqttPacket(type: UInt8, payload: Data) -> Data {
        var packet = Data([type])
        packet += remainingLength(payload.count)
        packet += payload
        return packet
    }

    private func remainingLength(_ value: Int) -> Data {
        var data = Data()
        var n = value
        repeat {
            var byte = UInt8(n & 0x7F)
            n >>= 7
            if n > 0 { byte |= 0x80 }
            data.append(byte)
        } while n > 0
        return data
    }

    private func mqttString(_ s: String) -> Data {
        let bytes = s.data(using: .utf8) ?? Data()
        return uint16BE(UInt16(bytes.count)) + bytes
    }

    private func uint16BE(_ v: UInt16) -> Data {
        Data([UInt8(v >> 8), UInt8(v & 0xFF)])
    }
}