import Foundation
import CryptoKit

/// 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 connectionStateChanged(connected: Bool)
}

/// 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(set) var sessionToken: String?
    private(set) var userId: String?
    private(set) var username: String = ""
    private(set) var email: String = ""
    private(set) var loginRejected = false

    // 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 var cacheKey: Data?   // for encrypting message cache
    private 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: - Request/Response Tracking

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

    // MARK: - Notification Stream

    private var notificationContinuation: AsyncStream<ChatNotification>.Continuation?
    nonisolated let notifications: AsyncStream<ChatNotification>

    // MARK: - Init

    init() {
        var continuation: AsyncStream<ChatNotification>.Continuation!
        notifications = AsyncStream { cont in
            continuation = cont
        }
        notificationContinuation = continuation
    }

    // MARK: - Connection

    func connect(host: String = Constants.defaultHost, port: UInt16 = Constants.defaultPort,
                 useTLS: Bool = false, tlsInsecure: Bool = false) async throws {
        try await connectionManager.connect(host: host, port: port, useTLS: useTLS, tlsInsecure: tlsInsecure)
        isConnected = true
        notificationContinuation?.yield(.connectionStateChanged(connected: true))
    }

    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)
        }
        notificationContinuation?.yield(.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()

        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)
                        }
                    }
                }
            }
            return response
        } catch {
            pendingRequests.removeValue(forKey: requestId)
            return [
                "type": type,
                "status": "error",
                "data": ["message": error.localizedDescription]
            ]
        }
    }

    // MARK: - Background Listener

    func startBackgroundListener() {
        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
                    await handleDisconnect()
                    break
                }
                await routeMessage(msg)
            } catch {
                await 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)
        }
        notificationContinuation?.yield(.connectionStateChanged(connected: false))
    }

    private func routeMessage(_ msg: [String: Any]) {
        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"
        ])

        if notificationTypes.contains(msgType) {
            let data = msg["data"] as? [String: Any] ?? msg
            switch msgType {
            case "new_message":
                notificationContinuation?.yield(.newMessage(data: data))
            case "messages_read":
                notificationContinuation?.yield(.messagesRead(data: data))
            case "message_deleted":
                notificationContinuation?.yield(.messageDeleted(data: data))
            case "conversation_created":
                notificationContinuation?.yield(.conversationCreated(data: data))
            case "member_added":
                notificationContinuation?.yield(.memberAdded(data: data))
            case "member_removed":
                notificationContinuation?.yield(.memberRemoved(data: data))
            case "user_online":
                if let uid = data["user_id"] as? String {
                    notificationContinuation?.yield(.userOnline(userId: uid))
                }
            case "user_offline":
                if let uid = data["user_id"] as? String {
                    notificationContinuation?.yield(.userOffline(userId: uid))
                }
            case "online_users":
                if let uids = data["user_ids"] as? [String] {
                    notificationContinuation?.yield(.onlineUsers(userIds: uids))
                }
            case "group_invitation":
                notificationContinuation?.yield(.groupInvitation(data: data))
            case "conversation_renamed":
                notificationContinuation?.yield(.conversationRenamed(data: data))
            case "session_reset":
                notificationContinuation?.yield(.sessionReset(data: data))
            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, err) = 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,
        ])

        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) {
        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)
        }

        // Upload prekeys
        await generateAndUploadPrekeys()

        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)
            for opk in opks {
                opkPrivates[opk.id] = opk.privateKey
                try? KeyStorage.saveOPKPrivate(email: email, opkId: opk.id, privateKey: opk.privateKey)
            }

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

            _ = await sendAndReceive(type: "upload_prekeys", params: [
                "signed_prekey": spkData,
                "one_time_prekeys": otpData,
            ])
        } catch {
            // Log error but don't fail
            print("Prekey generation error: \(error)")
        }
    }

    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 {
            let formatter = ISO8601DateFormatter()
            formatter.formatOptions = [.withInternetDateTime, .withFractionalSeconds]
            if let created = formatter.date(from: spkCreatedAt) ?? ISO8601DateFormatter().date(from: 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
        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
            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)
        }

        // Start background listener
        startBackgroundListener()

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

        // Ensure prekeys in background
        Task { await ensurePrekeys() }

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

    // MARK: - Reconnect

    func reconnect() async -> Bool {
        guard rsaPrivate != nil else { return false }

        await disconnect()

        do {
            try await connect()
        } catch {
            return false
        }

        // 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 {
            return false
        }

        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" else { return false }

        startBackgroundListener()
        return true
    }

    // 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] = []

        // 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) {
                    bundles.append(bundle)
                }
            }
        }
        // Legacy single bundle
        else if let ikHex = data["identity_key"] as? String {
            let bundle = try DeviceBundle.fromDict(data)
            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)"

        // Check memory
        if let session = sessions[sessionKey] {
            return session
        }

        // Check disk
        if let session = KeyStorage.loadSession(
            email: email,
            peerUserId: peerUserId,
            localKey: localKey,
            peerDeviceId: peerDeviceId
        ) {
            sessions[sessionKey] = session
            return session
        }

        // 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, ekPriv, 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
        var x3dhHeader: [String: Any] = [
            "ik": Ed25519Crypto.serializePublic(identityPublic!).hexString,
            "ek": X25519Crypto.serializePublic(ekPub).hexString,
            "spk_id": bundle.spkId,
        ]
        if let opkId = bundle.opkId {
            x3dhHeader["opk_id"] = opkId
        }
        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 ikHex = x3dhHeader["ik"] as? String,
              let ikData = Data(hexString: ikHex),
              let ekHex = x3dhHeader["ek"] as? String,
              let ekData = Data(hexString: ekHex),
              let spkIdStr = x3dhHeader["spk_id"] as? String else {
            throw CryptoError.x3dhFailed("Invalid X3DH header")
        }

        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 {
            spkToUse = override
        } else 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")
        }

        // OPK
        var opkPriv: Curve25519.KeyAgreement.PrivateKey?
        if let opkIdStr = x3dhHeader["opk_id"] as? String {
            opkPriv = opkPrivates[opkIdStr] ?? KeyStorage.loadOPKPrivate(email: email, opkId: opkIdStr)
            if opkPriv != nil {
                opkPrivates.removeValue(forKey: opkIdStr)
                KeyStorage.deleteOPKPrivate(email: email, opkId: opkIdStr)
            }
        }

        let sharedSecret = try X3DH.respond(
            ikPrivateEd: identityPrivate!,
            spkPrivate: spkToUse,
            ikRemoteEd: remoteIkEd,
            ekRemote: ekRemote,
            opkPrivate: opkPriv
        )

        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)

        return ratchet
    }

    // MARK: - Send Message

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

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

    // MARK: - Send DM

    private func sendDM(convId: String, text: String, members: [ConversationMember], replyTo: String? = nil) async -> (success: Bool, message: String) {
        guard let identityPrivate = identityPrivate else {
            return (false, "Identity key not loaded")
        }

        let plaintext = Data(text.utf8)
        var payload: [String: Any] = ["text": text]
        if let replyTo = replyTo {
            payload["reply_to"] = replyTo
        }

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

        // Encrypt for each member's devices
        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
                    )

                    // 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,
                        "ciphertext": ProtocolHandler.encodeBinary(encrypted.ciphertext),
                        "nonce": ProtocolHandler.encodeBinary(encrypted.nonce),
                        "ratchet_header": encrypted.header,
                    ]
                    if let x3dh = x3dhHeader {
                        recipientEntry["x3dh_header"] = x3dh
                    }
                    recipients.append(recipientEntry)
                }
            } catch {
                return (false, "Encryption failed for \(member.username): \(error.localizedDescription)")
            }
        }

        // Self-encrypted copy
        let selfKey = CryptoUtils.deriveSelfEncryptionKey(identityPrivateRaw: identityPrivate.rawData)
        if let (_, nonce, ct, tag) = try? CryptoUtils.aesEncrypt(plaintext, key: selfKey) {
            let selfCiphertext = ct + tag
            let dummyHeader: [String: Any] = [
                "dh_pub": String(repeating: "00", count: 32),
                "n": 0,
                "pn": 0,
            ]
            recipients.append([
                "user_id": userId!,
                "device_id": Constants.selfDeviceId,
                "ciphertext": ProtocolHandler.encodeBinary(selfCiphertext),
                "nonce": ProtocolHandler.encodeBinary(nonce),
                "ratchet_header": dummyHeader,
            ])
        }

        // 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
        }

        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)
        }

        // Cache the sent message
        if let msgData = resp.dict(for: "data"), let messageId = msgData.string(for: "message_id") {
            var cacheEntry = payload
            cacheEntry["sender_id"] = userId
            cacheEntry["sender_username"] = username
            cacheEntry["created_at"] = ISO8601DateFormatter().string(from: Date())
            try? MessageCache.save(email: email, convId: convId, messages: [cacheEntry], cacheKey: cacheKey)
        }

        return (true, "Message sent")
    }

    // MARK: - Send Group Message

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

        // 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: members)
        }

        // Encrypt with sender key
        let plaintext = Data(text.utf8)
        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 members where member.userId != userId {
                recipients.append([
                    "user_id": member.userId,
                    "device_id": Constants.selfDeviceId,  // group messages use sentinel
                    "ciphertext": ProtocolHandler.encodeBinary(encrypted.ciphertext),
                    "nonce": ProtocolHandler.encodeBinary(encrypted.nonce),
                ])
            }

            // Self copy
            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,
                    "ciphertext": ProtocolHandler.encodeBinary(ct + tag),
                    "nonce": ProtocolHandler.encodeBinary(nonce),
                ])
            }

            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,
                "sender_chain_id": ProtocolHandler.encodeBinary(encrypted.ciphertext.prefix(0)),  // placeholder
            ]

            // Include sender key metadata for group routing
            params["sender_chain_id"] = encrypted.chainIdHex
            params["sender_chain_n"] = encrypted.n

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

            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)
            }

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

    // 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,
                        "ciphertext": 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 {
                print("Failed to distribute sender key to \(member.userId): \(error)")
            }
        }
    }

    // MARK: - Decrypt

    func decryptDMRecipientData(
        senderData: [String: Any],
        senderId: String,
        senderDeviceId: String
    ) -> Data? {
        guard let ctB64 = senderData["ciphertext"] as? String,
              let nonceB64 = senderData["nonce"] as? String,
              let ct = try? ProtocolHandler.decodeBinary(ctB64),
              let nonce = try? ProtocolHandler.decodeBinary(nonceB64) else {
            return nil
        }

        // Self-encrypted copy
        if senderDeviceId == Constants.selfDeviceId || senderId == userId {
            if let cacheKey = cacheKey {
                // ct = ciphertext + tag(16)
                guard ct.count >= 16 else { return nil }
                let ciphertext = ct.prefix(ct.count - 16)
                let tag = ct.suffix(16)
                return try? CryptoUtils.aesDecrypt(key: cacheKey, nonce: nonce, ciphertext: Data(ciphertext), tag: Data(tag))
            }
            return nil
        }

        // Regular DM decryption
        let headerDict = senderData["ratchet_header"] as? [String: Any]
        let x3dhHeader = senderData["x3dh_header"] as? [String: Any]

        let sessionKey = "\(senderId):\(senderDeviceId)"
        var ratchet = sessions[sessionKey]
            ?? KeyStorage.loadSession(email: email, peerUserId: senderId, localKey: localKey, peerDeviceId: senderDeviceId)

        // Handle X3DH header (new session)
        if let x3dh = x3dhHeader {
            do {
                ratchet = try processX3DHHeader(
                    senderId: senderId,
                    x3dhHeader: x3dh,
                    senderDeviceId: senderDeviceId
                )
            } catch {
                // Try with previous SPK (grace period)
                if let prevSpk = prevSpkPrivate {
                    ratchet = try? processX3DHHeader(
                        senderId: senderId,
                        x3dhHeader: x3dh,
                        senderDeviceId: senderDeviceId,
                        spkOverride: prevSpk
                    )
                }
                if ratchet == nil { return nil }
            }
        }

        guard let ratchet = ratchet, let header = headerDict else { return nil }

        do {
            let plaintext = try ratchet.decrypt(headerDict: header, ciphertext: ct, nonce: nonce)
            sessions[sessionKey] = ratchet
            try? KeyStorage.saveSession(email: email, peerUserId: senderId, ratchet: ratchet, localKey: localKey, peerDeviceId: senderDeviceId)

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

            return plaintext
        } catch {
            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 {
            print("Failed to import sender key: \(error)")
        }
    }

    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 {
            return nil
        }

        let senderDeviceId = data.string(for: "sender_device_id") ?? Constants.selfDeviceId

        // Find our device's entry
        var recipientData: [String: Any]?
        if let deviceEntries = data["device_entries"] as? [[String: Any]] {
            recipientData = deviceEntries.first(where: {
                ($0["device_id"] as? String) == deviceId || ($0["device_id"] as? String) == Constants.selfDeviceId
            })
        }
        // Fallback: use data directly if it has ciphertext
        if recipientData == nil, data["ciphertext"] != nil {
            recipientData = data
        }

        guard let recipientData = recipientData else { return nil }

        // Try DM decryption
        if let plaintext = decryptDMRecipientData(
            senderData: recipientData,
            senderId: senderId,
            senderDeviceId: senderDeviceId
        ) {
            let text = String(data: plaintext, encoding: .utf8)

            // Parse JSON payload
            var messageText = text
            var replyTo: String?
            var file: FileInfo?
            if let jsonObj = try? JSONSerialization.jsonObject(with: plaintext) as? [String: Any] {
                messageText = jsonObj["text"] as? String
                replyTo = jsonObj["reply_to"] as? String
                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 ?? ""
                    )
                }
            }

            let createdAt = data.string(for: "created_at").flatMap { ISO8601DateFormatter().date(from: $0) } ?? Date()
            let senderUsername = data.string(for: "sender_username") ?? userCache[senderId]?.username ?? "Unknown"

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

        return nil
    }

    // 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: "id") else { return nil }

            let membersRaw = dict["members"] as? [[String: Any]] ?? []
            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 { return nil }
                return ConversationMember(userId: uid, username: uname, email: uemail)
            }

            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] = ["emails": 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: - Messages

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

        var messages: [Message] = []
        for msgDict in messagesRaw {
            guard let msgId = msgDict.string(for: "id"),
                  let senderId = msgDict.string(for: "sender_id") else { continue }

            let senderDeviceId = msgDict.string(for: "sender_device_id") ?? Constants.selfDeviceId
            let isDeleted = msgDict["deleted_at"] != nil && !(msgDict["deleted_at"] is NSNull)

            if isDeleted {
                let createdAt = msgDict.string(for: "created_at").flatMap { ISO8601DateFormatter().date(from: $0) } ?? Date()
                messages.append(Message(
                    id: msgId, conversationId: convId, senderId: senderId,
                    senderUsername: msgDict.string(for: "sender_username") ?? "",
                    createdAt: createdAt, text: nil, isDeleted: true, readBy: []
                ))
                continue
            }

            // Try to decrypt
            if let plaintext = decryptDMRecipientData(
                senderData: msgDict,
                senderId: senderId,
                senderDeviceId: senderDeviceId
            ) {
                let text = String(data: plaintext, encoding: .utf8)
                var messageText = text
                var replyTo: String?
                var file: FileInfo?

                if let jsonObj = try? JSONSerialization.jsonObject(with: plaintext) as? [String: Any] {
                    messageText = jsonObj["text"] as? String
                    replyTo = jsonObj["reply_to"] as? String
                    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 ?? ""
                        )
                    }
                }

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

                let createdAt = msgDict.string(for: "created_at").flatMap { ISO8601DateFormatter().date(from: $0) } ?? Date()
                messages.append(Message(
                    id: msgId, conversationId: convId, senderId: senderId,
                    senderUsername: msgDict.string(for: "sender_username") ?? "",
                    createdAt: createdAt, text: messageText, replyTo: replyTo,
                    imageFileId: msgDict.string(for: "image_file_id"), file: file,
                    isDeleted: false, readBy: []
                ))
            }
        }

        return messages
    }

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

    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"
    }

    // 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 {
        let resp = await sendAndReceive(type: "update_avatar", params: [
            "avatar_data": ProtocolHandler.encodeBinary(imageData),
        ])
        return resp.string(for: "status") == "ok"
    }

    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: "avatar_data"),
              let avatarData = try? ProtocolHandler.decodeBinary(avatarB64) else {
            return nil
        }
        return avatarData
    }

    // MARK: - Group Avatar

    func updateGroupAvatar(convId: String, imageData: Data) async -> Bool {
        let resp = await sendAndReceive(type: "update_group_avatar", params: [
            "conversation_id": convId,
            "avatar_data": ProtocolHandler.encodeBinary(imageData),
        ])
        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: "avatar_data"),
              let avatarData = try? ProtocolHandler.decodeBinary(avatarB64) else {
            return nil
        }
        return avatarData
    }

    // MARK: - File Sharing

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

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

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

        // Upload chunks
        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", params: [
                "file_id": fileId,
                "chunk_data": ProtocolHandler.encodeBinary(Data(chunk)),
                "offset": offset,
            ])
            guard chunkResp.string(for: "status") == "ok" else {
                return (false, "Chunk upload failed")
            }
            offset = end
        }

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

        // Send message with file info
        let fileInfo: [String: Any] = [
            "file_id": fileId,
            "aes_key": aesKey.hexString,
            "iv": nonce.hexString,
            "filename": filename,
            "size": fileData.count,
            "mime_type": mimeType,
        ]

        var payload: [String: Any] = ["file": fileInfo]
        if let replyTo = replyTo {
            payload["reply_to"] = replyTo
        }

        let payloadData = try! JSONSerialization.data(withJSONObject: payload)
        let payloadText = String(data: payloadData, encoding: .utf8)!

        return await sendMessage(convId: convId, text: payloadText, members: members, replyTo: replyTo)
    }

    func downloadFile(fileId: String, aesKey: Data, iv: Data) async -> Data? {
        var allData = Data()
        var offset = 0

        while true {
            let resp = await sendAndReceive(type: "download_image", params: [
                "file_id": fileId,
                "offset": offset,
            ])
            guard resp.string(for: "status") == "ok",
                  let data = resp.dict(for: "data"),
                  let chunkB64 = data.string(for: "chunk_data"),
                  let chunk = try? ProtocolHandler.decodeBinary(chunkB64) else {
                break
            }

            if chunk.isEmpty { break }
            allData.append(chunk)
            offset += chunk.count

            if data.bool(for: "complete") == true { break }
        }

        guard !allData.isEmpty else { return nil }

        // Decrypt: allData = ciphertext + tag(16)
        guard allData.count >= 16 else { return nil }
        let ct = allData.prefix(allData.count - 16)
        let tag = allData.suffix(16)
        return try? CryptoUtils.aesDecrypt(key: aesKey, nonce: iv, ciphertext: Data(ct), tag: Data(tag))
    }

    // 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)
    }
}
