import Foundation
import CryptoKit

/// Local file storage for keys, sessions, and sender keys.
/// Matches Python: chat_core.py key storage functions.
///
/// Base directory: Application Support / EncryptedChat / {email}
/// Same file names as Python client for cross-platform compatibility.
enum KeyStorage {

    // MARK: - Base Directory

    /// Get or create the key storage directory for a user
    static func getKeyDir(email: String) throws -> URL {
        let appSupport = FileManager.default.urls(for: .applicationSupportDirectory, in: .userDomainMask).first!
        let dir = appSupport.appendingPathComponent("EncryptedChat").appendingPathComponent(email)
        try FileManager.default.createDirectory(at: dir, withIntermediateDirectories: true)
        // iOS file protection
        try (dir as NSURL).setResourceValue(URLFileProtection.complete, forKey: .fileProtectionKey)
        return dir
    }

    // MARK: - RSA Keys

    /// Save RSA keypair
    static func saveRSAKeys(email: String, privateKey: SecKey, publicKey: SecKey, password: Data? = nil) throws {
        let dir = try getKeyDir(email: email)
        let privData = try RSACrypto.serializePrivateKey(privateKey, password: password)
        let pubData = try RSACrypto.serializePublicKey(publicKey)
        try writeProtected(privData, to: dir.appendingPathComponent("private.pem"))
        try pubData.write(to: dir.appendingPathComponent("public.pem"))
    }

    /// Load RSA keypair. Returns (private, public, error).
    static func loadRSAKeys(email: String, password: Data? = nil) -> (SecKey?, SecKey?, String?) {
        guard let dir = try? getKeyDir(email: email) else {
            return (nil, nil, "Cannot access key directory")
        }
        let privPath = dir.appendingPathComponent("private.pem")
        let pubPath = dir.appendingPathComponent("public.pem")

        guard FileManager.default.fileExists(atPath: privPath.path) else {
            return (nil, nil, "No local keys found.")
        }

        guard let privData = try? Data(contentsOf: privPath),
              let pubData = try? Data(contentsOf: pubPath) else {
            return (nil, nil, "Cannot read key files.")
        }

        do {
            let privateKey = try RSACrypto.loadPrivateKey(privData, password: password)
            let publicKey = try RSACrypto.loadPublicKey(pubData)
            return (privateKey, publicKey, nil)
        } catch {
            // Try without password (unencrypted)
            do {
                let privateKey = try RSACrypto.loadPrivateKey(privData, password: nil)
                let publicKey = try RSACrypto.loadPublicKey(pubData)
                // Re-save with password if provided
                if let password = password {
                    try? saveRSAKeys(email: email, privateKey: privateKey, publicKey: publicKey, password: password)
                }
                return (privateKey, publicKey, nil)
            } catch {
                return (nil, nil, "Invalid or missing password.")
            }
        }
    }

    // MARK: - Identity Keys (Ed25519)

    static func saveIdentityKeys(
        email: String,
        privateKey: Curve25519.Signing.PrivateKey,
        publicKey: Curve25519.Signing.PublicKey,
        password: Data? = nil
    ) throws {
        let dir = try getKeyDir(email: email)
        let privData = try Ed25519Crypto.serializePrivate(privateKey, password: password)
        let pubData = Ed25519Crypto.serializePublic(publicKey)
        try writeProtected(privData, to: dir.appendingPathComponent("identity_private.bin"))
        try pubData.write(to: dir.appendingPathComponent("identity_public.bin"))
    }

    static func loadIdentityKeys(
        email: String,
        password: Data? = nil
    ) -> (Curve25519.Signing.PrivateKey?, Curve25519.Signing.PublicKey?) {
        guard let dir = try? getKeyDir(email: email) else { return (nil, nil) }
        let privPath = dir.appendingPathComponent("identity_private.bin")
        let pubPath = dir.appendingPathComponent("identity_public.bin")

        guard FileManager.default.fileExists(atPath: privPath.path),
              let privData = try? Data(contentsOf: privPath),
              let pubData = try? Data(contentsOf: pubPath) else {
            return (nil, nil)
        }

        do {
            let priv = try Ed25519Crypto.loadPrivate(privData, password: password)
            let pub = try Ed25519Crypto.loadPublic(pubData)
            return (priv, pub)
        } catch {
            return (nil, nil)
        }
    }

    // MARK: - Signed Pre-Key

    static func saveSPK(email: String, privateKey: Curve25519.KeyAgreement.PrivateKey, spkId: String) throws {
        let dir = try getKeyDir(email: email)
        try writeProtected(X25519Crypto.serializePrivate(privateKey), to: dir.appendingPathComponent("spk_private.bin"))
        try spkId.write(to: dir.appendingPathComponent("spk_id.txt"), atomically: true, encoding: .utf8)
    }

    static func loadSPK(email: String) -> (Curve25519.KeyAgreement.PrivateKey?, String?) {
        guard let dir = try? getKeyDir(email: email) else { return (nil, nil) }
        let privPath = dir.appendingPathComponent("spk_private.bin")
        let idPath = dir.appendingPathComponent("spk_id.txt")

        guard FileManager.default.fileExists(atPath: privPath.path),
              let privData = try? Data(contentsOf: privPath),
              let priv = try? X25519Crypto.loadPrivate(privData) else {
            return (nil, nil)
        }
        let spkId = (try? String(contentsOf: idPath, encoding: .utf8))?.trimmed ?? ""
        return (priv, spkId)
    }

    // MARK: - Previous SPK (Grace Period)

    static func savePrevSPK(email: String, privateKey: Curve25519.KeyAgreement.PrivateKey, spkId: String) throws {
        let dir = try getKeyDir(email: email)
        try writeProtected(X25519Crypto.serializePrivate(privateKey), to: dir.appendingPathComponent("prev_spk_private.bin"))
        try spkId.write(to: dir.appendingPathComponent("prev_spk_id.txt"), atomically: true, encoding: .utf8)
    }

    static func loadPrevSPK(email: String) -> (Curve25519.KeyAgreement.PrivateKey?, String?) {
        guard let dir = try? getKeyDir(email: email) else { return (nil, nil) }
        let privPath = dir.appendingPathComponent("prev_spk_private.bin")
        let idPath = dir.appendingPathComponent("prev_spk_id.txt")

        guard FileManager.default.fileExists(atPath: privPath.path),
              let privData = try? Data(contentsOf: privPath),
              let priv = try? X25519Crypto.loadPrivate(privData) else {
            return (nil, nil)
        }
        let spkId = (try? String(contentsOf: idPath, encoding: .utf8))?.trimmed ?? ""
        return (priv, spkId)
    }

    // MARK: - One-Time Pre-Keys

    static func saveOPKPrivate(email: String, opkId: String, privateKey: Curve25519.KeyAgreement.PrivateKey) throws {
        let dir = try getKeyDir(email: email).appendingPathComponent("opk_private")
        try FileManager.default.createDirectory(at: dir, withIntermediateDirectories: true)
        try writeProtected(X25519Crypto.serializePrivate(privateKey), to: dir.appendingPathComponent("\(opkId).bin"))
    }

    static func loadOPKPrivate(email: String, opkId: String) -> Curve25519.KeyAgreement.PrivateKey? {
        guard let dir = try? getKeyDir(email: email) else { return nil }
        let path = dir.appendingPathComponent("opk_private").appendingPathComponent("\(opkId).bin")
        guard let data = try? Data(contentsOf: path) else { return nil }
        return try? X25519Crypto.loadPrivate(data)
    }

    static func deleteOPKPrivate(email: String, opkId: String) {
        guard let dir = try? getKeyDir(email: email) else { return }
        let path = dir.appendingPathComponent("opk_private").appendingPathComponent("\(opkId).bin")
        try? FileManager.default.removeItem(at: path)
    }

    // MARK: - Device ID

    static func saveDeviceId(email: String, deviceId: String) throws {
        let dir = try getKeyDir(email: email)
        try writeProtected(Data(deviceId.utf8), to: dir.appendingPathComponent("device_id.txt"))
    }

    static func loadDeviceId(email: String) -> String? {
        guard let dir = try? getKeyDir(email: email) else { return nil }
        let path = dir.appendingPathComponent("device_id.txt")
        guard let data = try? Data(contentsOf: path) else { return nil }
        let str = String(data: data, encoding: .utf8)?.trimmed
        return (str?.isEmpty ?? true) ? nil : str
    }

    // MARK: - Sessions (Double Ratchet)

    static func saveSession(
        email: String,
        peerUserId: String,
        ratchet: DoubleRatchet,
        localKey: Data? = nil,
        peerDeviceId: String? = nil
    ) throws {
        let dir = try getKeyDir(email: email).appendingPathComponent("sessions")
        try FileManager.default.createDirectory(at: dir, withIntermediateDirectories: true)

        let filename: String
        if let deviceId = peerDeviceId {
            filename = "\(peerUserId)_\(deviceId).bin"
        } else {
            filename = "\(peerUserId).bin"
        }

        var data = try ratchet.exportState()
        if let localKey = localKey {
            data = try CryptoUtils.encryptLocal(data, key: localKey)
        }
        try writeProtected(data, to: dir.appendingPathComponent(filename))
    }

    static func loadSession(
        email: String,
        peerUserId: String,
        localKey: Data? = nil,
        peerDeviceId: String? = nil
    ) -> DoubleRatchet? {
        guard let dir = try? getKeyDir(email: email) else { return nil }
        let sessionsDir = dir.appendingPathComponent("sessions")

        let filename: String
        if let deviceId = peerDeviceId {
            filename = "\(peerUserId)_\(deviceId).bin"
        } else {
            filename = "\(peerUserId).bin"
        }

        let path = sessionsDir.appendingPathComponent(filename)
        return loadSessionFile(path, localKey: localKey)
    }

    static func deleteSession(email: String, peerUserId: String, peerDeviceId: String? = nil) {
        guard let dir = try? getKeyDir(email: email) else { return }
        let sessionsDir = dir.appendingPathComponent("sessions")

        if let deviceId = peerDeviceId {
            let path = sessionsDir.appendingPathComponent("\(peerUserId)_\(deviceId).bin")
            try? FileManager.default.removeItem(at: path)
        } else {
            // Delete all sessions for this user
            if let files = try? FileManager.default.contentsOfDirectory(atPath: sessionsDir.path) {
                for file in files where file.hasPrefix(peerUserId) {
                    try? FileManager.default.removeItem(at: sessionsDir.appendingPathComponent(file))
                }
            }
        }
    }

    private static func loadSessionFile(_ path: URL, localKey: Data?) -> DoubleRatchet? {
        guard let raw = try? Data(contentsOf: path) else { return nil }

        if let localKey = localKey {
            // Try encrypted first
            if let decrypted = try? CryptoUtils.decryptLocal(raw, key: localKey) {
                return try? DoubleRatchet.importState(decrypted)
            }
            // Fallback: plaintext (transparent migration)
            if let ratchet = try? DoubleRatchet.importState(raw) {
                // Re-save encrypted
                try? writeProtected(CryptoUtils.encryptLocal(try ratchet.exportState(), key: localKey), to: path)
                return ratchet
            }
            return nil
        }

        return try? DoubleRatchet.importState(raw)
    }

    // MARK: - Sender Keys

    static func saveSenderKeyState(
        email: String,
        convId: String,
        state: SenderKeyState,
        localKey: Data? = nil
    ) throws {
        let dir = try getKeyDir(email: email).appendingPathComponent("sender_keys")
        try FileManager.default.createDirectory(at: dir, withIntermediateDirectories: true)

        var data = state.exportState()
        if let localKey = localKey {
            data = try CryptoUtils.encryptLocal(data, key: localKey)
        }
        try writeProtected(data, to: dir.appendingPathComponent("\(convId).bin"))
    }

    static func loadSenderKeyState(
        email: String,
        convId: String,
        localKey: Data? = nil
    ) -> SenderKeyState? {
        guard let dir = try? getKeyDir(email: email) else { return nil }
        let path = dir.appendingPathComponent("sender_keys").appendingPathComponent("\(convId).bin")
        guard let raw = try? Data(contentsOf: path) else { return nil }

        if let localKey = localKey {
            if let decrypted = try? CryptoUtils.decryptLocal(raw, key: localKey) {
                return try? SenderKeyState.importState(decrypted)
            }
            // Plaintext fallback
            if let state = try? SenderKeyState.importState(raw) {
                try? writeProtected(CryptoUtils.encryptLocal(state.exportState(), key: localKey), to: path)
                return state
            }
            return nil
        }

        return try? SenderKeyState.importState(raw)
    }

    static func deleteSenderKeyState(email: String, convId: String) {
        guard let dir = try? getKeyDir(email: email) else { return }
        let path = dir.appendingPathComponent("sender_keys").appendingPathComponent("\(convId).bin")
        try? FileManager.default.removeItem(at: path)
    }

    // MARK: - Received Sender Keys

    static func saveRecvSenderKey(
        email: String,
        convId: String,
        senderId: String,
        senderDeviceId: String,
        state: SenderKeyState,
        localKey: Data? = nil
    ) throws {
        let dir = try getKeyDir(email: email).appendingPathComponent("sender_keys_recv")
        try FileManager.default.createDirectory(at: dir, withIntermediateDirectories: true)

        var data = state.exportState()
        if let localKey = localKey {
            data = try CryptoUtils.encryptLocal(data, key: localKey)
        }
        try writeProtected(data, to: dir.appendingPathComponent("\(convId)_\(senderId)_\(senderDeviceId).bin"))
    }

    static func loadRecvSenderKey(
        email: String,
        convId: String,
        senderId: String,
        senderDeviceId: String,
        localKey: Data? = nil
    ) -> SenderKeyState? {
        guard let dir = try? getKeyDir(email: email) else { return nil }
        let path = dir.appendingPathComponent("sender_keys_recv").appendingPathComponent("\(convId)_\(senderId)_\(senderDeviceId).bin")
        guard let raw = try? Data(contentsOf: path) else { return nil }

        if let localKey = localKey {
            if let decrypted = try? CryptoUtils.decryptLocal(raw, key: localKey) {
                return try? SenderKeyState.importState(decrypted)
            }
            if let state = try? SenderKeyState.importState(raw) {
                try? writeProtected(CryptoUtils.encryptLocal(state.exportState(), key: localKey), to: path)
                return state
            }
            return nil
        }

        return try? SenderKeyState.importState(raw)
    }

    static func deleteRecvSenderKeys(email: String, convId: String) {
        guard let dir = try? getKeyDir(email: email) else { return }
        let recvDir = dir.appendingPathComponent("sender_keys_recv")
        guard let files = try? FileManager.default.contentsOfDirectory(atPath: recvDir.path) else { return }
        for file in files where file.hasPrefix(convId) {
            try? FileManager.default.removeItem(at: recvDir.appendingPathComponent(file))
        }
    }

    // MARK: - Favorites

    static func saveFavorites(email: String, favorites: Set<String>) throws {
        let dir = try getKeyDir(email: email)
        let data = try JSONSerialization.data(withJSONObject: Array(favorites))
        try data.write(to: dir.appendingPathComponent("favorites.json"))
    }

    static func loadFavorites(email: String) -> Set<String> {
        guard let dir = try? getKeyDir(email: email) else { return [] }
        let path = dir.appendingPathComponent("favorites.json")
        guard let data = try? Data(contentsOf: path),
              let array = try? JSONSerialization.jsonObject(with: data) as? [String] else {
            return []
        }
        return Set(array)
    }

    // MARK: - Helpers

    private static func writeProtected(_ data: Data, to url: URL) throws {
        try data.write(to: url, options: .completeFileProtection)
    }
}
