XuqmGroup-iOSSDK/Sources/XuqmSDK/IM/ImClient.swift

import Foundation

public final class ImClient: NSObject, URLSessionWebSocketDelegate, @unchecked Sendable {

    public weak var delegate: ImEventDelegate?

    private var webSocketTask: URLSessionWebSocketTask?
    private var session: URLSession?
    private var reconnectWorkItem: DispatchWorkItem?
    private var reconnectDelay: TimeInterval = 3
    private var shouldReconnect = true
    private let subscriptionId = "sub-user-queue"
    private var groupSubscriptions = Set<String>()

    private let wsURLOverride: URL?
    private let tokenOverride: String?
    private let appIdOverride: String?

    private var activeWsURL: URL?
    private var activeToken: String?
    private var activeAppId: String?
    private var activeUserId: String?

    public init(wsURL: URL? = nil, token: String? = nil, appId: String? = nil) {
        self.wsURLOverride = wsURL
        self.tokenOverride = token
        self.appIdOverride = appId
        super.init()
    }

    public func setCurrentUserId(_ userId: String) {
        activeUserId = userId
    }

    public func connect() {
        shouldReconnect = true
        reconnectWorkItem?.cancel()
        reconnectWorkItem = nil

        activeWsURL = wsURLOverride
        activeToken = tokenOverride
        activeAppId = appIdOverride

        guard let activeWsURL, let activeToken else {
            delegate?.imClientDidError("IM config or token not found")
            return
        }

        if webSocketTask != nil {
            webSocketTask?.cancel(with: .goingAway, reason: nil)
        }

        let request = URLRequest(url: activeWsURL)
        session = URLSession(configuration: .default, delegate: self, delegateQueue: nil)
        webSocketTask = session?.webSocketTask(with: request)
        webSocketTask?.resume()
        receiveMessage()

        // Keep the token in sync in case a direct caller provided only the stored token.
        self.activeToken = activeToken
    }

    public func sendMessage(toId: String, chatType: ChatType, msgType: MsgType, content: String) -> ImMessage {
        sendMessage(toId: toId, chatType: chatType, msgType: msgType, content: content, mentionedUserIds: nil)
    }

    public func sendMessage(
        toId: String,
        chatType: ChatType,
        msgType: MsgType,
        content: String,
        mentionedUserIds: String?
    ) -> ImMessage {
        let messageId = UUID().uuidString
        let now = Int64(Date().timeIntervalSince1970 * 1000)
        let message = ImMessage(
            id: messageId,
            appId: activeAppId ?? "",
            fromUserId: activeUserId ?? "",
            toId: toId,
            chatType: chatType,
            msgType: msgType,
            content: content,
            status: .sending,
            mentionedUserIds: mentionedUserIds?.isEmpty == false ? mentionedUserIds : nil,
            groupReadCount: nil,
            createdAt: now
        )
        guard let activeAppId else {
            delegate?.imClientDidError("IM appId not configured")
            return message.failedCopy()
        }
        let sent = sendFrame(
            command: "SEND",
            headers: [
                "destination": "/app/chat.send",
                "content-type": "application/json",
            ],
            body: encodeJSONString([
                "appId": activeAppId,
                "messageId": messageId,
                "toId": toId,
                "chatType": chatType.rawValue,
                "msgType": msgType.rawValue,
                "content": content,
                "mentionedUserIds": mentionedUserIds ?? "",
            ]),
        )
        return sent ? message : message.failedCopy()
    }

    public func revoke(messageId: String) {
        guard let activeAppId else {
            delegate?.imClientDidError("IM appId not configured")
            return
        }
        sendFrame(
            command: "SEND",
            headers: [
                "destination": "/app/chat.revoke",
                "content-type": "application/json",
            ],
            body: encodeJSONString([
                "appId": activeAppId,
                "messageId": messageId,
            ]),
        )
    }

    public func subscribeGroup(_ groupId: String) {
        let subscriptionKey = "group-\(groupId)"
        let isNew = groupSubscriptions.insert(groupId).inserted
        if isNew, webSocketTask?.state == .running {
            sendFrame(command: "SUBSCRIBE", headers: [
                "destination": "/topic/group/\(groupId)",
                "id": subscriptionKey,
            ])
        }
    }

    public func unsubscribeGroup(_ groupId: String) {
        groupSubscriptions.remove(groupId)
        if webSocketTask?.state == .running {
            sendFrame(command: "UNSUBSCRIBE", headers: [
                "id": "group-\(groupId)",
            ])
        }
    }

    public func disconnect() {
        shouldReconnect = false
        reconnectWorkItem?.cancel()
        reconnectWorkItem = nil
        webSocketTask?.cancel(with: .normalClosure, reason: nil)
        webSocketTask = nil
        session = nil
    }

    public func isConnected() -> Bool {
        webSocketTask?.state == .running
    }

    private func receiveMessage() {
        webSocketTask?.receive { [weak self] result in
            guard let self else { return }
            switch result {
            case .success(let message):
                if case .string(let text) = message {
                    self.handleMessage(text)
                }
                self.receiveMessage()
            case .failure(let error):
                self.delegate?.imClientDidError(error.localizedDescription)
                self.scheduleReconnect()
            }
        }
    }

    private func handleMessage(_ text: String) {
        guard let data = text.data(using: .utf8) else { return }
        let raw = String(decoding: data, as: UTF8.self)
        for frame in parseFrames(raw) {
            switch frame.command {
            case "CONNECTED":
                reconnectDelay = 3
                sendFrame(command: "SUBSCRIBE", headers: [
                    "destination": "/user/queue/messages",
                    "id": subscriptionId,
                ])
                for groupId in groupSubscriptions {
                    sendFrame(command: "SUBSCRIBE", headers: [
                        "destination": "/topic/group/\(groupId)",
                        "id": "group-\(groupId)",
                    ])
                }
                delegate?.imClientDidConnect()
            case "MESSAGE":
                guard let messageData = frame.body.data(using: .utf8),
                      let msg = try? JSONDecoder().decode(ImMessage.self, from: messageData) else {
                    continue
                }
                if msg.chatType == .group {
                    delegate?.imClientDidReceiveGroupMessage(msg)
                } else {
                    delegate?.imClientDidReceiveMessage(msg)
                }
            case "ERROR":
                delegate?.imClientDidError(frame.body.isEmpty ? "WebSocket error" : frame.body)
            default:
                break
            }
        }
    }

    private func scheduleReconnect() {
        guard shouldReconnect else { return }
        reconnectWorkItem?.cancel()
        let workItem = DispatchWorkItem { [weak self] in
            guard let self else { return }
            self.reconnectDelay = min(self.reconnectDelay * 2, 30)
            self.connect()
        }
        reconnectWorkItem = workItem
        DispatchQueue.global().asyncAfter(deadline: .now() + reconnectDelay, execute: workItem)
    }

    private func sendFrame(command: String, headers: [String: String] = [:], body: String = "") -> Bool {
        guard let webSocketTask, webSocketTask.state == .running else { return false }
        let headerLines = headers.map { "\($0.key):\($0.value)" }.joined(separator: "\n")
        let frame = headerLines.isEmpty
            ? "\(command)\n\n\(body)\u{0000}"
            : "\(command)\n\(headerLines)\n\n\(body)\u{0000}"
        webSocketTask.send(.string(frame)) { _ in }
        return true
    }

    private func encodeJSONString(_ object: [String: String]) -> String {
        guard let data = try? JSONSerialization.data(withJSONObject: object, options: []),
              let text = String(data: data, encoding: .utf8) else {
            return "{}"
        }
        return text
    }

    private func parseFrames(_ raw: String) -> [StompFrame] {
        raw
            .components(separatedBy: "\u{0000}")
            .map { $0.trimmingCharacters(in: .whitespacesAndNewlines) }
            .filter { !$0.isEmpty }
            .map { frame in
                let separator = frame.range(of: "\n\n")
                let headerBlock = separator.map { String(frame[..<$0.lowerBound]) } ?? frame
                let body = separator.map { String(frame[$0.upperBound...]) } ?? ""
                let lines = headerBlock.split(separator: "\n").map(String.init)
                let command = lines.first ?? ""
                let headers: [String: String] = Dictionary(uniqueKeysWithValues: lines.dropFirst().compactMap { line in
                    guard let index = line.firstIndex(of: ":") else { return nil }
                    let key = String(line[..<index])
                    let value = String(line[line.index(after: index)...])
                    return (key, value)
                })
                return StompFrame(command: command, headers: headers, body: body)
            }
    }

    public func urlSession(_ session: URLSession, webSocketTask: URLSessionWebSocketTask,
                           didOpenWithProtocol protocol: String?) {
        guard let token = activeToken else {
            delegate?.imClientDidError("IM token not found")
            return
        }
        activeToken = token
        sendFrame(command: "CONNECT", headers: [
            "accept-version": "1.2",
            "Authorization": "Bearer \(token)",
            "heart-beat": "10000,10000",
        ])
    }

    public func urlSession(_ session: URLSession, webSocketTask: URLSessionWebSocketTask,
                           didCloseWith closeCode: URLSessionWebSocketTask.CloseCode, reason: Data?) {
        let reasonStr = reason.flatMap { String(data: $0, encoding: .utf8) }
        delegate?.imClientDidDisconnect(reason: reasonStr)
        scheduleReconnect()
    }
}

private extension ImMessage {
    func failedCopy() -> ImMessage {
        ImMessage(
            id: id,
            appId: appId,
            fromUserId: fromUserId,
            toId: toId,
            chatType: chatType,
            msgType: msgType,
            content: content,
            status: .failed,
            mentionedUserIds: mentionedUserIds,
            groupReadCount: groupReadCount,
            createdAt: createdAt
        )
    }
}

private struct StompFrame {
    let command: String
    let headers: [String: String]
    let body: String
}