import Foundation
import CryptoKit
import UIKit

/// Notification types from the server
enum ChatNotification {
    case newMessage(data: [String: Any])
    case messagesRead(data: [String: Any])
    case messageDeleted(data: [String: Any])
    case conversationCreated(data: [String: Any])
    case memberAdded(data: [String: Any])
    case memberRemoved(data: [String: Any])
    case userOnline(userId: String)
    case userOffline(userId: String)
    case onlineUsers(userIds: [String])
    case groupInvitation(data: [String: Any])
    case conversationRenamed(data: [String: Any])
    case sessionReset(data: [String: Any])
    case messageReacted(data: [String: Any])
    case messagePinned(data: [String: Any])
    case messageUnpinned(data: [String: Any])
    case messageDelivered(data: [String: Any])
    case conversationDeleted(data: [String: Any])
    case connectionStateChanged(connected: Bool)
    case reconnected
}

/// Result of a reconnect attempt
enum ReconnectResult {
    case success
    case authFailed    // Keys invalid (e.g. rotated) — should logout immediately
    case networkError  // TCP/connection issue — can retry
}

/// Main chat client — handles all server communication and crypto operations.
/// Thread-safe via Swift actor isolation.
/// Port of Python ChatClient class from chat_core.py
actor ChatClient {

    // MARK: - Connection

    let connectionManager = ConnectionManager()
    private(set) var isConnected = false
    private var lastHost = Constants.defaultHost
    private var lastPort = Constants.defaultPort
    private(set) var sessionToken: String?
    private(set) var userId: String?
    private(set) var username: String = ""
    private(set) var email: String = ""
    private(set) var loginRejected = false
    private(set) var onlineUserIds: Set<String> = []

    // MARK: - Keys

    private var rsaPrivate: SecKey?
    private var rsaPublic: SecKey?
    private(set) var identityPrivate: Curve25519.Signing.PrivateKey?
    private(set) var identityPublic: Curve25519.Signing.PublicKey?
    private var spkPrivate: Curve25519.KeyAgreement.PrivateKey?
    private var spkId: String = ""
    private var prevSpkPrivate: Curve25519.KeyAgreement.PrivateKey?
    private var prevSpkId: String = ""
    private var opkPrivates: [String: Curve25519.KeyAgreement.PrivateKey] = [:]

    // MARK: - Sessions & Sender Keys

    private var sessions: [String: DoubleRatchet] = [:]         // "userId:deviceId" -> ratchet
    private var senderKeyStates: [String: SenderKeyState] = [:] // convId -> own sender key
    private var recvSenderKeys: [String: SenderKeyState] = [:]  // "convId:senderId:deviceId" -> their key

    // MARK: - Derived Keys

    private(set) var cacheKey: Data?   // for encrypting message cache
    private(set) var localKey: Data?   // for encrypting session/sender key files

    // MARK: - Multi-Device

    private(set) var deviceId: String?

    // MARK: - Caches

    private var userCache: [String: User] = [:]
    private var deviceBundleCache: [String: (timestamp: Date, bundles: [DeviceBundle])] = [:]

    // MARK: - Self-Encrypt Queue

    private var pendingSelfEncrypt: [(messageId: String, plaintext: Data)] = []

    // MARK: - TOFU / Contact Verification

    private var knownIdentityKeys: [String: [String: String]] = [:]  // userId -> {identity_key, first_seen, last_seen}
    private var verifiedContacts: [String: [String: String]] = [:]   // userId -> {identity_key, verified_at, method}

    // MARK: - Brute-Force Lockout

    private static let lockoutBaseSeconds: TimeInterval = 2
    private static let lockoutMaxSeconds: TimeInterval = 300

    // MARK: - Request/Response Tracking

    private var pendingRequests: [String: CheckedContinuation<[String: Any], Error>] = [:]
    private var listenerTask: Task<Void, Never>?

    // MARK: - Notification Broadcast (supports multiple consumers)

    private var notificationSubscribers: [UUID: AsyncStream<ChatNotification>.Continuation] = [:]

    /// Create a new notification stream. Each subscriber gets ALL notifications.
    func makeNotificationStream() -> AsyncStream<ChatNotification> {
        let id = UUID()
        var captured: AsyncStream<ChatNotification>.Continuation!
        let stream = AsyncStream<ChatNotification> { cont in
            captured = cont
        }
        captured.onTermination = { [weak self] _ in
            Task { await self?.removeSubscriber(id: id) }
        }
        notificationSubscribers[id] = captured
        return stream
    }

    private func removeSubscriber(id: UUID) {
        notificationSubscribers.removeValue(forKey: id)
    }

    private func broadcastNotification(_ notification: ChatNotification) {
        for (_, continuation) in notificationSubscribers {
            continuation.yield(notification)
        }
    }

    // MARK: - Init

    init() {}

    // MARK: - Connection

    func connect(host: String = Constants.defaultHost, port: UInt16 = Constants.defaultPort) async throws {
        try await connectionManager.connect(host: host, port: port)
        lastHost = host
        lastPort = port
        isConnected = true
        startBackgroundListener()
        broadcastNotification(.connectionStateChanged(connected: true))
    }

    /// Check if the underlying TCP connection is actually alive
    func isConnectionAlive() async -> Bool {
        guard isConnected else { return false }
        return await connectionManager.isConnected
    }

    func disconnect() async {
        listenerTask?.cancel()
        listenerTask = nil
        await connectionManager.disconnect()
        isConnected = false
        // Fail all pending requests
        let pending = pendingRequests
        pendingRequests.removeAll()
        for (_, cont) in pending {
            cont.resume(throwing: NetworkError.notConnected)
        }
        broadcastNotification(.connectionStateChanged(connected: false))
    }

    // MARK: - Send and Receive

    /// Send a request and wait for the matching response.
    func sendAndReceive(type: String, timeout: TimeInterval = 30, params: [String: Any] = [:]) async -> [String: Any] {
        let requestId = ProtocolHandler.newRequestId()
        #if DEBUG
        print("DEBUG sendAndReceive: type=\(type), requestId=\(requestId)")
        #endif

        // Timeout task — resumes continuation with error if server doesn't respond
        let timeoutTask = Task {
            try await Task.sleep(nanoseconds: UInt64(timeout * 1_000_000_000))
            if let cont = pendingRequests.removeValue(forKey: requestId) {
                cont.resume(throwing: NetworkError.timeout)
            }
        }

        do {
            let response: [String: Any] = try await withCheckedThrowingContinuation { continuation in
                pendingRequests[requestId] = continuation

                Task {
                    do {
                        try await connectionManager.sendMessage(type: type, requestId: requestId, params: params)
                    } catch {
                        if let cont = pendingRequests.removeValue(forKey: requestId) {
                            cont.resume(throwing: error)
                        }
                    }
                }
            }
            timeoutTask.cancel()
            return response
        } catch {
            timeoutTask.cancel()
            pendingRequests.removeValue(forKey: requestId)
            return [
                "type": type,
                "status": "error",
                "data": ["message": error.localizedDescription]
            ]
        }
    }

    // MARK: - Background Listener

    func startBackgroundListener() {
        listenerTask?.cancel()
        listenerTask = Task { [weak self] in
            guard let self = self else { return }
            await self.backgroundListenerLoop()
        }
    }

    private func backgroundListenerLoop() async {
        while !Task.isCancelled {
            do {
                guard let msg = try await connectionManager.readMessage() else {
                    // EOF — connection closed
                    handleDisconnect()
                    break
                }
                routeMessage(msg)
            } catch {
                handleDisconnect()
                break
            }
        }
    }

    private func handleDisconnect() {
        isConnected = false
        // Fail all pending futures
        let pending = pendingRequests
        pendingRequests.removeAll()
        for (_, cont) in pending {
            cont.resume(throwing: NetworkError.notConnected)
        }
        broadcastNotification(.connectionStateChanged(connected: false))
    }

    private func routeMessage(_ msg: [String: Any]) {
        #if DEBUG
        print("DEBUG routeMessage received: \(msg)")
        #endif
        let msgType = msg["type"] as? String ?? ""

        // Notification types (no request_id expected from client)
        let notificationTypes = Set([
            "new_message", "messages_read", "message_deleted",
            "conversation_created", "member_added", "member_removed",
            "user_online", "user_offline", "online_users",
            "group_invitation", "conversation_renamed", "session_reset",
            "keys_updated", "message_reacted", "message_pinned", "message_unpinned",
            "message_delivered", "conversation_deleted"
        ])

        if notificationTypes.contains(msgType) {
            let data = msg["data"] as? [String: Any] ?? msg
            switch msgType {
            case "new_message":
                broadcastNotification(.newMessage(data: data))
            case "messages_read":
                broadcastNotification(.messagesRead(data: data))
            case "message_deleted":
                broadcastNotification(.messageDeleted(data: data))
            case "conversation_created":
                broadcastNotification(.conversationCreated(data: data))
            case "member_added":
                broadcastNotification(.memberAdded(data: data))
            case "member_removed":
                broadcastNotification(.memberRemoved(data: data))
            case "user_online":
                if let uid = data["user_id"] as? String {
                    onlineUserIds.insert(uid)
                    broadcastNotification(.userOnline(userId: uid))
                }
            case "user_offline":
                if let uid = data["user_id"] as? String {
                    onlineUserIds.remove(uid)
                    broadcastNotification(.userOffline(userId: uid))
                }
            case "online_users":
                if let uids = data["user_ids"] as? [String] {
                    onlineUserIds = Set(uids)
                    broadcastNotification(.onlineUsers(userIds: uids))
                }
            case "group_invitation":
                broadcastNotification(.groupInvitation(data: data))
            case "conversation_renamed":
                broadcastNotification(.conversationRenamed(data: data))
            case "session_reset":
                broadcastNotification(.sessionReset(data: data))
            case "message_reacted":
                broadcastNotification(.messageReacted(data: data))
            case "message_pinned":
                broadcastNotification(.messagePinned(data: data))
            case "message_unpinned":
                broadcastNotification(.messageUnpinned(data: data))
            case "message_delivered":
                broadcastNotification(.messageDelivered(data: data))
            case "conversation_deleted":
                broadcastNotification(.conversationDeleted(data: data))
            case "keys_updated":
                // Peer uploaded new prekeys (new device or rotation) — invalidate bundle cache
                if let uid = data["user_id"] as? String {
                    deviceBundleCache.removeValue(forKey: uid)
                    #if DEBUG
                    print("DEBUG keys_updated: invalidated bundle cache for \(uid)")
                    #endif
                }
            default:
                break
            }
        } else {
            // Response to a pending request
            if let requestId = msg["request_id"] as? String,
               let cont = pendingRequests.removeValue(forKey: requestId) {
                cont.resume(returning: msg)
            }
        }
    }

    // MARK: - User Info Cache

    func getUserInfo(userId: String = "", userEmail: String = "") async -> User? {
        if !userId.isEmpty, let cached = userCache[userId] {
            return cached
        }
        var params: [String: Any] = [:]
        if !userId.isEmpty { params["user_id"] = userId }
        else if !userEmail.isEmpty { params["email"] = userEmail }
        else { return nil }

        let resp = await sendAndReceive(type: "get_user_info", params: params)
        guard resp.string(for: "status") == "ok",
              let data = resp.dict(for: "data") else { return nil }

        var ikData: Data?
        if let ikB64 = data["identity_key"] as? String {
            ikData = try? ProtocolHandler.decodeBinary(ikB64)
        }

        let user = User(
            id: data.string(for: "user_id") ?? "",
            username: data.string(for: "username") ?? "",
            email: data.string(for: "email") ?? "",
            identityKey: ikData
        )
        userCache[user.id] = user
        return user
    }

    // MARK: - Registration

    func register(username: String, password: String, email: String) async -> (success: Bool, message: String) {
        self.username = username
        self.email = email
        var pwdBytes = Array(password.utf8)
        defer { pwdBytes.withUnsafeMutableBytes { ptr in _ = memset(ptr.baseAddress!, 0, ptr.count) } }

        let pwdData = Data(pwdBytes)

        do {
            // RSA keys
            let (rsaPriv, rsaPub, _) = KeyStorage.loadRSAKeys(email: email, password: pwdData)
            if let rsaPriv = rsaPriv, let rsaPub = rsaPub {
                self.rsaPrivate = rsaPriv
                self.rsaPublic = rsaPub
            } else {
                let (newPriv, newPub) = try RSACrypto.generateKeypair()
                try KeyStorage.saveRSAKeys(email: email, privateKey: newPriv, publicKey: newPub, password: pwdData)
                self.rsaPrivate = newPriv
                self.rsaPublic = newPub
            }

            // Ed25519 identity keys
            let (edPriv, edPub) = KeyStorage.loadIdentityKeys(email: email, password: pwdData)
            if let edPriv = edPriv, let edPub = edPub {
                self.identityPrivate = edPriv
                self.identityPublic = edPub
            } else {
                let (newPriv, newPub) = Ed25519Crypto.generateKeypair()
                try KeyStorage.saveIdentityKeys(email: email, privateKey: newPriv, publicKey: newPub, password: pwdData)
                self.identityPrivate = newPriv
                self.identityPublic = newPub
            }

            self.cacheKey = CryptoUtils.deriveSelfEncryptionKey(identityPrivateRaw: identityPrivate!.rawData)
            self.localKey = CryptoUtils.deriveLocalStorageKey(identityPrivateRaw: identityPrivate!.rawData)
        } catch {
            return (false, "Key generation failed: \(error.localizedDescription)")
        }

        // Send registration request
        let pubPem = String(data: try! RSACrypto.serializePublicKey(rsaPublic!), encoding: .utf8)!
        let ikB64 = ProtocolHandler.encodeBinary(Ed25519Crypto.serializePublic(identityPublic!))

        let resp = await sendAndReceive(type: "register", params: [
            "username": username,
            "public_key": pubPem,
            "email": email,
            "identity_key": ikB64,
        ])

        #if DEBUG
        print("DEBUG register response: \(resp)")
        #endif

        // Handle PoW challenge if required by server
        if resp.string(for: "status") == "pow_required",
           let powData = resp.dict(for: "data"),
           let powChallenge = powData.string(for: "challenge"),
           let powMac = powData.string(for: "mac"),
           let powDifficulty = powData.int(for: "difficulty") {
            #if DEBUG
            print("DEBUG register: PoW required, difficulty=\(powDifficulty)")
            #endif
            let powNonce = ChatClient.solvePow(challenge: powChallenge, difficulty: powDifficulty)
            #if DEBUG
            print("DEBUG register: PoW solved, nonce=\(powNonce)")
            #endif

            // Retry with PoW solution
            let retryResp = await sendAndReceive(type: "register", params: [
                "username": username,
                "public_key": pubPem,
                "email": email,
                "identity_key": ikB64,
                "pow_challenge": powChallenge,
                "pow_mac": powMac,
                "pow_nonce": powNonce,
            ])

            guard retryResp.string(for: "status") == "ok" else {
                let msg = retryResp.dict(for: "data")?.string(for: "message") ?? "Registration failed after PoW"
                return (false, msg)
            }
            let retryData = retryResp.dict(for: "data") ?? [:]
            if let code = retryData.string(for: "code") {
                return (true, code)
            }
            return (true, retryData.string(for: "message") ?? "Check your email for the code.")
        }

        guard resp.string(for: "status") == "ok" else {
            let msg = resp.dict(for: "data")?.string(for: "message") ?? "Registration failed"
            return (false, msg)
        }

        let data = resp.dict(for: "data") ?? [:]
        if let code = data.string(for: "code") {
            return (true, code)
        }
        return (true, data.string(for: "message") ?? "Check your email for the code.")
    }

    func confirmRegistration(email: String, username: String, code: String) async -> (success: Bool, message: String) {
        #if DEBUG
        print("DEBUG confirmRegistration: email=\(email), code=\(code)")
        #endif
        let resp = await sendAndReceive(type: "register_confirm", params: [
            "email": email,
            "code": code,
        ])

        guard resp.string(for: "status") == "ok" else {
            let msg = resp.dict(for: "data")?.string(for: "message") ?? "Confirmation failed"
            return (false, msg)
        }

        // NOTE: Don't upload prekeys here - user isn't logged in yet.
        // ensurePrekeys() will be called after login.

        let uid = resp.dict(for: "data")?.string(for: "user_id") ?? ""
        return (true, "Registered as '\(username)' (ID: \(uid))")
    }

    // MARK: - Prekeys

    private func generateAndUploadPrekeys(keepSPK: Bool = false) async {
        guard let identityPrivate = identityPrivate else { return }

        do {
            let spkData: [String: Any]

            if keepSPK, let spkPriv = spkPrivate, !spkId.isEmpty {
                let spkPubBytes = X25519Crypto.serializePublic(spkPriv.publicKey)
                let sig = try Ed25519Crypto.sign(identityPrivate, data: spkPubBytes)
                spkData = [
                    "id": spkId,
                    "public_key": ProtocolHandler.encodeBinary(spkPubBytes),
                    "signature": ProtocolHandler.encodeBinary(sig),
                ]
            } else {
                // Save current as previous (grace period)
                if let spkPriv = spkPrivate, !spkId.isEmpty {
                    prevSpkPrivate = spkPriv
                    prevSpkId = spkId
                    try? KeyStorage.savePrevSPK(email: email, privateKey: spkPriv, spkId: spkId)
                }

                let spk = try X3DH.generateSignedPrekey(identityPrivate: identityPrivate)
                self.spkPrivate = spk.privateKey
                self.spkId = spk.id
                try? KeyStorage.saveSPK(email: email, privateKey: spk.privateKey, spkId: spk.id)

                spkData = [
                    "id": spk.id,
                    "public_key": ProtocolHandler.encodeBinary(X25519Crypto.serializePublic(spk.publicKey)),
                    "signature": ProtocolHandler.encodeBinary(spk.signature),
                ]
            }

            // Generate OPKs
            let opks = X3DH.generateOneTimePrekeys(count: Constants.opkBatchSize)
            #if DEBUG
            print("DEBUG generatePrekeys: generating \(opks.count) OPKs")
            #endif
            for opk in opks {
                opkPrivates[opk.id] = opk.privateKey
                do {
                    try KeyStorage.saveOPKPrivate(email: email, opkId: opk.id, privateKey: opk.privateKey)
                } catch {
                    #if DEBUG
                    print("DEBUG generatePrekeys: FAILED to save OPK \(opk.id): \(error)")
                    #endif
                }
            }
            #if DEBUG
            print("DEBUG generatePrekeys: OPK IDs = \(opks.map { $0.id })")
            #endif

            let otpData = opks.map { opk -> [String: Any] in
                [
                    "id": opk.id,
                    "public_key": ProtocolHandler.encodeBinary(X25519Crypto.serializePublic(opk.publicKey)),
                ]
            }

            let uploadResp = await sendAndReceive(type: "ensure_prekeys", params: [
                "signed_prekey": spkData,
                "one_time_prekeys": otpData,
            ])
            let uploadStatus = uploadResp.string(for: "status") ?? "nil"
            #if DEBUG
            print("DEBUG generatePrekeys: ensure_prekeys response status=\(uploadStatus) deviceId=\(deviceId ?? "nil")")
            #endif
            if uploadStatus != "ok" {
                let errMsg = uploadResp.dict(for: "data")?.string(for: "message") ?? "unknown"
                #if DEBUG
                print("DEBUG generatePrekeys: upload FAILED: \(errMsg)")
                #endif
            }
        } catch {
            // Log error but don't fail
            #if DEBUG
            print("Prekey generation error: \(error)")
            #endif
        }
    }

    private func ensurePrekeys() async {
        let resp = await sendAndReceive(type: "get_prekey_count")
        guard resp.string(for: "status") == "ok",
              let data = resp.dict(for: "data") else { return }

        let count = data.int(for: "count") ?? 0
        let spkCreatedAt = data.string(for: "spk_created_at") ?? ""

        var needNewSPK = false
        if !spkCreatedAt.isEmpty {
            if let created = DateParsing.parse(spkCreatedAt) {
                let ageDays = Calendar.current.dateComponents([.day], from: created, to: Date()).day ?? 0
                if ageDays >= Constants.spkRotationDays {
                    needNewSPK = true
                }
            }
        }

        if count < Constants.opkReplenishThreshold || needNewSPK {
            await generateAndUploadPrekeys()
        }
    }

    // MARK: - Login

    func login(email: String, password: String) async -> (success: Bool, message: String) {
        self.email = email

        // Check brute-force lockout
        let lockoutRemaining = ChatClient.checkLockout(email: email)
        if lockoutRemaining > 0 {
            return (false, "Too many failed attempts. Wait \(Int(lockoutRemaining))s.")
        }

        var pwdBytes = Array(password.utf8)
        defer { pwdBytes.withUnsafeMutableBytes { ptr in _ = memset(ptr.baseAddress!, 0, ptr.count) } }
        let pwdData = Data(pwdBytes)

        // Load RSA keys
        let (rsaPriv, rsaPub, err) = KeyStorage.loadRSAKeys(email: email, password: pwdData)
        guard let rsaPriv = rsaPriv, let rsaPub = rsaPub else {
            return (false, err ?? "No local keys found. Register first.")
        }
        self.rsaPrivate = rsaPriv
        self.rsaPublic = rsaPub

        // Load identity keys
        let (edPriv, edPub) = KeyStorage.loadIdentityKeys(email: email, password: pwdData)
        if let edPriv = edPriv, let edPub = edPub {
            self.identityPrivate = edPriv
            self.identityPublic = edPub
            self.cacheKey = CryptoUtils.deriveSelfEncryptionKey(identityPrivateRaw: edPriv.rawData)
            self.localKey = CryptoUtils.deriveLocalStorageKey(identityPrivateRaw: edPriv.rawData)
        }

        // Load SPK
        let (spkP, spkI) = KeyStorage.loadSPK(email: email)
        if let spkP = spkP {
            self.spkPrivate = spkP
            self.spkId = spkI ?? ""
        }

        // Load previous SPK (grace period)
        let (prevP, prevI) = KeyStorage.loadPrevSPK(email: email)
        if let prevP = prevP {
            self.prevSpkPrivate = prevP
            self.prevSpkId = prevI ?? ""
        }

        // Load device ID
        self.deviceId = KeyStorage.loadDeviceId(email: email)

        // RSA challenge-response login
        let startResp = await sendAndReceive(type: "login_start", params: ["email": email])
        guard startResp.string(for: "status") == "ok",
              let startData = startResp.dict(for: "data"),
              let challengeB64 = startData.string(for: "challenge") else {
            let msg = startResp.dict(for: "data")?.string(for: "message") ?? "Login failed"
            return (false, msg)
        }

        let challengeData: Data
        do {
            challengeData = try ProtocolHandler.decodeBinary(challengeB64)
        } catch {
            return (false, "Invalid challenge data")
        }

        let signature: Data
        do {
            signature = try RSACrypto.sign(rsaPriv, data: challengeData)
        } catch {
            return (false, "RSA signing failed: \(error.localizedDescription)")
        }

        var finishParams: [String: Any] = [
            "email": email,
            "signature": ProtocolHandler.encodeBinary(signature),
            "client_version": Constants.version,
        ]
        if let deviceId = deviceId {
            finishParams["device_id"] = deviceId
        }

        let finishResp = await sendAndReceive(type: "login_finish", params: finishParams)
        guard finishResp.string(for: "status") == "ok",
              let finishData = finishResp.dict(for: "data") else {
            let msg = finishResp.dict(for: "data")?.string(for: "message") ?? "Login failed"
            loginRejected = true
            ChatClient.recordFailedAttempt(email: email)
            return (false, msg)
        }

        self.userId = finishData.string(for: "user_id")
        self.username = finishData.string(for: "username") ?? ""
        self.sessionToken = finishData.string(for: "session_token")

        // Save device ID from server
        if let newDeviceId = finishData.string(for: "device_id") {
            self.deviceId = newDeviceId
            try? KeyStorage.saveDeviceId(email: email, deviceId: newDeviceId)
        }

        // Handle online_users if included
        if let onlineUserIds = finishData["online_user_ids"] as? [String] {
            self.onlineUserIds = Set(onlineUserIds)
            broadcastNotification(.onlineUsers(userIds: onlineUserIds))
        }

        // Check if we have local OPK private keys.
        // After device pairing, the new device has no local OPKs — must generate fresh ones.
        let hasLocalOPKs: Bool
        if let dir = try? KeyStorage.getKeyDir(email: email) {
            let opkDir = dir.appendingPathComponent("opk_private")
            let files = (try? FileManager.default.contentsOfDirectory(atPath: opkDir.path)) ?? []
            hasLocalOPKs = !files.isEmpty
        } else {
            hasLocalOPKs = false
        }

        if hasLocalOPKs {
            // Existing device — check/replenish prekeys in background
            Task { await ensurePrekeys() }
        } else {
            // New device — MUST upload prekeys before returning so other clients
            // can encrypt for this device immediately. Fire-and-forget would create
            // a race where senders fetch bundles before prekeys exist on server.
            #if DEBUG
            print("DEBUG login: no local OPKs (likely new device). Generating fresh prekeys (synchronous).")
            #endif
            await generateAndUploadPrekeys(keepSPK: true)
        }

        // Load previous SPK for grace period (M4)
        let (prevSPK, prevSPKId) = KeyStorage.loadPrevSPK(email: email)
        if let prevSPK = prevSPK {
            prevSpkPrivate = prevSPK
            prevSpkId = prevSPKId ?? ""
        }

        // Load TOFU and verification stores
        loadVerificationStores()

        // Clear lockout on successful login
        ChatClient.clearLockout(email: email)

        return (true, "Logged in as \(username)")
    }

    // MARK: - Reconnect

    func reconnect() async -> ReconnectResult {
        guard rsaPrivate != nil else { return .authFailed }

        let host = lastHost
        let port = lastPort

        await disconnect()

        do {
            try await connect(host: host, port: port)
        } catch {
            return .networkError
        }

        // TCP connected — try RSA challenge-response with in-memory keys
        let startResp = await sendAndReceive(type: "login_start", params: ["email": email])
        guard startResp.string(for: "status") == "ok",
              let startData = startResp.dict(for: "data"),
              let challengeB64 = startData.string(for: "challenge"),
              let challengeData = try? ProtocolHandler.decodeBinary(challengeB64),
              let signature = try? RSACrypto.sign(rsaPrivate!, data: challengeData) else {
            // TCP was fine but login_start failed — auth issue
            return .authFailed
        }

        var finishParams: [String: Any] = [
            "email": email,
            "signature": ProtocolHandler.encodeBinary(signature),
            "client_version": Constants.version,
        ]
        if let deviceId = deviceId {
            finishParams["device_id"] = deviceId
        }

        let finishResp = await sendAndReceive(type: "login_finish", params: finishParams)
        guard finishResp.string(for: "status") == "ok",
              let finishData = finishResp.dict(for: "data") else {
            // TCP connected, got challenge, but signature rejected — keys rotated
            return .authFailed
        }

        // Update session token if server returned a new one
        if let token = finishData.string(for: "session_token") {
            self.sessionToken = token
        }

        // Refresh online users (matches Python: self.session = finish["data"])
        if let onlineUserIds = finishData["online_user_ids"] as? [String] {
            self.onlineUserIds = Set(onlineUserIds)
            broadcastNotification(.onlineUsers(userIds: onlineUserIds))
        }

        // Replenish prekeys if needed (matches Python: asyncio.create_task(self._ensure_prekeys()))
        Task { await ensurePrekeys() }

        broadcastNotification(.reconnected)
        return .success
    }

    // MARK: - Device Bundles

    private func getDeviceBundles(userId: String) async throws -> [DeviceBundle] {
        // Check cache (5-min TTL)
        if let cached = deviceBundleCache[userId],
           Date().timeIntervalSince(cached.timestamp) < Constants.deviceBundleCacheTTL {
            return cached.bundles
        }

        let resp = await sendAndReceive(type: "get_key_bundle", params: ["user_id": userId])
        guard resp.string(for: "status") == "ok",
              let data = resp.dict(for: "data") else {
            throw ChatError.operationFailed("Failed to get key bundle")
        }

        var bundles: [DeviceBundle] = []

        // Get identity key from top level (shared across all device bundles)
        var identityKey: Data?
        if let ikB64 = data["identity_key"] as? String {
            identityKey = Data(base64Encoded: ikB64)
        }

        // Per-device bundles (new format)
        if let deviceBundlesRaw = data["device_bundles"] as? [[String: Any]] {
            for bundleDict in deviceBundlesRaw {
                if let bundle = try? DeviceBundle.fromDict(bundleDict, identityKey: identityKey) {
                    bundles.append(bundle)
                }
            }
        }
        // Legacy single bundle
        else if identityKey != nil {
            let bundle = try DeviceBundle.fromDict(data, identityKey: identityKey)
            bundles.append(bundle)
        }

        deviceBundleCache[userId] = (Date(), bundles)
        return bundles
    }

    // MARK: - Session Management

    private func getOrCreateSession(
        peerUserId: String,
        peerDeviceId: String,
        bundle: DeviceBundle
    ) async throws -> DoubleRatchet {
        let sessionKey = "\(peerUserId):\(peerDeviceId)"
        #if DEBUG
        print("DEBUG getOrCreateSession: looking for session \(sessionKey)")
        #endif

        // Check memory
        if let session = sessions[sessionKey] {
            #if DEBUG
            print("DEBUG getOrCreateSession: found in memory")
            #endif
            return session
        }

        // Check disk
        if let session = KeyStorage.loadSession(
            email: email,
            peerUserId: peerUserId,
            localKey: localKey,
            peerDeviceId: peerDeviceId
        ) {
            #if DEBUG
            print("DEBUG getOrCreateSession: found on disk")
            #endif
            sessions[sessionKey] = session
            return session
        }

        #if DEBUG
        print("DEBUG getOrCreateSession: creating new session via X3DH")
        #endif

        // Create new via X3DH
        let remoteIkEd = try Ed25519Crypto.loadPublic(bundle.identityKey)
        let spkRemote = try X25519Crypto.loadPublic(bundle.spk)
        var opkRemote: Curve25519.KeyAgreement.PublicKey?
        if let opkData = bundle.opk {
            opkRemote = try X25519Crypto.loadPublic(opkData)
        }

        let (sharedSecret, _, ekPub) = try X3DH.initiate(
            ikPrivateEd: identityPrivate!,
            ikPublicRemoteEd: remoteIkEd,
            spkRemote: spkRemote,
            spkSignature: bundle.spkSignature,
            opkRemote: opkRemote
        )

        let ratchet = try DoubleRatchet.initAlice(sharedSecret: sharedSecret, bobSpkPub: spkRemote)

        // Build X3DH header for first message
        // Keys must be base64 encoded (not hex) to match Python server
        let myIkBytes = Ed25519Crypto.serializePublic(identityPublic!)
        let ekPubBytes = X25519Crypto.serializePublic(ekPub)
        #if DEBUG
        print("DEBUG getOrCreateSession: sending ik len=\(myIkBytes.count)")
        print("DEBUG getOrCreateSession: sending ek len=\(ekPubBytes.count)")
        #endif
        var x3dhHeader: [String: Any] = [
            "ik": ProtocolHandler.encodeBinary(myIkBytes),
            "ek": ProtocolHandler.encodeBinary(ekPubBytes),
            "spk_id": bundle.spkId,
        ]
        if let opkId = bundle.opkId {
            x3dhHeader["opk_id"] = opkId
            #if DEBUG
            print("DEBUG getOrCreateSession: using opk_id = \(opkId)")
            #endif
        }
        ratchet.x3dhHeader = x3dhHeader

        sessions[sessionKey] = ratchet
        try? KeyStorage.saveSession(email: email, peerUserId: peerUserId, ratchet: ratchet, localKey: localKey, peerDeviceId: peerDeviceId)

        return ratchet
    }

    // MARK: - X3DH Response (Bob Side)

    private func processX3DHHeader(
        senderId: String,
        x3dhHeader: [String: Any],
        senderDeviceId: String,
        spkOverride: Curve25519.KeyAgreement.PrivateKey? = nil
    ) throws -> DoubleRatchet {
        guard let ikB64 = x3dhHeader["ik"] as? String,
              let ikData = try? ProtocolHandler.decodeBinary(ikB64),
              let ekB64 = x3dhHeader["ek"] as? String,
              let ekData = try? ProtocolHandler.decodeBinary(ekB64) else {
            throw CryptoError.x3dhFailed("Invalid X3DH header - missing ik or ek")
        }

        // spk_id is optional - if missing, use current SPK
        let spkIdStr = x3dhHeader["spk_id"] as? String
        #if DEBUG
        print("DEBUG processX3DHHeader: ik=\(ikData.count)B, ek=\(ekData.count)B, spk_id=\(spkIdStr ?? "nil")")
        #endif

        let remoteIkEd = try Ed25519Crypto.loadPublic(ikData)
        let ekRemote = try X25519Crypto.loadPublic(ekData)

        // Determine which SPK to use
        let spkToUse: Curve25519.KeyAgreement.PrivateKey
        if let override = spkOverride {
            #if DEBUG
            print("DEBUG processX3DHHeader: using spkOverride")
            #endif
            spkToUse = override
        } else if let spkIdStr = spkIdStr {
            // spk_id provided - match against known SPKs
            #if DEBUG
            print("DEBUG processX3DHHeader: spk_id provided, mySpkId=\(spkId), prevSpkId=\(prevSpkId)")
            #endif
            if spkIdStr == spkId, let spk = spkPrivate {
                spkToUse = spk
            } else if spkIdStr == prevSpkId, let prevSpk = prevSpkPrivate {
                spkToUse = prevSpk
            } else {
                throw CryptoError.x3dhFailed("SPK \(spkIdStr) not found")
            }
        } else {
            // No spk_id provided - use current SPK (for backward compatibility)
            #if DEBUG
            print("DEBUG processX3DHHeader: no spk_id, using current SPK")
            #endif
            guard let spk = spkPrivate else {
                throw CryptoError.x3dhFailed("No SPK available")
            }
            spkToUse = spk
        }

        // OPK
        var opkPriv: Curve25519.KeyAgreement.PrivateKey?
        if let opkIdStr = x3dhHeader["opk_id"] as? String {
            #if DEBUG
            print("DEBUG processX3DHHeader: looking for OPK \(opkIdStr)")
            print("DEBUG processX3DHHeader: opkPrivates has \(opkPrivates.count) keys: \(Array(opkPrivates.keys).prefix(5))...")
            #endif

            // Check file system
            if let dir = try? KeyStorage.getKeyDir(email: email) {
                let opkDir = dir.appendingPathComponent("opk_private")
                let opkFile = opkDir.appendingPathComponent("\(opkIdStr).bin")
                #if DEBUG
                print("DEBUG processX3DHHeader: checking file \(opkFile.path)")
                print("DEBUG processX3DHHeader: file exists = \(FileManager.default.fileExists(atPath: opkFile.path))")
                #endif
                if let files = try? FileManager.default.contentsOfDirectory(atPath: opkDir.path) {
                    #if DEBUG
                    print("DEBUG processX3DHHeader: opk_private dir has \(files.count) files: \(files.prefix(5))...")
                    #endif
                }
            }

            opkPriv = opkPrivates[opkIdStr] ?? KeyStorage.loadOPKPrivate(email: email, opkId: opkIdStr)
            if opkPriv != nil {
                #if DEBUG
                print("DEBUG processX3DHHeader: OPK found")
                #endif
                opkPrivates.removeValue(forKey: opkIdStr)
                KeyStorage.deleteOPKPrivate(email: email, opkId: opkIdStr)
            } else {
                #if DEBUG
                print("DEBUG processX3DHHeader: OPK NOT found - continuing without OPK")
                #endif
            }
        }

        #if DEBUG
        print("DEBUG processX3DHHeader: performing X3DH respond")
        #endif
        let sharedSecret = try X3DH.respond(
            ikPrivateEd: identityPrivate!,
            spkPrivate: spkToUse,
            ikRemoteEd: remoteIkEd,
            ekRemote: ekRemote,
            opkPrivate: opkPriv
        )
        #if DEBUG
        print("DEBUG processX3DHHeader: X3DH respond success, sharedSecret=\(sharedSecret.count) bytes")
        #endif

        let ratchet = DoubleRatchet.initBob(
            sharedSecret: sharedSecret,
            spkPair: (spkToUse, spkToUse.publicKey)
        )

        let sessionKey = "\(senderId):\(senderDeviceId)"
        sessions[sessionKey] = ratchet
        try? KeyStorage.saveSession(email: email, peerUserId: senderId, ratchet: ratchet, localKey: localKey, peerDeviceId: senderDeviceId)
        #if DEBUG
        print("DEBUG processX3DHHeader: session created for \(sessionKey)")
        #endif

        return ratchet
    }

    // MARK: - Send Message

    func sendMessage(convId: String, text: String, members: [ConversationMember], replyTo: String? = nil,
                     extraPayload: [String: Any]? = nil, imageFileId: String? = nil) async -> (success: Bool, message: String, sentMessage: Message?) {
        let isGroup = members.count > 2

        if isGroup {
            return await sendGroupMessage(convId: convId, text: text, members: members, replyTo: replyTo, extraPayload: extraPayload, imageFileId: imageFileId)
        } else {
            return await sendDM(convId: convId, text: text, members: members, replyTo: replyTo, extraPayload: extraPayload, imageFileId: imageFileId)
        }
    }

    // MARK: - Send DM

    private func sendDM(convId: String, text: String, members: [ConversationMember], replyTo: String? = nil,
                        extraPayload: [String: Any]? = nil, imageFileId: String? = nil) async -> (success: Bool, message: String, sentMessage: Message?) {
        #if DEBUG
        print("DEBUG sendDM: convId=\(convId), members=\(members.map { $0.userId })")
        #endif
        guard let identityPrivate = identityPrivate else {
            #if DEBUG
            print("DEBUG sendDM: Identity key not loaded")
            #endif
            return (false, "Identity key not loaded", nil)
        }

        // Build JSON payload matching Python format for cross-client compatibility
        var payload: [String: Any] = [
            "sender": username,
            "text": text,
            "reply_to": (replyTo as Any?) ?? NSNull(),
            "timestamp": ISO8601DateFormatter().string(from: Date()),
        ]
        // Merge extra fields (image, file) at top level
        if let extra = extraPayload {
            for (key, value) in extra { payload[key] = value }
        }
        let rawPlaintext: Data
        if let jsonData = try? JSONSerialization.data(withJSONObject: payload, options: [.sortedKeys]) {
            rawPlaintext = jsonData
        } else {
            rawPlaintext = Data(text.utf8)
        }
        let plaintext = MessagePadding.pad(rawPlaintext)

        var recipients: [[String: Any]] = []

        // Ensure we have the other member(s) — refetch if members only contains self
        var actualMembers = members
        let otherMembers = members.filter { $0.userId != userId }
        if otherMembers.isEmpty {
            #if DEBUG
            print("DEBUG sendDM: WARNING — no other members in conversation, refetching from server")
            #endif
            let allConvs = await listConversations()
            if let conv = allConvs.first(where: { $0.id == convId }) {
                actualMembers = conv.members
                #if DEBUG
                print("DEBUG sendDM: refetched \(actualMembers.count) members for conv \(convId)")
                #endif
            }
        }

        // Encrypt for each member's devices
        #if DEBUG
        print("DEBUG sendDM: encrypting for members, my userId=\(userId ?? "nil")")
        #endif
        for member in actualMembers where member.userId != userId {
            #if DEBUG
            print("DEBUG sendDM: processing member \(member.userId)")
            #endif
            do {
                let bundles = try await getDeviceBundles(userId: member.userId)
                #if DEBUG
                print("DEBUG sendDM: got \(bundles.count) device bundles for \(member.userId)")
                #endif
                for (index, bundle) in bundles.enumerated() {
                    #if DEBUG
                    print("DEBUG sendDM: encrypting for bundle \(index + 1)/\(bundles.count), deviceId=\(bundle.deviceId)")
                    #endif
                    let ratchet = try await getOrCreateSession(
                        peerUserId: member.userId,
                        peerDeviceId: bundle.deviceId,
                        bundle: bundle
                    )

                    // Consume X3DH header if present (first message only)
                    let x3dhHeader = ratchet.x3dhHeader
                    ratchet.x3dhHeader = nil

                    let encrypted = try ratchet.encrypt(plaintext)
                    try? KeyStorage.saveSession(email: email, peerUserId: member.userId, ratchet: ratchet, localKey: localKey, peerDeviceId: bundle.deviceId)

                    var recipientEntry: [String: Any] = [
                        "user_id": member.userId,
                        "device_id": bundle.deviceId,
                        "encrypted_content": ProtocolHandler.encodeBinary(encrypted.ciphertext),
                        "nonce": ProtocolHandler.encodeBinary(encrypted.nonce),
                        "ratchet_header": encrypted.header,
                    ]
                    if let x3dh = x3dhHeader {
                        recipientEntry["x3dh_header"] = x3dh
                        #if DEBUG
                        print("DEBUG sendDM: x3dh_header for \(bundle.deviceId) = \(x3dh)")
                        #endif
                    }
                    #if DEBUG
                    print("DEBUG sendDM: ratchet_header for \(bundle.deviceId) = \(encrypted.header)")
                    #endif
                    recipients.append(recipientEntry)
                }
            } catch {
                #if DEBUG
                print("DEBUG sendDM: ERROR encrypting for member \(member.userId): \(error)")
                #endif
                return (false, "Encryption failed for \(member.username): \(error.localizedDescription)", nil)
            }
        }

        // Self-encrypted copy (readable by all own devices via shared identity key)
        let selfKey = CryptoUtils.deriveSelfEncryptionKey(identityPrivateRaw: identityPrivate.rawData)
        if let (_, nonce, ct, tag) = try? CryptoUtils.aesEncrypt(plaintext, key: selfKey) {
            let selfCiphertext = ct + tag
            recipients.append([
                "user_id": userId!,
                "device_id": Constants.selfDeviceId,
                "encrypted_content": ProtocolHandler.encodeBinary(selfCiphertext),
                "nonce": ProtocolHandler.encodeBinary(nonce),
                "ratchet_header": ["self": true],
            ])
        }

        // Build ratchet header for message table (use first recipient's or dummy)
        let ratchetHeader: [String: Any]
        if let first = recipients.first {
            ratchetHeader = first["ratchet_header"] as? [String: Any] ?? [:]
        } else {
            ratchetHeader = ["dh_pub": String(repeating: "00", count: 32), "n": 0, "pn": 0]
        }

        var params: [String: Any] = [
            "conversation_id": convId,
            "ratchet_header": ratchetHeader,
            "recipients": recipients,
        ]
        if let replyTo = replyTo {
            params["reply_to"] = replyTo
        }
        if let imageFileId = imageFileId {
            params["image_file_id"] = imageFileId
        }

        #if DEBUG
        print("DEBUG sendDM: sending with \(recipients.count) recipients")
        #endif
        let resp = await sendAndReceive(type: "send_message", params: params)
        #if DEBUG
        print("DEBUG sendDM: response status=\(resp.string(for: "status") ?? "nil")")
        #endif
        guard resp.string(for: "status") == "ok" else {
            let msg = resp.dict(for: "data")?.string(for: "message") ?? "Send failed"
            return (false, msg, nil)
        }

        // Build sent Message and cache the plaintext
        var sentMessage: Message?
        if let msgData = resp.dict(for: "data"), let messageId = msgData.string(for: "message_id") {
            // Cache plaintext for future getMessages calls
            MessageCache.cacheDecryptedMessage(
                email: email, convId: convId, messageId: messageId,
                plaintext: plaintext, cacheKey: cacheKey
            )

            var forwardedFrom: ForwardedFrom?
            if let fwd = extraPayload?["forwarded_from"] as? [String: Any],
               let fwdSender = fwd["sender"] as? String {
                forwardedFrom = ForwardedFrom(sender: fwdSender,
                    conversationId: fwd["conversation_id"] as? String ?? "",
                    messageId: fwd["message_id"] as? String ?? "")
            }

            // Parse image/file info from extraPayload so the sent message displays correctly
            var sentImage: ImageInfo?
            if let imgDict = extraPayload?["image"] as? [String: Any],
               let imgFileId = imgDict["file_id"] as? String, !imgFileId.isEmpty {
                sentImage = ImageInfo(
                    fileId: imgFileId,
                    aesKey: imgDict["aes_key"] as? String ?? "",
                    iv: imgDict["iv"] as? String ?? "",
                    thumbnail: imgDict["thumbnail"] as? String,
                    filename: imgDict["filename"] as? String ?? "image.jpg",
                    size: imgDict["size"] as? Int ?? 0
                )
            }
            var sentFile: FileInfo?
            if let fileDict = extraPayload?["file"] as? [String: Any],
               let fFileId = fileDict["file_id"] as? String, !fFileId.isEmpty {
                sentFile = FileInfo(
                    fileId: fFileId,
                    aesKey: fileDict["aes_key"] as? String ?? "",
                    iv: fileDict["iv"] as? String ?? "",
                    filename: fileDict["filename"] as? String ?? "",
                    size: fileDict["size"] as? Int ?? 0,
                    mimeType: fileDict["mime_type"] as? String ?? ""
                )
            }

            let createdAt = msgData.string(for: "created_at").flatMap { DateParsing.parse($0) } ?? Date()
            sentMessage = Message(
                id: messageId, conversationId: convId, senderId: userId!,
                senderUsername: username, createdAt: createdAt,
                text: text.isEmpty ? nil : text, replyTo: replyTo,
                imageFileId: imageFileId, file: sentFile, image: sentImage,
                isDeleted: false, readBy: [],
                reactions: [], forwardedFrom: forwardedFrom, pinnedAt: nil, pinnedBy: nil
            )
        }

        return (true, "Message sent", sentMessage)
    }

    // MARK: - Send Group Message

    private func sendGroupMessage(convId: String, text: String, members: [ConversationMember], replyTo: String? = nil,
                                  extraPayload: [String: Any]? = nil, imageFileId: String? = nil) async -> (success: Bool, message: String, sentMessage: Message?) {
        guard let identityPrivate = identityPrivate, let userId = userId, deviceId != nil else {
            return (false, "Not properly logged in", nil)
        }

        // Ensure we have all members — refetch if only self is present
        var actualMembers = members
        let otherMembers = members.filter { $0.userId != userId }
        if otherMembers.isEmpty {
            #if DEBUG
            print("DEBUG sendGroupMessage: WARNING — no other members, refetching from server")
            #endif
            let allConvs = await listConversations()
            if let conv = allConvs.first(where: { $0.id == convId }) {
                actualMembers = conv.members
            }
        }

        // Get or create sender key for this group
        var senderKeyState = senderKeyStates[convId]
        if senderKeyState == nil {
            senderKeyState = KeyStorage.loadSenderKeyState(email: email, convId: convId, localKey: localKey)
        }

        var needDistribute = false
        if senderKeyState == nil {
            senderKeyState = SenderKeyState()
            needDistribute = true
        }

        senderKeyStates[convId] = senderKeyState

        // Distribute sender key if new
        if needDistribute {
            await distributeSenderKey(convId: convId, members: actualMembers)
        }

        // Encrypt with sender key — JSON payload for cross-client compatibility
        var groupPayload: [String: Any] = [
            "sender": username,
            "text": text,
            "reply_to": (replyTo as Any?) ?? NSNull(),
            "timestamp": ISO8601DateFormatter().string(from: Date()),
        ]
        if let extra = extraPayload {
            for (key, value) in extra { groupPayload[key] = value }
        }
        let rawPlaintext: Data
        if let jsonData = try? JSONSerialization.data(withJSONObject: groupPayload, options: [.sortedKeys]) {
            rawPlaintext = jsonData
        } else {
            rawPlaintext = Data(text.utf8)
        }
        let plaintext = MessagePadding.pad(rawPlaintext)
        do {
            let encrypted = try senderKeyState!.encrypt(plaintext)
            try? KeyStorage.saveSenderKeyState(email: email, convId: convId, state: senderKeyState!, localKey: localKey)

            // Build recipients (same ciphertext for all)
            var recipients: [[String: Any]] = []
            for member in actualMembers where member.userId != userId {
                recipients.append([
                    "user_id": member.userId,
                    "device_id": Constants.selfDeviceId,  // group messages use sentinel
                    "encrypted_content": ProtocolHandler.encodeBinary(encrypted.ciphertext),
                    "nonce": ProtocolHandler.encodeBinary(encrypted.nonce),
                ])
            }

            // Self copy (readable by all own devices via shared identity key)
            let selfKey = CryptoUtils.deriveSelfEncryptionKey(identityPrivateRaw: identityPrivate.rawData)
            if let (_, nonce, ct, tag) = try? CryptoUtils.aesEncrypt(plaintext, key: selfKey) {
                recipients.append([
                    "user_id": userId,
                    "device_id": Constants.selfDeviceId,
                    "encrypted_content": ProtocolHandler.encodeBinary(ct + tag),
                    "nonce": ProtocolHandler.encodeBinary(nonce),
                    "ratchet_header": ["self": true],
                ])
            }

            let dummyHeader: [String: Any] = [
                "dh_pub": String(repeating: "00", count: 32),
                "n": 0,
                "pn": 0,
            ]

            var params: [String: Any] = [
                "conversation_id": convId,
                "ratchet_header": dummyHeader,
                "recipients": recipients,
            ]

            // Include sender key metadata for group routing (base64-encoded chain ID bytes)
            if let chainIdBytes = Data(hexString: encrypted.chainIdHex) {
                params["sender_chain_id"] = ProtocolHandler.encodeBinary(chainIdBytes)
            }
            params["sender_chain_n"] = encrypted.n

            if let replyTo = replyTo {
                params["reply_to"] = replyTo
            }
            if let imageFileId = imageFileId {
                params["image_file_id"] = imageFileId
            }

            let resp = await sendAndReceive(type: "send_message", params: params)
            guard resp.string(for: "status") == "ok" else {
                let msg = resp.dict(for: "data")?.string(for: "message") ?? "Send failed"
                return (false, msg, nil)
            }

            // Build sent Message and cache the plaintext
            var sentMessage: Message?
            if let msgData = resp.dict(for: "data"), let messageId = msgData.string(for: "message_id") {
                MessageCache.cacheDecryptedMessage(
                    email: email, convId: convId, messageId: messageId,
                    plaintext: plaintext, cacheKey: cacheKey
                )

                var forwardedFrom: ForwardedFrom?
                if let fwd = extraPayload?["forwarded_from"] as? [String: Any],
                   let fwdSender = fwd["sender"] as? String {
                    forwardedFrom = ForwardedFrom(sender: fwdSender,
                        conversationId: fwd["conversation_id"] as? String ?? "",
                        messageId: fwd["message_id"] as? String ?? "")
                }

                // Parse image/file info from extraPayload so the sent message displays correctly
                var sentImage: ImageInfo?
                if let imgDict = extraPayload?["image"] as? [String: Any],
                   let imgFileId = imgDict["file_id"] as? String, !imgFileId.isEmpty {
                    sentImage = ImageInfo(
                        fileId: imgFileId,
                        aesKey: imgDict["aes_key"] as? String ?? "",
                        iv: imgDict["iv"] as? String ?? "",
                        thumbnail: imgDict["thumbnail"] as? String,
                        filename: imgDict["filename"] as? String ?? "image.jpg",
                        size: imgDict["size"] as? Int ?? 0
                    )
                }
                var sentFile: FileInfo?
                if let fileDict = extraPayload?["file"] as? [String: Any],
                   let fFileId = fileDict["file_id"] as? String, !fFileId.isEmpty {
                    sentFile = FileInfo(
                        fileId: fFileId,
                        aesKey: fileDict["aes_key"] as? String ?? "",
                        iv: fileDict["iv"] as? String ?? "",
                        filename: fileDict["filename"] as? String ?? "",
                        size: fileDict["size"] as? Int ?? 0,
                        mimeType: fileDict["mime_type"] as? String ?? ""
                    )
                }

                let createdAt = msgData.string(for: "created_at").flatMap { DateParsing.parse($0) } ?? Date()
                sentMessage = Message(
                    id: messageId, conversationId: convId, senderId: userId,
                    senderUsername: username, createdAt: createdAt,
                    text: text.isEmpty ? nil : text, replyTo: replyTo,
                    imageFileId: imageFileId, file: sentFile, image: sentImage,
                    isDeleted: false, readBy: [],
                    reactions: [], forwardedFrom: forwardedFrom, pinnedAt: nil, pinnedBy: nil
                )
            }

            return (true, "Message sent", sentMessage)
        } catch {
            return (false, "Encryption failed: \(error.localizedDescription)", nil)
        }
    }

    // MARK: - Distribute Sender Key

    private func distributeSenderKey(convId: String, members: [ConversationMember]) async {
        guard let senderKeyState = senderKeyStates[convId],
              let userId = userId,
              let deviceId = deviceId else { return }

        let exportedKey = senderKeyState.exportKey()

        for member in members where member.userId != userId {
            do {
                let bundles = try await getDeviceBundles(userId: member.userId)
                for bundle in bundles {
                    let ratchet = try await getOrCreateSession(
                        peerUserId: member.userId,
                        peerDeviceId: bundle.deviceId,
                        bundle: bundle
                    )

                    let x3dhHeader = ratchet.x3dhHeader
                    ratchet.x3dhHeader = nil

                    // Payload includes sender key + metadata
                    let controlPayload: [String: Any] = [
                        "_sender_key": [
                            "conv_id": convId,
                            "key": ProtocolHandler.encodeBinary(exportedKey),
                            "sender_device_id": deviceId,
                        ]
                    ]
                    let controlData = try JSONSerialization.data(withJSONObject: controlPayload)
                    let encrypted = try ratchet.encrypt(controlData)
                    try? KeyStorage.saveSession(email: email, peerUserId: member.userId, ratchet: ratchet, localKey: localKey, peerDeviceId: bundle.deviceId)

                    var recipientEntry: [String: Any] = [
                        "user_id": member.userId,
                        "device_id": bundle.deviceId,
                        "encrypted_content": ProtocolHandler.encodeBinary(encrypted.ciphertext),
                        "nonce": ProtocolHandler.encodeBinary(encrypted.nonce),
                        "ratchet_header": encrypted.header,
                    ]
                    if let x3dh = x3dhHeader {
                        recipientEntry["x3dh_header"] = x3dh
                    }

                    let dummyHeader: [String: Any] = [
                        "dh_pub": String(repeating: "00", count: 32),
                        "n": 0,
                        "pn": 0,
                    ]

                    _ = await sendAndReceive(type: "send_message", params: [
                        "conversation_id": convId,
                        "ratchet_header": dummyHeader,
                        "recipients": [recipientEntry],
                    ])
                }
            } catch {
                #if DEBUG
                print("Failed to distribute sender key to \(member.userId): \(error)")
                #endif
            }
        }
    }

    // MARK: - Decrypt

    /// Decrypt DM using Double Ratchet, or static key for self-copies.
    /// Matches Python: _decrypt_dm() lines 1547-1652
    func decryptDMRecipientData(
        senderData: [String: Any],
        senderId: String,
        senderDeviceId: String
    ) -> Data? {
        // Server uses "encrypted_content", fallback to "ciphertext" for compatibility
        let ctB64 = senderData["encrypted_content"] as? String ?? senderData["ciphertext"] as? String
        guard let ctB64 = ctB64,
              let nonceB64 = senderData["nonce"] as? String,
              let ciphertext = try? ProtocolHandler.decodeBinary(ctB64),
              let nonce = try? ProtocolHandler.decodeBinary(nonceB64) else {
            #if DEBUG
            print("DEBUG decryptDM: missing encrypted_content or nonce")
            #endif
            return nil
        }
        #if DEBUG
        print("DEBUG decryptDM: ct=\(ciphertext.count)B nonce=\(nonce.count)B sender=\(senderId) senderDev=\(senderDeviceId)")
        #endif

        let ratchetHeader = senderData["ratchet_header"] as? [String: Any] ?? [:]

        // Self-encrypted message (ratchet_header.self == true)
        // Matches Python lines 1562-1566
        let isSelf = (ratchetHeader["self"] as? Bool) == true || (ratchetHeader["self"] as? Int) == 1
        if isSelf {
            #if DEBUG
            print("DEBUG decryptDM: self-encrypted copy")
            #endif
            guard let cacheKey = cacheKey, ciphertext.count >= 16 else {
                #if DEBUG
                print("DEBUG decryptDM: no cacheKey or ct too short")
                #endif
                return nil
            }
            let ct = ciphertext.prefix(ciphertext.count - 16)
            let tag = ciphertext.suffix(16)
            let result = try? CryptoUtils.aesDecrypt(key: cacheKey, nonce: nonce, ciphertext: Data(ct), tag: Data(tag))
            #if DEBUG
            print("DEBUG decryptDM: self-decrypt \(result != nil ? "OK" : "FAILED")")
            #endif
            return result
        }

        // Regular DM — Double Ratchet decryption
        // Matches Python lines 1568-1631
        let x3dhHeader = senderData["x3dh_header"] as? [String: Any]
        let sessionKey = "\(senderId):\(senderDeviceId)"

        // Try to load existing session (memory → disk)
        let ratchet = sessions[sessionKey]
            ?? KeyStorage.loadSession(email: email, peerUserId: senderId, localKey: localKey, peerDeviceId: senderDeviceId)
        if ratchet != nil {
            sessions[sessionKey] = ratchet
        }
        #if DEBUG
        print("DEBUG decryptDM: sessionKey=\(sessionKey) existingSession=\(ratchet != nil) hasX3DH=\(x3dhHeader != nil)")
        #endif

        var plaintext: Data?

        if ratchet != nil && x3dhHeader == nil {
            // Normal case: existing session, no X3DH header (Python line 1581-1585)
            do {
                plaintext = try ratchet!.decrypt(headerDict: ratchetHeader, ciphertext: ciphertext, nonce: nonce)
                try? KeyStorage.saveSession(email: email, peerUserId: senderId, ratchet: ratchet!, localKey: localKey, peerDeviceId: senderDeviceId)
                #if DEBUG
                print("DEBUG decryptDM: normal decrypt OK")
                #endif
            } catch {
                #if DEBUG
                print("DEBUG decryptDM: normal decrypt FAILED: \(error)")
                #endif
                return nil
            }
        } else if let x3dh = x3dhHeader, let existingRatchet = ratchet {
            // Existing session + X3DH: sender may have reset. Try existing first, fallback to new X3DH
            // Matches Python lines 1587-1613
            #if DEBUG
            print("DEBUG decryptDM: existing session + X3DH header, trying existing first")
            #endif
            let backup = try? existingRatchet.exportState()
            do {
                plaintext = try existingRatchet.decrypt(headerDict: ratchetHeader, ciphertext: ciphertext, nonce: nonce)
                try? KeyStorage.saveSession(email: email, peerUserId: senderId, ratchet: existingRatchet, localKey: localKey, peerDeviceId: senderDeviceId)
                #if DEBUG
                print("DEBUG decryptDM: existing session decrypt OK (X3DH ignored)")
                #endif
            } catch {
                // Restore from backup and try X3DH
                #if DEBUG
                print("DEBUG decryptDM: existing session failed, trying X3DH")
                #endif
                if let backup = backup, let restored = try? DoubleRatchet.importState(backup) {
                    sessions[sessionKey] = restored
                    try? KeyStorage.saveSession(email: email, peerUserId: senderId, ratchet: restored, localKey: localKey, peerDeviceId: senderDeviceId)
                }
                plaintext = tryX3DHDecrypt(x3dh: x3dh, ratchetHeader: ratchetHeader, ciphertext: ciphertext, nonce: nonce, senderId: senderId, senderDeviceId: senderDeviceId)
            }
        } else if let x3dh = x3dhHeader {
            // No existing session, process X3DH
            // Matches Python lines 1614-1629
            #if DEBUG
            print("DEBUG decryptDM: no existing session, processing X3DH")
            #endif
            plaintext = tryX3DHDecrypt(x3dh: x3dh, ratchetHeader: ratchetHeader, ciphertext: ciphertext, nonce: nonce, senderId: senderId, senderDeviceId: senderDeviceId)
        } else {
            // No session and no X3DH
            #if DEBUG
            print("DEBUG decryptDM: no session and no X3DH for \(senderId)")
            #endif
            return nil
        }

        guard let rawDecrypted = plaintext else { return nil }
        let unpadded = MessagePadding.unpad(rawDecrypted)

        sessions[sessionKey] = ratchet ?? sessions[sessionKey]
        #if DEBUG
        print("DEBUG decryptDM: decrypt OK, \(unpadded.count) bytes")
        #endif

        // Check for sender key distribution (control message)
        if let jsonObj = try? JSONSerialization.jsonObject(with: unpadded) as? [String: Any],
           let senderKeyInfo = jsonObj["_sender_key"] as? [String: Any] {
            handleSenderKeyDistribution(senderKeyInfo, senderId: senderId)
            return nil  // Control message
        }

        return unpadded
    }

    /// Helper: try X3DH header processing + decrypt, with prev SPK fallback.
    /// Matches Python's try/except pattern in _decrypt_dm lines 1615-1629.
    private func tryX3DHDecrypt(
        x3dh: [String: Any], ratchetHeader: [String: Any],
        ciphertext: Data, nonce: Data,
        senderId: String, senderDeviceId: String
    ) -> Data? {
        let sessionKey = "\(senderId):\(senderDeviceId)"
        do {
            let ratchet = try processX3DHHeader(
                senderId: senderId, x3dhHeader: x3dh, senderDeviceId: senderDeviceId)
            let plaintext = try ratchet.decrypt(headerDict: ratchetHeader, ciphertext: ciphertext, nonce: nonce)
            sessions[sessionKey] = ratchet
            try? KeyStorage.saveSession(email: email, peerUserId: senderId, ratchet: ratchet, localKey: localKey, peerDeviceId: senderDeviceId)
            #if DEBUG
            print("DEBUG tryX3DHDecrypt: OK with current SPK")
            #endif
            return plaintext
        } catch {
            #if DEBUG
            print("DEBUG tryX3DHDecrypt: failed with current SPK: \(error)")
            #endif
            // Try with previous SPK (grace period, M4)
            if let prevSpk = prevSpkPrivate {
                do {
                    let ratchet = try processX3DHHeader(
                        senderId: senderId, x3dhHeader: x3dh, senderDeviceId: senderDeviceId,
                        spkOverride: prevSpk)
                    let plaintext = try ratchet.decrypt(headerDict: ratchetHeader, ciphertext: ciphertext, nonce: nonce)
                    sessions[sessionKey] = ratchet
                    try? KeyStorage.saveSession(email: email, peerUserId: senderId, ratchet: ratchet, localKey: localKey, peerDeviceId: senderDeviceId)
                    #if DEBUG
                    print("DEBUG tryX3DHDecrypt: OK with prev SPK")
                    #endif
                    return plaintext
                } catch {
                    #if DEBUG
                    print("DEBUG tryX3DHDecrypt: failed with prev SPK too: \(error)")
                    #endif
                }
            }
            return nil
        }
    }

    private func handleSenderKeyDistribution(_ info: [String: Any], senderId: String) {
        guard let convId = info["conv_id"] as? String,
              let keyB64 = info["key"] as? String,
              let keyData = try? ProtocolHandler.decodeBinary(keyB64) else { return }

        let senderDeviceId = info["sender_device_id"] as? String ?? Constants.selfDeviceId

        do {
            let senderKey = try SenderKeyState.fromKey(keyData)
            let stateKey = "\(convId):\(senderId):\(senderDeviceId)"
            recvSenderKeys[stateKey] = senderKey
            try? KeyStorage.saveRecvSenderKey(
                email: email,
                convId: convId,
                senderId: senderId,
                senderDeviceId: senderDeviceId,
                state: senderKey,
                localKey: localKey
            )
        } catch {
            #if DEBUG
            print("Failed to import sender key: \(error)")
            #endif
        }
    }

    // MARK: - Group Decrypt

    /// Decrypt a group message using the sender's distributed Sender Key.
    /// Matches Python: _decrypt_group() in chat_core.py
    func decryptGroupMessage(
        msgDict: [String: Any],
        senderId: String,
        senderDeviceId: String,
        convId: String
    ) -> Data? {
        guard let chainIdB64 = msgDict["sender_chain_id"] as? String,
              let chainN = msgDict["sender_chain_n"] as? Int,
              let chainIdData = try? ProtocolHandler.decodeBinary(chainIdB64) else {
            #if DEBUG
            print("DEBUG decryptGroupMessage: missing sender_chain_id or sender_chain_n")
            #endif
            return nil
        }
        let chainIdHex = chainIdData.hexString

        let ctB64 = msgDict["encrypted_content"] as? String ?? msgDict["ciphertext"] as? String
        guard let ctB64 = ctB64,
              let nonceB64 = msgDict["nonce"] as? String,
              let ciphertext = try? ProtocolHandler.decodeBinary(ctB64),
              let nonce = try? ProtocolHandler.decodeBinary(nonceB64) else {
            #if DEBUG
            print("DEBUG decryptGroupMessage: missing encrypted_content or nonce")
            #endif
            return nil
        }

        // Look up received sender key — try with sender_device_id first
        var sk: SenderKeyState?
        let keyWithDevice = "\(convId):\(senderId):\(senderDeviceId)"
        sk = recvSenderKeys[keyWithDevice]
        if sk == nil {
            sk = KeyStorage.loadRecvSenderKey(
                email: email, convId: convId, senderId: senderId,
                senderDeviceId: senderDeviceId, localKey: localKey
            )
            if let loaded = sk {
                recvSenderKeys[keyWithDevice] = loaded
            }
        }

        // Fallback: try with self device ID (legacy/default)
        if sk == nil && senderDeviceId != Constants.selfDeviceId {
            let keyLegacy = "\(convId):\(senderId):\(Constants.selfDeviceId)"
            sk = recvSenderKeys[keyLegacy]
            if sk == nil {
                sk = KeyStorage.loadRecvSenderKey(
                    email: email, convId: convId, senderId: senderId,
                    senderDeviceId: Constants.selfDeviceId, localKey: localKey
                )
                if let loaded = sk {
                    recvSenderKeys[keyLegacy] = loaded
                }
            }
        }

        guard let senderKey = sk else {
            #if DEBUG
            print("DEBUG decryptGroupMessage: no sender key for \(senderId) in conv \(convId)")
            #endif
            return nil
        }

        do {
            let rawPlaintext = try senderKey.decrypt(chainIdHex: chainIdHex, n: chainN, ciphertext: ciphertext, nonce: nonce)
            let plaintext = MessagePadding.unpad(rawPlaintext)
            // Save updated state (chain has advanced)
            try? KeyStorage.saveRecvSenderKey(
                email: email, convId: convId, senderId: senderId,
                senderDeviceId: senderDeviceId, state: senderKey, localKey: localKey
            )
            #if DEBUG
            print("DEBUG decryptGroupMessage: success, \(plaintext.count) bytes")
            #endif
            return plaintext
        } catch {
            #if DEBUG
            print("DEBUG decryptGroupMessage: decrypt failed: \(error)")
            #endif
            return nil
        }
    }

    /// Decrypt a new_message notification. Matches Python: decrypt_notification().
    /// Supports multi-device format (device_entries array) and legacy flat format.
    func decryptNotification(_ data: [String: Any]) -> Message? {
        guard let senderId = data.string(for: "sender_id"),
              let conversationId = data.string(for: "conversation_id"),
              let messageId = data.string(for: "message_id") else {
            #if DEBUG
            print("DEBUG decryptNotification: missing sender_id/conversation_id/message_id")
            #endif
            return nil
        }

        let senderDeviceId = data.string(for: "sender_device_id") ?? Constants.selfDeviceId
        let myUserId = userId ?? ""
        #if DEBUG
        print("DEBUG decryptNotification: msgId=\(messageId) senderId=\(senderId) senderDeviceId=\(senderDeviceId) convId=\(conversationId)")
        #endif

        // --- Step 1: Extract per-device encrypted content ---
        // Matches Python decrypt_notification lines 1862-1897
        var encryptedContent: String = ""
        var nonce: String = ""
        var ratchetHeader: [String: Any] = [:]
        var x3dhHeader: [String: Any]?

        if let deviceEntries = data["device_entries"] as? [[String: Any]] {
            // Multi-device format: pick entry matching our device_id or SELF_DEVICE_ID
            let entryDeviceIds = deviceEntries.compactMap { $0["device_id"] as? String }
            #if DEBUG
            print("DEBUG decryptNotification: myDeviceId=\(deviceId ?? "nil") device_entries=\(entryDeviceIds)")
            #endif

            var chosen: [String: Any]?
            var selfEntry: [String: Any]?
            for entry in deviceEntries {
                let eid = entry["device_id"] as? String ?? ""
                if eid == deviceId {
                    chosen = entry
                    break
                }
                if eid == Constants.selfDeviceId {
                    selfEntry = entry
                }
            }

            // If sender is us, prefer self-encrypted entry (matches Python lines 1878-1882)
            if senderId == myUserId {
                chosen = selfEntry ?? chosen
            } else if chosen == nil {
                chosen = selfEntry
            }

            guard let chosen = chosen else {
                #if DEBUG
                print("DEBUG decryptNotification: NO matching device_entry for device \(deviceId ?? "nil")")
                #endif
                return nil
            }

            encryptedContent = chosen["encrypted_content"] as? String ?? ""
            nonce = chosen["nonce"] as? String ?? ""
            ratchetHeader = (chosen["ratchet_header"] as? [String: Any]) ?? (data["ratchet_header"] as? [String: Any]) ?? [:]
            x3dhHeader = (chosen["x3dh_header"] as? [String: Any]) ?? (data["x3dh_header"] as? [String: Any])
        } else {
            // Legacy flat format (no device_entries)
            encryptedContent = data["encrypted_content"] as? String ?? data["ciphertext"] as? String ?? ""
            nonce = data["nonce"] as? String ?? ""
            ratchetHeader = data["ratchet_header"] as? [String: Any] ?? [:]
            x3dhHeader = data["x3dh_header"] as? [String: Any]
        }

        guard !encryptedContent.isEmpty, !nonce.isEmpty else {
            #if DEBUG
            print("DEBUG decryptNotification: empty encrypted_content or nonce")
            #endif
            return nil
        }

        // --- Step 2: Build msg_data and decrypt ---
        // Matches Python decrypt_notification lines 1899-1910
        let msgData: [String: Any] = [
            "sender_id": senderId,
            "sender_device_id": senderDeviceId,
            "conversation_id": conversationId,
            "ratchet_header": ratchetHeader,
            "encrypted_content": encryptedContent,
            "nonce": nonce,
            "x3dh_header": x3dhHeader as Any,
            "sender_chain_id": data["sender_chain_id"] as Any,
            "sender_chain_n": data["sender_chain_n"] as Any,
        ]

        // Dispatch: matches Python _decrypt_message() lines 1533-1545
        // Check self-encrypted FIRST (before group check) — after re-encryption,
        // group messages have {"self": true} ratchet_header but still sender_chain_id.
        var plaintext: Data?
        let isSelfEncrypted = (ratchetHeader["self"] as? Bool) == true
            || (ratchetHeader["self"] as? Int) == 1

        if isSelfEncrypted {
            // Self-encrypted copy → static self-key (same path as _decrypt_dm self check)
            #if DEBUG
            print("DEBUG decryptNotification: self-encrypted message")
            #endif
            if let cacheKey = cacheKey {
                if let ct = try? ProtocolHandler.decodeBinary(encryptedContent),
                   let n = try? ProtocolHandler.decodeBinary(nonce),
                   ct.count >= 16 {
                    let ciphertext = ct.prefix(ct.count - 16)
                    let tag = ct.suffix(16)
                    plaintext = try? CryptoUtils.aesDecrypt(key: cacheKey, nonce: n, ciphertext: Data(ciphertext), tag: Data(tag))
                    #if DEBUG
                    print("DEBUG decryptNotification: self-decrypt result=\(plaintext != nil ? "OK" : "FAILED")")
                    #endif
                }
            }
        } else if let chainId = msgData["sender_chain_id"] as? String, !chainId.isEmpty, senderId != myUserId {
            // Group message from someone else → sender key decryption
            #if DEBUG
            print("DEBUG decryptNotification: group message")
            #endif
            plaintext = decryptGroupMessage(
                msgDict: msgData, senderId: senderId,
                senderDeviceId: senderDeviceId, convId: conversationId
            )
        } else {
            // DM → ratchet decryption
            #if DEBUG
            print("DEBUG decryptNotification: DM message, calling decryptDMRecipientData")
            #endif
            plaintext = decryptDMRecipientData(
                senderData: msgData,
                senderId: senderId,
                senderDeviceId: senderDeviceId
            )
        }

        guard let plaintext = plaintext else {
            #if DEBUG
            print("DEBUG decryptNotification: decrypt returned nil for msgId=\(messageId)")
            #endif
            return nil
        }

        // --- Step 3: Cache and parse ---
        MessageCache.cacheDecryptedMessage(
            email: email, convId: conversationId, messageId: messageId,
            plaintext: plaintext, cacheKey: cacheKey
        )

        // Queue for self-encryption if from another user (so other devices can read it)
        if senderId != myUserId && !messageId.isEmpty {
            pendingSelfEncrypt.append((messageId: messageId, plaintext: plaintext))
        }

        var messageText: String? = String(data: plaintext, encoding: .utf8)
        var replyTo: String?
        var file: FileInfo?
        var image: ImageInfo?
        var forwardedFrom: ForwardedFrom?
        var senderName = data.string(for: "sender_username") ?? userCache[senderId]?.username ?? "Unknown"

        if let jsonObj = try? JSONSerialization.jsonObject(with: plaintext) as? [String: Any] {
            messageText = jsonObj["text"] as? String
            replyTo = jsonObj["reply_to"] as? String
            if senderName == "Unknown", let s = jsonObj["sender"] as? String {
                senderName = s
            }
            if let fileDict = jsonObj["file"] as? [String: Any] {
                file = FileInfo(
                    fileId: fileDict["file_id"] as? String ?? "",
                    aesKey: fileDict["aes_key"] as? String ?? "",
                    iv: fileDict["iv"] as? String ?? "",
                    filename: fileDict["filename"] as? String ?? "",
                    size: fileDict["size"] as? Int ?? 0,
                    mimeType: fileDict["mime_type"] as? String ?? ""
                )
            }
            image = parseImageInfo(from: jsonObj)
            if let fwd = jsonObj["forwarded_from"] as? [String: Any],
               let fwdSender = fwd["sender"] as? String {
                forwardedFrom = ForwardedFrom(sender: fwdSender,
                    conversationId: fwd["conversation_id"] as? String ?? "",
                    messageId: fwd["message_id"] as? String ?? "")
            }

            // Check for sender key distribution (control message)
            if jsonObj["_sender_key"] != nil {
                #if DEBUG
                print("DEBUG decryptNotification: control message (sender key distribution)")
                #endif
                return nil
            }
        }

        if messageText == nil && file == nil && image == nil {
            #if DEBUG
            print("DEBUG decryptNotification: control message (no text/file/image)")
            #endif
            return nil
        }

        let createdAt = data.string(for: "created_at").flatMap { DateParsing.parse($0) } ?? Date()

        return Message(
            id: messageId,
            conversationId: conversationId,
            senderId: senderId,
            senderUsername: senderName,
            createdAt: createdAt,
            text: messageText,
            replyTo: replyTo,
            imageFileId: data.string(for: "image_file_id"),
            file: file,
            image: image,
            isDeleted: false,
            readBy: [],
            reactions: [],
            forwardedFrom: forwardedFrom,
            pinnedAt: nil,
            pinnedBy: nil
        )
    }

    // MARK: - Self-Encrypt Flush

    /// Encrypt received messages with self-encryption key and upload to server.
    /// Allows other devices of the same user to read received messages.
    /// Matches Python: _flush_self_encrypt()
    func flushSelfEncrypt() async {
        guard !pendingSelfEncrypt.isEmpty, let edPriv = identityPrivate else { return }

        let selfKey = CryptoUtils.deriveSelfEncryptionKey(identityPrivateRaw: edPriv.rawData)
        var updates: [[String: Any]] = []

        for item in pendingSelfEncrypt {
            guard let (_, nonce, ct, tag) = try? CryptoUtils.aesEncrypt(item.plaintext, key: selfKey) else {
                continue
            }
            let selfCiphertext = ct + tag
            updates.append([
                "message_id": item.messageId,
                "encrypted_content": ProtocolHandler.encodeBinary(selfCiphertext),
                "nonce": ProtocolHandler.encodeBinary(nonce),
            ])
        }
        pendingSelfEncrypt.removeAll()

        // Send in batches of 500 (matches Python batch size)
        let batchSize = 500
        for batchStart in stride(from: 0, to: updates.count, by: batchSize) {
            let batchEnd = min(batchStart + batchSize, updates.count)
            let batch = Array(updates[batchStart..<batchEnd])
            _ = await sendAndReceive(type: "reencrypt_messages", params: ["updates": batch])
        }
    }

    // MARK: - Conversations

    func listConversations() async -> [Conversation] {
        let resp = await sendAndReceive(type: "list_conversations")
        guard resp.string(for: "status") == "ok",
              let data = resp.dict(for: "data"),
              let convList = data["conversations"] as? [[String: Any]] else {
            return []
        }

        return convList.compactMap { dict -> Conversation? in
            guard let id = dict.string(for: "conversation_id") else { return nil }

            let membersRaw = dict["members"] as? [[String: Any]] ?? []
            #if DEBUG
            print("DEBUG listConversations: conv \(id) raw members count=\(membersRaw.count), raw=\(membersRaw)")
            #endif
            let members = membersRaw.compactMap { m -> ConversationMember? in
                guard let uid = m.string(for: "user_id"),
                      let uname = m.string(for: "username"),
                      let uemail = m.string(for: "email") else {
                    #if DEBUG
                    print("DEBUG listConversations: FAILED to parse member: \(m)")
                    #endif
                    return nil
                }
                return ConversationMember(userId: uid, username: uname, email: uemail)
            }
            #if DEBUG
            print("DEBUG listConversations: conv \(id) parsed members count=\(members.count)")
            #endif

            let unreadCount = dict.int(for: "unread_count") ?? 0

            return Conversation(
                id: id,
                name: dict.string(for: "name"),
                members: members,
                createdBy: dict.string(for: "created_by"),
                avatarFile: dict.string(for: "avatar_file"),
                unreadCount: unreadCount,
                isFavorite: false,
                lastMessageTime: nil
            )
        }
    }

    func createConversation(emails: [String], name: String? = nil) async -> (convId: String?, message: String) {
        var params: [String: Any] = ["members": emails]
        if let name = name {
            params["name"] = name
        }

        let resp = await sendAndReceive(type: "create_conversation", params: params)
        guard resp.string(for: "status") == "ok",
              let data = resp.dict(for: "data"),
              let convId = data.string(for: "conversation_id") else {
            let msg = resp.dict(for: "data")?.string(for: "message") ?? "Failed to create conversation"
            return (nil, msg)
        }

        return (convId, "Conversation created")
    }

    func findConversation(email: String) async -> String? {
        let resp = await sendAndReceive(type: "find_conversation", params: ["email": email])
        guard resp.string(for: "status") == "ok",
              let data = resp.dict(for: "data") else { return nil }
        return data.string(for: "conversation_id")
    }

    // MARK: - Payload Parsing Helpers

    private func parseImageInfo(from jsonObj: [String: Any]) -> ImageInfo? {
        guard let imgDict = jsonObj["image"] as? [String: Any],
              let fileId = imgDict["file_id"] as? String, !fileId.isEmpty else { return nil }
        return ImageInfo(
            fileId: fileId,
            aesKey: imgDict["aes_key"] as? String ?? "",
            iv: imgDict["iv"] as? String ?? "",
            thumbnail: imgDict["thumbnail"] as? String,
            filename: imgDict["filename"] as? String ?? "image.jpg",
            size: imgDict["size"] as? Int ?? 0
        )
    }

    // MARK: - Messages

    func getMessages(convId: String, limit: Int = 50, offset: Int = 0, afterTs: String? = nil) async -> [Message] {
        var params: [String: Any] = [
            "conversation_id": convId,
            "limit": limit,
            "offset": offset,
        ]
        if let afterTs = afterTs {
            params["after_ts"] = afterTs
        }
        let resp = await sendAndReceive(type: "get_messages", params: params)
        guard resp.string(for: "status") == "ok",
              let data = resp.dict(for: "data"),
              let messagesRaw = data["messages"] as? [[String: Any]] else {
            return []
        }

        // Server returns DESC (newest first), reverse to ASC (oldest first) — matches Python client
        let sortedMessages = Array(messagesRaw.reversed())

        // Deduplicate: server can return both device-specific and SELF_DEVICE_ID rows
        // for the same message. Prefer self-encrypted rows (new devices after pairing
        // can't decrypt old device-specific ratchet data, but CAN decrypt self-encrypted).
        var dedupedMessages: [[String: Any]] = []
        var msgIndex: [String: Int] = [:]
        for msgDict in sortedMessages {
            guard let msgId = msgDict.string(for: "message_id") else { continue }
            let rh = msgDict["ratchet_header"] as? [String: Any]
            let isSelf = (rh?["self"] as? Bool) == true || (rh?["self"] as? Int) == 1
            if let idx = msgIndex[msgId] {
                if isSelf { dedupedMessages[idx] = msgDict }
            } else {
                msgIndex[msgId] = dedupedMessages.count
                dedupedMessages.append(msgDict)
            }
        }

        var messages: [Message] = []
        #if DEBUG
        print("DEBUG getMessages: processing \(dedupedMessages.count) messages (from \(sortedMessages.count) rows)")
        #endif
        for msgDict in dedupedMessages {
            guard let msgId = msgDict.string(for: "message_id"),
                  let senderId = msgDict.string(for: "sender_id") else {
                #if DEBUG
                print("DEBUG getMessages: skipping message without message_id or sender_id")
                #endif
                continue
            }

            let senderDeviceId = msgDict.string(for: "sender_device_id") ?? Constants.selfDeviceId
            let isDeleted = msgDict["deleted_at"] != nil && !(msgDict["deleted_at"] is NSNull)
            #if DEBUG
            print("DEBUG getMessages: msgId=\(msgId), senderId=\(senderId), senderDeviceId=\(senderDeviceId), myUserId=\(userId ?? "nil")")
            print("DEBUG getMessages: msgDict keys=\(Array(msgDict.keys))")
            #endif

            // Parse read_by from server response (matches Python: m.get("read_by", []))
            let readBy = Set(msgDict["read_by"] as? [String] ?? [])

            // Parse reactions from server response
            var reactions: [MessageReaction] = []
            if let reactionsArr = msgDict["reactions"] as? [[String: Any]] {
                reactions = reactionsArr.compactMap { r in
                    guard let rUserId = r["user_id"] as? String,
                          let reaction = r["reaction"] as? String else { return nil }
                    let rCreatedAt = (r["created_at"] as? String).flatMap { DateParsing.parse($0) } ?? Date()
                    return MessageReaction(userId: rUserId, reaction: reaction, createdAt: rCreatedAt)
                }
            }

            // Parse pin info from server response
            let pinnedAt = (msgDict["pinned_at"] as? String).flatMap { DateParsing.parse($0) }
            let pinnedBy = msgDict["pinned_by"] as? String

            if isDeleted {
                let createdAt = msgDict.string(for: "created_at").flatMap { DateParsing.parse($0) } ?? Date()
                messages.append(Message(
                    id: msgId, conversationId: convId, senderId: senderId,
                    senderUsername: "",
                    createdAt: createdAt, text: nil, isDeleted: true, readBy: readBy,
                    reactions: [], forwardedFrom: nil, pinnedAt: pinnedAt, pinnedBy: pinnedBy
                ))
                continue
            }

            // Check per-message cache first (ratchet keys are one-time — cannot re-decrypt)
            if let cachedPlaintext = MessageCache.getCachedMessage(
                email: email, convId: convId, messageId: msgId, cacheKey: cacheKey
            ) {
                var messageText: String? = String(data: cachedPlaintext, encoding: .utf8)
                var replyTo: String?
                var file: FileInfo?
                var image: ImageInfo?
                var forwardedFrom: ForwardedFrom?
                var senderName = userCache[senderId]?.username ?? ""
                if let jsonObj = try? JSONSerialization.jsonObject(with: cachedPlaintext) as? [String: Any] {
                    messageText = jsonObj["text"] as? String
                    replyTo = jsonObj["reply_to"] as? String
                    if senderName.isEmpty, let s = jsonObj["sender"] as? String {
                        senderName = s
                    }
                    if let fileDict = jsonObj["file"] as? [String: Any] {
                        file = FileInfo(
                            fileId: fileDict["file_id"] as? String ?? "",
                            aesKey: fileDict["aes_key"] as? String ?? "",
                            iv: fileDict["iv"] as? String ?? "",
                            filename: fileDict["filename"] as? String ?? "",
                            size: fileDict["size"] as? Int ?? 0,
                            mimeType: fileDict["mime_type"] as? String ?? ""
                        )
                    }
                    image = parseImageInfo(from: jsonObj)
                    if let fwd = jsonObj["forwarded_from"] as? [String: Any],
                       let fwdSender = fwd["sender"] as? String {
                        forwardedFrom = ForwardedFrom(sender: fwdSender,
                            conversationId: fwd["conversation_id"] as? String ?? "",
                            messageId: fwd["message_id"] as? String ?? "")
                    }
                }
                if messageText == nil && file == nil && image == nil { continue }  // Control message
                let createdAt = msgDict.string(for: "created_at").flatMap { DateParsing.parse($0) } ?? Date()
                messages.append(Message(
                    id: msgId, conversationId: convId, senderId: senderId,
                    senderUsername: senderName,
                    createdAt: createdAt, text: messageText, replyTo: replyTo,
                    imageFileId: msgDict.string(for: "image_file_id"), file: file,
                    image: image,
                    isDeleted: false, readBy: readBy,
                    reactions: reactions, forwardedFrom: forwardedFrom,
                    pinnedAt: pinnedAt, pinnedBy: pinnedBy
                ))
                #if DEBUG
                print("DEBUG getMessages: loaded from cache msgId=\(msgId)")
                #endif
                continue
            }

            // Try to decrypt (only for messages not in cache)
            // Dispatch: self-encrypted (ratchet_header.self) → static key
            //           group messages (sender_chain_id present, from others) → sender key decrypt
            //           DM or self-copy → ratchet / self-key decrypt
            var plaintext: Data?
            let ratchetHeader = msgDict["ratchet_header"] as? [String: Any]
            let isSelfEncrypted = (ratchetHeader?["self"] as? Bool) == true
                || (ratchetHeader?["self"] as? Int) == 1

            #if DEBUG
            print("DEBUG getMessages: decrypt msgId=\(msgId) senderId=\(senderId) isSelfEncrypted=\(isSelfEncrypted) ratchetHeader=\(ratchetHeader as Any) hasCacheKey=\(cacheKey != nil)")
            #endif

            if isSelfEncrypted, let cacheKey = cacheKey {
                // Re-encrypted message (from reencryptHistory) — decrypt with self-key
                if let ctB64 = (msgDict["encrypted_content"] as? String ?? msgDict["ciphertext"] as? String),
                   let nonceB64 = msgDict["nonce"] as? String,
                   let ct = try? ProtocolHandler.decodeBinary(ctB64),
                   let nonce = try? ProtocolHandler.decodeBinary(nonceB64),
                   ct.count >= 16 {
                    let ciphertext = ct.prefix(ct.count - 16)
                    let tag = ct.suffix(16)
                    plaintext = try? CryptoUtils.aesDecrypt(key: cacheKey, nonce: nonce, ciphertext: Data(ciphertext), tag: Data(tag))
                    #if DEBUG
                    print("DEBUG getMessages: self-decrypt msgId=\(msgId) ctLen=\(ct.count) nonceLen=\(nonce.count) result=\(plaintext != nil ? "OK(\(plaintext!.count)B)" : "FAILED")")
                    #endif
                } else {
                    #if DEBUG
                    print("DEBUG getMessages: self-decrypt msgId=\(msgId) SKIPPED (missing ct/nonce)")
                    #endif
                }
            } else if let chainId = msgDict["sender_chain_id"] as? String, !chainId.isEmpty, senderId != userId {
                plaintext = decryptGroupMessage(
                    msgDict: msgDict, senderId: senderId,
                    senderDeviceId: senderDeviceId, convId: convId
                )
            } else {
                plaintext = decryptDMRecipientData(
                    senderData: msgDict, senderId: senderId,
                    senderDeviceId: senderDeviceId
                )
            }

            if let plaintext = plaintext {
                // Cache the decrypted plaintext for future use
                MessageCache.cacheDecryptedMessage(
                    email: email, convId: convId, messageId: msgId,
                    plaintext: plaintext, cacheKey: cacheKey
                )

                // Queue for self-encryption if from another user
                if senderId != userId && !msgId.isEmpty {
                    pendingSelfEncrypt.append((messageId: msgId, plaintext: plaintext))
                }

                let text = String(data: plaintext, encoding: .utf8)
                var messageText = text
                var replyTo: String?
                var file: FileInfo?
                var image: ImageInfo?
                var forwardedFrom: ForwardedFrom?

                var senderName = userCache[senderId]?.username ?? ""
                if let jsonObj = try? JSONSerialization.jsonObject(with: plaintext) as? [String: Any] {
                    messageText = jsonObj["text"] as? String
                    replyTo = jsonObj["reply_to"] as? String
                    if senderName.isEmpty, let s = jsonObj["sender"] as? String {
                        senderName = s
                    }
                    if let fileDict = jsonObj["file"] as? [String: Any] {
                        file = FileInfo(
                            fileId: fileDict["file_id"] as? String ?? "",
                            aesKey: fileDict["aes_key"] as? String ?? "",
                            iv: fileDict["iv"] as? String ?? "",
                            filename: fileDict["filename"] as? String ?? "",
                            size: fileDict["size"] as? Int ?? 0,
                            mimeType: fileDict["mime_type"] as? String ?? ""
                        )
                    }
                    image = parseImageInfo(from: jsonObj)
                    if let fwd = jsonObj["forwarded_from"] as? [String: Any],
                       let fwdSender = fwd["sender"] as? String {
                        forwardedFrom = ForwardedFrom(sender: fwdSender,
                            conversationId: fwd["conversation_id"] as? String ?? "",
                            messageId: fwd["message_id"] as? String ?? "")
                    }
                }

                if messageText == nil && file == nil && image == nil { continue }  // Control message

                let createdAt = msgDict.string(for: "created_at").flatMap { DateParsing.parse($0) } ?? Date()
                messages.append(Message(
                    id: msgId, conversationId: convId, senderId: senderId,
                    senderUsername: senderName,
                    createdAt: createdAt, text: messageText, replyTo: replyTo,
                    imageFileId: msgDict.string(for: "image_file_id"), file: file,
                    image: image,
                    isDeleted: false, readBy: readBy,
                    reactions: reactions, forwardedFrom: forwardedFrom,
                    pinnedAt: pinnedAt, pinnedBy: pinnedBy
                ))
                #if DEBUG
                print("DEBUG getMessages: decrypted and cached msgId=\(msgId)")
                #endif
            } else {
                #if DEBUG
                print("DEBUG getMessages: decryption failed for msgId=\(msgId)")
                #endif
            }
        }

        // Confirm delivery for messages from others (fire-and-forget)
        let deliverIds = messages.filter { $0.senderId != userId && !$0.isDeleted }.map(\.id)
        if !deliverIds.isEmpty {
            Task { await confirmDelivery(convId: convId, messageIds: deliverIds) }
        }

        // Flush any queued self-encryptions in background
        if !pendingSelfEncrypt.isEmpty {
            Task { await flushSelfEncrypt() }
        }

        #if DEBUG
        print("DEBUG getMessages: returning \(messages.count) messages")
        #endif
        return messages
    }

    func markRead(convId: String, messageIds: [String]) async {
        _ = await sendAndReceive(type: "mark_read", params: [
            "conversation_id": convId,
            "message_ids": messageIds,
        ])
    }

    func markConversationRead(convId: String) async {
        _ = await sendAndReceive(type: "mark_conversation_read", params: [
            "conversation_id": convId,
        ])
    }

    func deleteMessage(messageId: String, convId: String) async -> Bool {
        let resp = await sendAndReceive(type: "delete_message", params: [
            "message_id": messageId,
            "conversation_id": convId,
        ])
        return resp.string(for: "status") == "ok"
    }

    func getDeletedSince(convId: String, sinceTs: String) async -> [String] {
        let resp = await sendAndReceive(type: "get_deleted_since", params: [
            "conversation_id": convId,
            "since_ts": sinceTs,
        ])
        guard resp.string(for: "status") == "ok",
              let data = resp.dict(for: "data"),
              let ids = data["deleted_ids"] as? [String] else { return [] }
        return ids
    }

    func reactMessage(messageId: String, conversationId: String,
                      reaction: String, action: String) async -> Bool {
        let resp = await sendAndReceive(type: "react_message", params: [
            "message_id": messageId,
            "conversation_id": conversationId,
            "reaction": reaction,
            "action": action,
        ])
        return resp.string(for: "status") == "ok"
    }

    func pinMessage(messageId: String, conversationId: String,
                    action: String) async -> Bool {
        let resp = await sendAndReceive(type: "pin_message", params: [
            "message_id": messageId,
            "conversation_id": conversationId,
            "action": action,
        ])
        return resp.string(for: "status") == "ok"
    }

    func getPinnedMessages(conversationId: String) async -> [[String: Any]] {
        let resp = await sendAndReceive(type: "get_pinned_messages", params: [
            "conversation_id": conversationId,
        ])
        guard resp.string(for: "status") == "ok",
              let data = resp.dict(for: "data"),
              let messages = data["messages"] as? [[String: Any]] else { return [] }
        return messages
    }

    // MARK: - Group Operations

    func addMember(convId: String, email: String) async -> (success: Bool, message: String) {
        let resp = await sendAndReceive(type: "add_member", params: [
            "conversation_id": convId,
            "email": email,
        ])
        let msg = resp.dict(for: "data")?.string(for: "message") ?? ""
        return (resp.string(for: "status") == "ok", msg)
    }

    func removeMember(convId: String, userId: String) async -> (success: Bool, message: String) {
        let resp = await sendAndReceive(type: "remove_member", params: [
            "conversation_id": convId,
            "user_id": userId,
        ])
        let msg = resp.dict(for: "data")?.string(for: "message") ?? ""
        return (resp.string(for: "status") == "ok", msg)
    }

    func leaveGroup(convId: String) async -> (success: Bool, message: String) {
        let resp = await sendAndReceive(type: "leave_group", params: [
            "conversation_id": convId,
        ])
        if resp.string(for: "status") == "ok" {
            // Clean up local sender keys
            senderKeyStates.removeValue(forKey: convId)
            KeyStorage.deleteSenderKeyState(email: email, convId: convId)
            KeyStorage.deleteRecvSenderKeys(email: email, convId: convId)
            return (true, "Left group")
        }
        let msg = resp.dict(for: "data")?.string(for: "message") ?? "Failed"
        return (false, msg)
    }

    func renameConversation(convId: String, name: String) async -> (success: Bool, message: String) {
        let resp = await sendAndReceive(type: "rename_conversation", params: [
            "conversation_id": convId,
            "name": name,
        ])
        let msg = resp.dict(for: "data")?.string(for: "message") ?? ""
        return (resp.string(for: "status") == "ok", msg)
    }

    func deleteConversation(convId: String) async -> (success: Bool, message: String) {
        let resp = await sendAndReceive(type: "delete_conversation", params: [
            "conversation_id": convId,
        ])
        if resp.string(for: "status") == "ok" {
            senderKeyStates.removeValue(forKey: convId)
            KeyStorage.deleteSenderKeyState(email: email, convId: convId)
            KeyStorage.deleteRecvSenderKeys(email: email, convId: convId)
            return (true, "Deleted")
        }
        let msg = resp.dict(for: "data")?.string(for: "message") ?? "Failed"
        return (false, msg)
    }

    // MARK: - Invitations

    func acceptInvitation(convId: String) async -> (success: Bool, message: String) {
        let resp = await sendAndReceive(type: "accept_invitation", params: [
            "conversation_id": convId,
        ])
        let msg = resp.dict(for: "data")?.string(for: "message") ?? ""
        return (resp.string(for: "status") == "ok", msg)
    }

    func declineInvitation(convId: String) async -> (success: Bool, message: String) {
        let resp = await sendAndReceive(type: "decline_invitation", params: [
            "conversation_id": convId,
        ])
        let msg = resp.dict(for: "data")?.string(for: "message") ?? ""
        return (resp.string(for: "status") == "ok", msg)
    }

    func listInvitations() async -> [Invitation] {
        let resp = await sendAndReceive(type: "list_invitations")
        guard resp.string(for: "status") == "ok",
              let data = resp.dict(for: "data"),
              let invList = data["invitations"] as? [[String: Any]] else {
            return []
        }

        return invList.compactMap { dict -> Invitation? in
            guard let convId = dict.string(for: "conversation_id") else { return nil }
            return Invitation(
                id: dict.string(for: "id") ?? convId,
                conversationId: convId,
                conversationName: dict.string(for: "conversation_name") ?? "Group",
                invitedBy: dict.string(for: "invited_by") ?? "",
                invitedByUsername: dict.string(for: "invited_by_username") ?? ""
            )
        }
    }

    // MARK: - Profile

    func getProfile(userId: String? = nil) async -> UserProfile? {
        var params: [String: Any] = [:]
        if let userId = userId {
            params["user_id"] = userId
        }
        let resp = await sendAndReceive(type: "get_profile", params: params)
        guard resp.string(for: "status") == "ok",
              let data = resp.dict(for: "data") else { return nil }

        return UserProfile(
            userId: data.string(for: "user_id") ?? userId ?? self.userId ?? "",
            username: data.string(for: "username"),
            email: data.string(for: "email"),
            phone: data.string(for: "phone"),
            phoneVisible: data.bool(for: "phone_visible") ?? false,
            location: data.string(for: "location"),
            locationVisible: data.bool(for: "location_visible") ?? false,
            avatarFile: data.string(for: "avatar_file")
        )
    }

    func updateProfile(phone: String? = nil, phoneVisible: Bool? = nil,
                       location: String? = nil, locationVisible: Bool? = nil) async -> Bool {
        var params: [String: Any] = [:]
        if let phone = phone { params["phone"] = phone }
        if let phoneVisible = phoneVisible { params["phone_visible"] = phoneVisible }
        if let location = location { params["location"] = location }
        if let locationVisible = locationVisible { params["location_visible"] = locationVisible }

        let resp = await sendAndReceive(type: "update_profile", params: params)
        return resp.string(for: "status") == "ok"
    }

    func updateAvatar(imageData: Data) async -> (Bool, String) {
        // Resize avatar to fit within protocol message limit (64KB total)
        guard let uiImage = UIImage(data: imageData) else { return (false, "Invalid image") }

        let maxDim: CGFloat = 256
        let scale = min(1.0, maxDim / max(uiImage.size.width, uiImage.size.height))
        let newSize = CGSize(width: uiImage.size.width * scale, height: uiImage.size.height * scale)
        let renderer = UIGraphicsImageRenderer(size: newSize)
        let processed = renderer.image { _ in
            uiImage.draw(in: CGRect(origin: .zero, size: newSize))
        }

        guard let jpegData = processed.jpegData(compressionQuality: 0.6) else { return (false, "JPEG encode failed") }
        #if DEBUG
        print("DEBUG updateAvatar: jpegData=\(jpegData.count) bytes, base64=\(jpegData.count * 4 / 3) bytes approx")
        #endif

        let resp = await sendAndReceive(type: "update_avatar", params: [
            "data": ProtocolHandler.encodeBinary(jpegData),
        ])
        #if DEBUG
        print("DEBUG updateAvatar: resp=\(resp)")
        #endif
        if resp.string(for: "status") == "ok" {
            return (true, "")
        }
        let errMsg = resp.dict(for: "data")?.string(for: "message") ?? "Upload failed"
        return (false, errMsg)
    }

    func getAvatar(userId: String) async -> Data? {
        let resp = await sendAndReceive(type: "get_avatar", params: ["user_id": userId])
        guard resp.string(for: "status") == "ok",
              let data = resp.dict(for: "data"),
              let avatarB64 = data.string(for: "data"),
              let avatarData = try? ProtocolHandler.decodeBinary(avatarB64) else {
            return nil
        }
        return avatarData
    }

    // MARK: - Group Avatar

    func updateGroupAvatar(convId: String, imageData: Data) async -> Bool {
        guard let uiImage = UIImage(data: imageData) else { return false }

        let maxDim: CGFloat = 256
        let scale = min(1.0, maxDim / max(uiImage.size.width, uiImage.size.height))
        let newSize = CGSize(width: uiImage.size.width * scale, height: uiImage.size.height * scale)
        let renderer = UIGraphicsImageRenderer(size: newSize)
        let processed = renderer.image { _ in
            uiImage.draw(in: CGRect(origin: .zero, size: newSize))
        }

        guard let jpegData = processed.jpegData(compressionQuality: 0.6) else { return false }

        let resp = await sendAndReceive(type: "update_group_avatar", params: [
            "conversation_id": convId,
            "data": ProtocolHandler.encodeBinary(jpegData),
        ])
        return resp.string(for: "status") == "ok"
    }

    func getGroupAvatar(convId: String) async -> Data? {
        let resp = await sendAndReceive(type: "get_group_avatar", params: ["conversation_id": convId])
        guard resp.string(for: "status") == "ok",
              let data = resp.dict(for: "data"),
              let avatarB64 = data.string(for: "data"),
              let avatarData = try? ProtocolHandler.decodeBinary(avatarB64) else {
            return nil
        }
        return avatarData
    }

    // MARK: - Key Rotation

    func rotateKeys(password: String) async -> (success: Bool, message: String) {
        guard rsaPrivate != nil, userId != nil else {
            return (false, "Not logged in.")
        }

        // Generate new RSA-4096 keypair
        let newPriv: SecKey
        let newPub: SecKey
        do {
            (newPriv, newPub) = try RSACrypto.generateKeypair()
        } catch {
            return (false, "Key generation failed: \(error.localizedDescription)")
        }

        // Save locally (password-encrypted)
        var pwdBytes = Array(password.utf8)
        defer { pwdBytes.withUnsafeMutableBytes { ptr in _ = memset(ptr.baseAddress!, 0, ptr.count) } }
        do {
            try KeyStorage.saveRSAKeys(email: email, privateKey: newPriv, publicKey: newPub, password: Data(pwdBytes))
        } catch {
            return (false, "Failed to save keys: \(error.localizedDescription)")
        }

        // Update in-memory
        self.rsaPrivate = newPriv
        self.rsaPublic = newPub

        // Serialize public key to PEM and send to server
        guard let pubPEM = try? RSACrypto.serializePublicKey(newPub),
              let pubPEMStr = String(data: pubPEM, encoding: .utf8) else {
            return (false, "Failed to serialize public key")
        }

        let resp = await sendAndReceive(type: "rotate_keys", params: ["public_key": pubPEMStr])
        if resp.string(for: "status") == "ok" {
            return (true, "Keys rotated. Other devices will be disconnected.")
        }
        let errMsg = resp.dict(for: "data")?.string(for: "message") ?? "Rotation failed"
        return (false, errMsg)
    }

    // MARK: - Device Pairing

    private var pairingTempPrivateKey: SecKey?
    private var pairingPollToken: String = ""

    /// Start pairing on NEW device: generates temp RSA-2048 key, sends to server, returns 8-digit code
    func pairingStart(email: String) async -> (success: Bool, codeOrMessage: String) {
        // Generate ephemeral RSA-2048 keypair
        let tempPriv: SecKey
        let tempPub: SecKey
        do {
            (tempPriv, tempPub) = try RSACrypto.generateKeypair2048()
        } catch {
            return (false, "Temp key generation failed")
        }
        self.pairingTempPrivateKey = tempPriv

        // Serialize temp public key to PEM
        guard let tempPubPEM = try? RSACrypto.serializePublicKey(tempPub),
              let tempPubStr = String(data: tempPubPEM, encoding: .utf8) else {
            return (false, "Failed to serialize temp key")
        }

        let resp = await sendAndReceive(type: "pairing_start", params: [
            "email": email,
            "temp_public_key": tempPubStr,
        ])

        guard resp.string(for: "status") == "ok",
              let data = resp.dict(for: "data"),
              let code = data.string(for: "code") else {
            let msg = resp.dict(for: "data")?.string(for: "message") ?? "Pairing start failed"
            return (false, msg)
        }

        self.pairingPollToken = data.string(for: "poll_token") ?? ""
        return (true, code)
    }

    /// Wait for pairing on NEW device: polls server every 2s, decrypts keys when ready
    func pairingWait(code: String, email: String, password: String, timeout: TimeInterval = 300) async -> (success: Bool, message: String) {
        guard pairingTempPrivateKey != nil else {
            return (false, "Pairing not started.")
        }

        let deadline = Date().addingTimeInterval(timeout)

        while Date() < deadline {
            let resp = await sendAndReceive(type: "pairing_poll", params: [
                "code": code,
                "poll_token": pairingPollToken,
            ])

            guard resp.string(for: "status") == "ok",
                  let data = resp.dict(for: "data") else {
                let msg = resp.dict(for: "data")?.string(for: "message") ?? "Poll failed"
                return (false, msg)
            }

            guard data.bool(for: "ready") == true else {
                try? await Task.sleep(for: .seconds(2))
                continue
            }

            // Payload received — decrypt keys
            guard let payload = data["payload"] as? [String: Any],
                  let encKeyB64 = payload.string(for: "encrypted_key"),
                  let ivB64 = payload.string(for: "iv"),
                  let ctB64 = payload.string(for: "ciphertext"),
                  let tagB64 = payload.string(for: "tag") else {
                return (false, "Invalid payload format")
            }

            do {
                let encAESKey = try ProtocolHandler.decodeBinary(encKeyB64)
                let iv = try ProtocolHandler.decodeBinary(ivB64)
                let ct = try ProtocolHandler.decodeBinary(ctB64)
                let tag = try ProtocolHandler.decodeBinary(tagB64)

                // RSA-OAEP decrypt the AES key
                let aesKey = try RSACrypto.decrypt(pairingTempPrivateKey!, ciphertext: encAESKey)

                // AES-GCM decrypt the keys JSON
                let keysJSON = try CryptoUtils.aesDecrypt(key: aesKey, nonce: iv, ciphertext: ct, tag: tag)

                guard let keysDict = try JSONSerialization.jsonObject(with: keysJSON) as? [String: String],
                      let rsaPrivPEM = keysDict["rsa_private"],
                      let identityHex = keysDict["identity_private"] else {
                    return (false, "Invalid keys JSON")
                }

                // Import RSA key
                var pwdBytes = Array(password.utf8)
                defer { pwdBytes.withUnsafeMutableBytes { ptr in _ = memset(ptr.baseAddress!, 0, ptr.count) } }
                let pwdData = Data(pwdBytes)

                let rsaPriv = try RSACrypto.loadPrivateKey(Data(rsaPrivPEM.utf8), password: nil)
                let rsaPub = SecKeyCopyPublicKey(rsaPriv)!

                // Import Ed25519 identity key
                guard let identityRaw = Data(hexString: identityHex), identityRaw.count == 32 else {
                    return (false, "Invalid identity key format")
                }
                let edPriv = try Curve25519.Signing.PrivateKey(rawRepresentation: identityRaw)
                let edPub = edPriv.publicKey

                // Save to disk (encrypted with password)
                try KeyStorage.saveRSAKeys(email: email, privateKey: rsaPriv, publicKey: rsaPub, password: pwdData)
                try KeyStorage.saveIdentityKeys(email: email, privateKey: edPriv, publicKey: edPub, password: pwdData)

                // Update in-memory state
                self.email = email
                self.rsaPrivate = rsaPriv
                self.rsaPublic = rsaPub
                self.identityPrivate = edPriv
                self.identityPublic = edPub
                self.cacheKey = CryptoUtils.deriveSelfEncryptionKey(identityPrivateRaw: edPriv.rawData)
                self.localKey = CryptoUtils.deriveLocalStorageKey(identityPrivateRaw: edPriv.rawData)

                // Clear pairing state
                self.pairingTempPrivateKey = nil
                self.pairingPollToken = ""

                return (true, "Pairing complete.")
            } catch {
                return (false, "Failed to import keys: \(error.localizedDescription)")
            }
        }

        return (false, "Pairing timed out.")
    }

    /// Re-encrypt all cached messages with self-encryption key so a paired device can read them.
    /// Called by authorizeDevice() before sending keys.
    func reencryptHistory() async {
        guard let edPriv = identityPrivate else { return }
        let selfKey = CryptoUtils.deriveSelfEncryptionKey(identityPrivateRaw: edPriv.rawData)

        // Phase 1: Fetch all messages to populate cache (decrypts and caches them)
        let convs = await listConversations()
        for conv in convs {
            var offset = 0
            while true {
                let msgs = await getMessages(convId: conv.id, limit: 200, offset: offset)
                if msgs.count < 200 { break }
                offset += msgs.count
            }
        }

        // Phase 2: Read per-message cache and re-encrypt with self-key
        var allUpdates: [[String: String]] = []
        for conv in convs {
            let cachedMessages = MessageCache.loadAllCachedMessages(email: email, convId: conv.id, cacheKey: cacheKey)
            for (msgId, plaintext) in cachedMessages {
                // Skip control messages
                if let jsonObj = try? JSONSerialization.jsonObject(with: plaintext) as? [String: Any] {
                    if jsonObj["_control"] != nil { continue }
                    let text = jsonObj["text"] as? String
                    let image = jsonObj["image"]
                    let file = jsonObj["file"]
                    if text == nil && image == nil && file == nil { continue }
                }

                // Re-encrypt with self-encryption key
                guard let (_, nonce, ct, tag) = try? CryptoUtils.aesEncrypt(plaintext, key: selfKey) else { continue }

                allUpdates.append([
                    "message_id": msgId,
                    "encrypted_content": ProtocolHandler.encodeBinary(ct + tag),
                    "nonce": ProtocolHandler.encodeBinary(nonce),
                ])
            }
        }

        if allUpdates.isEmpty { return }

        // Send in batches of 500
        let batchSize = 500
        for start in stride(from: 0, to: allUpdates.count, by: batchSize) {
            let end = min(start + batchSize, allUpdates.count)
            let batch = Array(allUpdates[start..<end])
            let resp = await sendAndReceive(type: "reencrypt_messages", params: ["updates": batch])
            if resp.string(for: "status") == "ok" {
                #if DEBUG
                print("DEBUG reencryptHistory: re-encrypted \(end)/\(allUpdates.count) messages")
                #endif
            } else {
                #if DEBUG
                print("DEBUG reencryptHistory: batch failed: \(resp.dict(for: "data")?.string(for: "message") ?? "unknown")")
                #endif
            }
        }
        #if DEBUG
        print("DEBUG reencryptHistory: complete, \(allUpdates.count) messages re-encrypted")
        #endif
    }

    /// Authorize a new device from the EXISTING device: encrypts and sends keys
    func authorizeDevice(code: String) async -> (success: Bool, message: String) {
        guard let rsaPriv = rsaPrivate, let edPriv = identityPrivate else {
            return (false, "Not logged in.")
        }

        // Claim the pairing code — get temp public key
        let claimResp = await sendAndReceive(type: "pairing_claim", params: ["code": code])
        guard claimResp.string(for: "status") == "ok",
              let claimData = claimResp.dict(for: "data"),
              let tempPubPEM = claimData.string(for: "temp_public_key") else {
            let msg = claimResp.dict(for: "data")?.string(for: "message") ?? "Claim failed"
            return (false, msg)
        }

        // Load temp public key
        let tempPub: SecKey
        do {
            tempPub = try RSACrypto.loadPublicKey(Data(tempPubPEM.utf8))
        } catch {
            return (false, "Invalid temp public key: \(error.localizedDescription)")
        }

        // Re-encrypt message history so new device can read old messages
        await reencryptHistory()

        // Build keys payload (RSA private unencrypted PEM + Ed25519 raw hex)
        let rsaPrivPEM: String
        do {
            let pemData = try RSACrypto.serializePrivateKey(rsaPriv, password: nil)
            rsaPrivPEM = String(data: pemData, encoding: .utf8) ?? ""
        } catch {
            return (false, "Failed to serialize RSA key")
        }

        let identityHex = edPriv.rawRepresentation.map { String(format: "%02x", $0) }.joined()

        let keysDict: [String: String] = [
            "rsa_private": rsaPrivPEM,
            "identity_private": identityHex,
        ]

        guard let keysJSON = try? JSONSerialization.data(withJSONObject: keysDict) else {
            return (false, "JSON serialization failed")
        }

        // Encrypt: AES-GCM for payload, RSA-OAEP for AES key
        do {
            let (aesKey, nonce, ct, tag) = try CryptoUtils.aesEncrypt(keysJSON)
            let encAESKey = try RSACrypto.encrypt(tempPub, plaintext: aesKey)

            let payload: [String: String] = [
                "encrypted_key": ProtocolHandler.encodeBinary(encAESKey),
                "iv": ProtocolHandler.encodeBinary(nonce),
                "ciphertext": ProtocolHandler.encodeBinary(ct),
                "tag": ProtocolHandler.encodeBinary(tag),
            ]

            let resp = await sendAndReceive(type: "pairing_send", params: [
                "code": code,
                "payload": payload,
            ])

            if resp.string(for: "status") == "ok" {
                return (true, "Device authorized.")
            }
            let msg = resp.dict(for: "data")?.string(for: "message") ?? "Send failed"
            return (false, msg)
        } catch {
            return (false, "Encryption failed: \(error.localizedDescription)")
        }
    }

    // MARK: - File Sharing

    func sendFile(convId: String, fileData: Data, filename: String, mimeType: String,
                  members: [ConversationMember], replyTo: String? = nil) async -> (success: Bool, message: String, sentMessage: Message?) {
        // Encrypt file with AES-GCM
        guard let (aesKey, nonce, ct, tag) = try? CryptoUtils.aesEncrypt(fileData) else {
            return (false, "File encryption failed", nil)
        }

        let encryptedData = ct + tag
        let fileType = mimeType.hasPrefix("image/") ? "image" : "file"

        // Start upload
        let fileId = UUID().uuidString.lowercased()
        let startResp = await sendAndReceive(type: "upload_image_start", params: [
            "conversation_id": convId,
            "file_id": fileId,
            "file_size": encryptedData.count,
            "file_type": fileType,
        ])
        guard startResp.string(for: "status") == "ok" else {
            let msg = startResp.dict(for: "data")?.string(for: "message") ?? "Upload start failed"
            return (false, msg, nil)
        }

        // Upload chunks
        var offset = 0
        while offset < encryptedData.count {
            let end = min(offset + Constants.imageChunkSize, encryptedData.count)
            let chunk = encryptedData[offset..<end]
            let chunkNum = offset / Constants.imageChunkSize + 1
            let chunkResp = await sendAndReceive(type: "upload_image_chunk", timeout: 60, params: [
                "file_id": fileId,
                "data": ProtocolHandler.encodeBinary(Data(chunk)),
                "offset": offset,
            ])
            guard chunkResp.string(for: "status") == "ok" else {
                return (false, "Chunk upload failed (chunk \(chunkNum))", nil)
            }
            offset = end
        }

        // End upload
        let endResp = await sendAndReceive(type: "upload_image_end", timeout: 60, params: [
            "file_id": fileId,
        ])
        guard endResp.string(for: "status") == "ok" else {
            return (false, "Upload completion failed", nil)
        }

        // Send message with file info
        // Keys must be base64 encoded to match Python server
        let fileInfo: [String: Any] = [
            "file_id": fileId,
            "aes_key": ProtocolHandler.encodeBinary(aesKey),
            "iv": ProtocolHandler.encodeBinary(nonce),
            "filename": filename,
            "size": fileData.count,
            "mime_type": mimeType,
        ]

        let extra: [String: Any] = ["file": fileInfo]
        let (success, msg, sentMessage) = await sendMessage(
            convId: convId, text: "", members: members, replyTo: replyTo,
            extraPayload: extra, imageFileId: fileId
        )
        return (success, msg, sentMessage)
    }

    func sendImage(convId: String, imageData: Data, members: [ConversationMember],
                   replyTo: String? = nil) async -> (success: Bool, message: String, sentMessage: Message?) {

        // Step 1: Process image on MainActor (UIKit requires main thread)
        let processed = await MainActor.run { () -> (imageBytes: Data, thumbnailB64: String, originalSize: Int)? in
            guard let uiImage = UIImage(data: imageData) else { return nil }

            var finalBytes = imageData
            var img = uiImage
            // AES-GCM overhead: 12 nonce + 16 tag = 28 bytes
            let maxPlaintextSize = Constants.maxImageBytes - 28

            // If too large, progressively compress
            if finalBytes.count > maxPlaintextSize {
                let qualities: [CGFloat] = [0.92, 0.85, 0.75, 0.60]
                var fits = false
                for quality in qualities {
                    if let compressed = img.jpegData(compressionQuality: quality) {
                        finalBytes = compressed
                        if compressed.count <= maxPlaintextSize { fits = true; break }
                    }
                }

                if !fits {
                    // Downscale dimensions
                    for maxDim: CGFloat in [3840, 2560, 1920, 1280] {
                        if max(img.size.width, img.size.height) > maxDim {
                            let scale = maxDim / max(img.size.width, img.size.height)
                            let newSize = CGSize(width: img.size.width * scale, height: img.size.height * scale)
                            let renderer = UIGraphicsImageRenderer(size: newSize)
                            img = renderer.image { _ in img.draw(in: CGRect(origin: .zero, size: newSize)) }
                        }
                        if let compressed = img.jpegData(compressionQuality: 0.75) {
                            finalBytes = compressed
                            if compressed.count <= maxPlaintextSize { break }
                        }
                    }
                }
            }

            // Generate thumbnail (200x200, quality 60)
            let thumbScale = min(200.0 / max(img.size.width, img.size.height), 1.0)
            let thumbDim = CGSize(width: img.size.width * thumbScale, height: img.size.height * thumbScale)
            let thumbRenderer = UIGraphicsImageRenderer(size: thumbDim)
            let thumbImg = thumbRenderer.image { _ in img.draw(in: CGRect(origin: .zero, size: thumbDim)) }
            let thumbB64 = thumbImg.jpegData(compressionQuality: 0.6)?.base64EncodedString() ?? ""

            return (finalBytes, thumbB64, imageData.count)
        }

        guard let processed = processed else {
            return (false, "Invalid image data", nil)
        }

        // Step 2: Encrypt on background (actor isolates this automatically)
        guard let (aesKey, nonce, ct, tag) = try? CryptoUtils.aesEncrypt(processed.imageBytes) else {
            return (false, "Image encryption failed", nil)
        }
        let encryptedData = ct + tag

        // Step 3: Chunked upload
        let fileId = UUID().uuidString.lowercased()
        let startResp = await sendAndReceive(type: "upload_image_start", params: [
            "conversation_id": convId,
            "file_id": fileId,
            "file_size": encryptedData.count,
            "file_type": "image",
        ])
        guard startResp.string(for: "status") == "ok" else {
            let msg = startResp.dict(for: "data")?.string(for: "message") ?? "Upload start failed"
            return (false, msg, nil)
        }

        var offset = 0
        while offset < encryptedData.count {
            let end = min(offset + Constants.imageChunkSize, encryptedData.count)
            let chunk = encryptedData[offset..<end]
            let chunkResp = await sendAndReceive(type: "upload_image_chunk", timeout: 60, params: [
                "file_id": fileId,
                "data": ProtocolHandler.encodeBinary(Data(chunk)),
                "offset": offset,
            ])
            guard chunkResp.string(for: "status") == "ok" else {
                return (false, "Chunk upload failed", nil)
            }
            offset = end
        }

        let endResp = await sendAndReceive(type: "upload_image_end", timeout: 60, params: ["file_id": fileId])
        guard endResp.string(for: "status") == "ok" else {
            return (false, "Upload completion failed", nil)
        }

        // Step 4: Send message with image info
        let imageInfo: [String: Any] = [
            "file_id": fileId,
            "aes_key": ProtocolHandler.encodeBinary(aesKey),
            "iv": ProtocolHandler.encodeBinary(nonce),
            "thumbnail": processed.thumbnailB64,
            "filename": "image.jpg",
            "size": processed.originalSize,
        ]

        let extra: [String: Any] = ["image": imageInfo]
        let (success, msg, sentMessage) = await sendMessage(
            convId: convId, text: "", members: members, replyTo: replyTo,
            extraPayload: extra, imageFileId: fileId
        )
        return (success, msg, sentMessage)
    }

    func downloadFile(fileId: String, aesKey: Data, iv: Data, conversationId: String? = nil) async -> Data? {
        #if DEBUG
        print("DEBUG downloadFile: starting download for fileId=\(fileId) convId=\(conversationId ?? "nil") aesKeyLen=\(aesKey.count) ivLen=\(iv.count)")
        #endif
        var allData = Data()
        var offset = 0

        while true {
            var params: [String: Any] = [
                "file_id": fileId,
                "offset": offset,
            ]
            if let conversationId = conversationId {
                params["conversation_id"] = conversationId
            }
            let resp = await sendAndReceive(type: "download_image", params: params)

            let status = resp.string(for: "status")
            guard status == "ok" else {
                let errMsg = resp.dict(for: "data")?.string(for: "message") ?? "unknown"
                #if DEBUG
                print("DEBUG downloadFile: server error status=\(status ?? "nil") msg=\(errMsg)")
                #endif
                break
            }
            guard let data = resp.dict(for: "data") else {
                #if DEBUG
                print("DEBUG downloadFile: missing 'data' dict in response, keys=\(Array(resp.keys))")
                #endif
                break
            }
            guard let chunkB64 = data.string(for: "data") else {
                #if DEBUG
                print("DEBUG downloadFile: missing 'data' string in data dict, keys=\(Array(data.keys))")
                #endif
                break
            }
            guard let chunk = try? ProtocolHandler.decodeBinary(chunkB64) else {
                #if DEBUG
                print("DEBUG downloadFile: base64 decode failed, chunkB64 length=\(chunkB64.count)")
                #endif
                break
            }

            if chunk.isEmpty {
                #if DEBUG
                print("DEBUG downloadFile: empty chunk received, done")
                #endif
                break
            }
            allData.append(chunk)
            offset += chunk.count
            #if DEBUG
            print("DEBUG downloadFile: received chunk \(chunk.count) bytes, total=\(allData.count)")
            #endif

            if data.bool(for: "done") == true {
                #if DEBUG
                print("DEBUG downloadFile: server signaled done")
                #endif
                break
            }
        }

        guard !allData.isEmpty else {
            #if DEBUG
            print("DEBUG downloadFile: no data downloaded for fileId=\(fileId)")
            #endif
            return nil
        }

        // Decrypt: allData = ciphertext + tag(16)
        guard allData.count >= 16 else {
            #if DEBUG
            print("DEBUG downloadFile: data too short (\(allData.count) bytes) for fileId=\(fileId)")
            #endif
            return nil
        }
        let ct = allData.prefix(allData.count - 16)
        let tag = allData.suffix(16)
        #if DEBUG
        print("DEBUG downloadFile: decrypting \(ct.count) bytes ciphertext + 16 bytes tag")
        #endif
        do {
            let result = try CryptoUtils.aesDecrypt(key: aesKey, nonce: iv, ciphertext: Data(ct), tag: Data(tag))
            #if DEBUG
            print("DEBUG downloadFile: decryption success, \(result.count) bytes")
            #endif
            return result
        } catch {
            #if DEBUG
            print("DEBUG downloadFile: decryption FAILED: \(error)")
            #endif
            return nil
        }
    }

    // MARK: - Devices

    func listDevices() async -> [[String: Any]] {
        let resp = await sendAndReceive(type: "list_devices")
        guard resp.string(for: "status") == "ok",
              let data = resp.dict(for: "data") else { return [] }
        return data["devices"] as? [[String: Any]] ?? []
    }

    func removeDevice(deviceIdToRemove: String) async -> Bool {
        let resp = await sendAndReceive(type: "remove_device", params: [
            "device_id": deviceIdToRemove,
        ])
        return resp.string(for: "status") == "ok"
    }

    // MARK: - Session Reset

    func resetSession(peerUserId: String, peerDeviceId: String? = nil) async {
        if let peerDeviceId = peerDeviceId {
            let sessionKey = "\(peerUserId):\(peerDeviceId)"
            sessions.removeValue(forKey: sessionKey)
            KeyStorage.deleteSession(email: email, peerUserId: peerUserId, peerDeviceId: peerDeviceId)
        } else {
            // Delete all sessions for this user
            for key in sessions.keys where key.hasPrefix(peerUserId) {
                sessions.removeValue(forKey: key)
            }
            KeyStorage.deleteSession(email: email, peerUserId: peerUserId)
        }

        _ = await sendAndReceive(type: "session_reset", params: [
            "peer_user_id": peerUserId,
            "peer_device_id": peerDeviceId ?? "",
        ])
    }

    func handleSessionResetNotification(fromUserId: String, fromDeviceId: String?) {
        if let deviceId = fromDeviceId {
            let sessionKey = "\(fromUserId):\(deviceId)"
            sessions.removeValue(forKey: sessionKey)
            KeyStorage.deleteSession(email: email, peerUserId: fromUserId, peerDeviceId: deviceId)
        } else {
            for key in sessions.keys where key.hasPrefix(fromUserId) {
                sessions.removeValue(forKey: key)
            }
            KeyStorage.deleteSession(email: email, peerUserId: fromUserId)
        }
    }

    // MARK: - Search

    func searchMessages(convId: String, query: String) -> [[String: Any]] {
        MessageCache.search(email: email, convId: convId, query: query, cacheKey: cacheKey)
    }

    // MARK: - Change Username

    func changeUsername(newUsername: String) async -> (success: Bool, message: String) {
        let trimmed = newUsername.trimmingCharacters(in: .whitespacesAndNewlines)
        guard !trimmed.isEmpty, trimmed.count <= 100 else {
            return (false, "Username must be 1-100 characters.")
        }
        let resp = await sendAndReceive(type: "change_username", params: ["username": trimmed])
        if resp.string(for: "status") == "ok" {
            if let data = resp.dict(for: "data"), let newName = data.string(for: "username") {
                username = newName
            } else {
                username = trimmed
            }
            return (true, "Username changed.")
        }
        let msg = resp.dict(for: "data")?.string(for: "message") ?? "Unknown error"
        return (false, msg)
    }

    // MARK: - Change Password

    func changePassword(oldPassword: String, newPassword: String) -> (success: Bool, message: String) {
        guard !email.isEmpty else {
            return (false, "Not logged in.")
        }

        let oldPwd = Data(oldPassword.utf8)
        let newPwd = Data(newPassword.utf8)

        // 1. Verify old password by loading keys
        let (rsaPriv, rsaPub, err) = KeyStorage.loadRSAKeys(email: email, password: oldPwd)
        guard let rsaPriv = rsaPriv, let rsaPub = rsaPub else {
            return (false, err ?? "Wrong current password.")
        }

        let (edPriv, edPub) = KeyStorage.loadIdentityKeys(email: email, password: oldPwd)
        guard let edPriv = edPriv, let edPub = edPub else {
            return (false, "Failed to load identity key.")
        }

        // 2. Re-save with new password
        do {
            try KeyStorage.saveRSAKeys(email: email, privateKey: rsaPriv, publicKey: rsaPub, password: newPwd)
            try KeyStorage.saveIdentityKeys(email: email, privateKey: edPriv, publicKey: edPub, password: newPwd)
            return (true, "Password changed successfully.")
        } catch {
            return (false, "Failed to save keys: \(error.localizedDescription)")
        }
    }

    // MARK: - Confirm Delivery

    func confirmDelivery(convId: String, messageIds: [String]) async {
        guard !messageIds.isEmpty else { return }
        // Fire-and-forget, non-critical
        _ = await sendAndReceive(type: "confirm_delivery", params: [
            "conversation_id": convId,
            "message_ids": messageIds,
        ])
    }

    // MARK: - Forward Message

    func forwardMessage(targetConvId: String, originalMsg: [String: Any],
                        targetMembers: [ConversationMember]) async -> (success: Bool, message: String, sentMessage: Message?) {
        let text = originalMsg["text"] as? String ?? ""

        var payload: [String: Any] = [
            "sender": username,
            "text": text,
            "forwarded_from": [
                "sender": originalMsg["sender"] as? String ?? "",
                "conversation_id": originalMsg["conversation_id"] as? String ?? "",
                "message_id": originalMsg["message_id"] as? String ?? "",
            ] as [String: Any],
            "timestamp": ISO8601DateFormatter().string(from: Date()),
        ]

        // Forward image/file metadata (the encrypted blob is already on the server)
        if let image = originalMsg["image"] {
            payload["image"] = image
            if text.isEmpty { payload["text"] = "" }
        }
        if let file = originalMsg["file"] {
            payload["file"] = file
            if text.isEmpty { payload["text"] = "" }
        }

        var extraPayload: [String: Any] = [:]
        if let fwd = payload["forwarded_from"] {
            extraPayload["forwarded_from"] = fwd
        }
        if let image = payload["image"] {
            extraPayload["image"] = image
        }
        if let file = payload["file"] {
            extraPayload["file"] = file
        }

        return await sendMessage(
            convId: targetConvId, text: text, members: targetMembers,
            extraPayload: extraPayload
        )
    }

    // MARK: - Proof of Work

    static func solvePow(challenge: String, difficulty: Int) -> String {
        let targetBytes = difficulty / 8
        let targetBits = difficulty % 8
        let mask: UInt8 = targetBits > 0 ? UInt8((0xFF << (8 - targetBits)) & 0xFF) : 0

        var nonce: UInt64 = 0
        while true {
            let input = "\(challenge)\(nonce)"
            let inputData = Data(input.utf8)
            let digest = SHA256.hash(data: inputData)
            let digestBytes = Array(digest)

            var ok = true
            for i in 0..<targetBytes {
                if digestBytes[i] != 0 {
                    ok = false
                    break
                }
            }
            if ok && targetBits > 0 {
                if digestBytes[targetBytes] & mask != 0 {
                    ok = false
                }
            }
            if ok {
                return String(nonce)
            }
            nonce += 1
        }
    }

    // MARK: - TOFU / Contact Verification

    /// Load TOFU and verification stores from disk
    private func loadVerificationStores() {
        guard !email.isEmpty else { return }
        knownIdentityKeys = KeyStorage.loadKnownIdentityKeys(email: email, localKey: localKey)
        verifiedContacts = KeyStorage.loadVerifiedContacts(email: email, localKey: localKey)
    }

    /// Check a user's identity key against TOFU registry
    func checkIdentityKey(userId: String, identityKeyBytes: Data) -> String {
        let ikHex = identityKeyBytes.hexString
        let now = ISO8601DateFormatter().string(from: Date())
        let known = knownIdentityKeys[userId]

        if known == nil {
            // TOFU: trust on first use
            knownIdentityKeys[userId] = [
                "identity_key": ikHex,
                "first_seen": now,
                "last_seen": now,
            ]
            try? KeyStorage.saveKnownIdentityKeys(email: email, keys: knownIdentityKeys, localKey: localKey)
            return "new"
        }

        if known?["identity_key"] == ikHex {
            knownIdentityKeys[userId]?["last_seen"] = now
            try? KeyStorage.saveKnownIdentityKeys(email: email, keys: knownIdentityKeys, localKey: localKey)
            let verified = verifiedContacts[userId]
            if verified != nil && verified?["identity_key"] == ikHex {
                return "verified"
            }
            return "trusted"
        }

        // Key has CHANGED
        let wasVerified = verifiedContacts[userId] != nil
        return wasVerified ? "changed_verified" : "changed"
    }

    /// Mark a contact's identity key as explicitly verified
    func verifyContact(userId: String, identityKey: Data, method: String = "manual") {
        let ikHex = identityKey.hexString
        let now = ISO8601DateFormatter().string(from: Date())
        verifiedContacts[userId] = [
            "identity_key": ikHex,
            "verified_at": now,
            "method": method,
        ]
        if knownIdentityKeys[userId] == nil {
            knownIdentityKeys[userId] = [
                "identity_key": ikHex,
                "first_seen": now,
                "last_seen": now,
            ]
        } else {
            knownIdentityKeys[userId]?["last_seen"] = now
        }
        try? KeyStorage.saveVerifiedContacts(email: email, contacts: verifiedContacts, localKey: localKey)
        try? KeyStorage.saveKnownIdentityKeys(email: email, keys: knownIdentityKeys, localKey: localKey)
    }

    /// Remove explicit verification for a contact
    func unverifyContact(userId: String) {
        verifiedContacts.removeValue(forKey: userId)
        try? KeyStorage.saveVerifiedContacts(email: email, contacts: verifiedContacts, localKey: localKey)
    }

    /// Accept a changed identity key
    func acceptKeyChange(userId: String, newIdentityKey: Data) {
        let ikHex = newIdentityKey.hexString
        let now = ISO8601DateFormatter().string(from: Date())
        knownIdentityKeys[userId] = [
            "identity_key": ikHex,
            "first_seen": now,
            "last_seen": now,
        ]
        verifiedContacts.removeValue(forKey: userId)
        try? KeyStorage.saveKnownIdentityKeys(email: email, keys: knownIdentityKeys, localKey: localKey)
        try? KeyStorage.saveVerifiedContacts(email: email, contacts: verifiedContacts, localKey: localKey)
    }

    /// Get verification status for a user
    func getVerificationStatus(userId: String) -> String {
        if let verified = verifiedContacts[userId] {
            let known = knownIdentityKeys[userId]
            if known?["identity_key"] == verified["identity_key"] {
                return "verified"
            }
        }
        if knownIdentityKeys[userId] != nil {
            return "trusted"
        }
        return "unverified"
    }

    /// Get formatted safety number for a peer
    func getSafetyNumber(peerUserId: String) -> String? {
        guard let idPub = identityPublic, let myId = userId else { return nil }
        let myIK = idPub.rawData
        guard let peerIK = getPeerIdentityKeySync(userId: peerUserId) else { return nil }
        return ContactVerification.computeSafetyNumber(
            myUserId: myId, myIdentityKey: myIK,
            theirUserId: peerUserId, theirIdentityKey: peerIK
        )
    }

    /// Get formatted fingerprint for own identity key
    func getMyFingerprint() -> String? {
        guard let idPub = identityPublic, let myId = userId else { return nil }
        let fp = ContactVerification.computeFingerprint(userId: myId, identityKey: idPub.rawData)
        return ContactVerification.formatFingerprint(fp)
    }

    /// Get formatted fingerprint for a peer's identity key
    func getPeerFingerprint(peerUserId: String) -> String? {
        guard let peerIK = getPeerIdentityKeySync(userId: peerUserId) else { return nil }
        let fp = ContactVerification.computeFingerprint(userId: peerUserId, identityKey: peerIK)
        return ContactVerification.formatFingerprint(fp)
    }

    /// Get QR code payload bytes for own identity
    func getVerificationQRData() -> Data? {
        guard let idPub = identityPublic, let myId = userId else { return nil }
        return ContactVerification.encodeVerificationQR(userId: myId, identityKey: idPub.rawData)
    }

    /// Verify a scanned QR code against known identity keys
    func verifyQRCode(qrData: Data) -> (success: Bool, userId: String, message: String) {
        let decoded: (userId: String, identityKey: Data)
        do {
            decoded = try ContactVerification.decodeVerificationQR(qrData)
        } catch {
            return (false, "", "Invalid QR code: \(error.localizedDescription)")
        }
        guard let peerIK = getPeerIdentityKeySync(userId: decoded.userId) else {
            return (false, decoded.userId, "Unknown user — not in your contacts.")
        }
        guard peerIK == decoded.identityKey else {
            return (false, decoded.userId, "Identity key MISMATCH — verification failed!")
        }
        verifyContact(userId: decoded.userId, identityKey: decoded.identityKey, method: "qr_code")
        let name = userCache[decoded.userId]?.username ?? String(decoded.userId.prefix(8))
        return (true, decoded.userId, "Verified \(name) via QR code.")
    }

    /// Get a peer's identity key bytes (from cache, TOFU registry, or server)
    func getPeerIdentityKey(userId: String) async -> Data? {
        // 1. Check local caches first
        if let ik = getPeerIdentityKeySync(userId: userId) {
            return ik
        }
        // 2. Fetch from server
        if let user = await getUserInfo(userId: userId) {
            return user.identityKey
        }
        return nil
    }

    /// Sync lookup: user cache → TOFU registry
    private func getPeerIdentityKeySync(userId: String) -> Data? {
        if let ik = userCache[userId]?.identityKey {
            return ik
        }
        if let ikHex = knownIdentityKeys[userId]?["identity_key"],
           let ikData = Data(hexString: ikHex) {
            return ikData
        }
        return nil
    }

    // MARK: - Brute-Force Lockout

    static func checkLockout(email: String) -> TimeInterval {
        guard let dir = try? KeyStorage.getKeyDir(email: email) else { return 0 }
        let path = dir.appendingPathComponent("login_lockout.json")
        guard let data = try? Data(contentsOf: path),
              let json = try? JSONSerialization.jsonObject(with: data) as? [String: Any],
              let lockedUntil = json["locked_until"] as? TimeInterval else { return 0 }
        return max(0, lockedUntil - Date().timeIntervalSince1970)
    }

    static func recordFailedAttempt(email: String) {
        guard let dir = try? KeyStorage.getKeyDir(email: email) else { return }
        let path = dir.appendingPathComponent("login_lockout.json")
        var failed = 0
        if let data = try? Data(contentsOf: path),
           let json = try? JSONSerialization.jsonObject(with: data) as? [String: Any] {
            failed = json["failed_attempts"] as? Int ?? 0
        }
        failed += 1
        let delay = min(pow(lockoutBaseSeconds, Double(failed)), lockoutMaxSeconds)
        let lockedUntil = Date().timeIntervalSince1970 + delay
        let json: [String: Any] = ["failed_attempts": failed, "locked_until": lockedUntil]
        if let data = try? JSONSerialization.data(withJSONObject: json) {
            try? data.write(to: path, options: .completeFileProtection)
        }
    }

    static func clearLockout(email: String) {
        guard let dir = try? KeyStorage.getKeyDir(email: email) else { return }
        let path = dir.appendingPathComponent("login_lockout.json")
        try? FileManager.default.removeItem(at: path)
    }
}