import Foundation
import CryptoKit
import BeamCore
import KeychainAccess

enum EncryptionManagerError: Error {
    case authenticationFailure
    case stringEncodingError
    case stringDecodingError
    case keyError
    case internalEncryptionError
}
extension EncryptionManagerError: LocalizedError {
    public var errorDescription: String? {
        switch self {
        case .authenticationFailure:
            return loc("Couldn't decrypt data, encrypted with a different private key?")
        case .stringEncodingError:
            return loc("Couldn't change String to Data.")
        case .stringDecodingError:
            return loc("Couldn't decode UTF-8")
        case .keyError:
            return loc("Key couldn't be read")
        case .internalEncryptionError:
            return loc("Internal encryption error")
        }
    }
}

// https://medium.com/swlh/common-cryptographic-operations-in-swift-with-cryptokit-b30a4becc895
class EncryptionManager {
    static let shared = EncryptionManager()
    private let lock = RWLock()

    enum Algorithm: String, CaseIterable {
        case ChaChaPoly
        case AES_GCM
    }

    func generateKey() -> SymmetricKey {
        SymmetricKey(size: .bits256)
    }

    // MARK: - Private Keys
    func clearPrivateKey(for email: String) {
        lock.write {
            guard var privateKeys = Persistence.Encryption.privateKeys else { return }

            privateKeys.removeValue(forKey: email)
            Persistence.Encryption.privateKeys = privateKeys
            Persistence.Encryption.updateDate = BeamDate.now

            Logger.shared.logWarning("Private key for user \(email) has been cleared", category: .encryption)
        }
    }

    func resetPrivateKeys(andMigrateOldSharedKey: Bool) {
        lock.write {
            Logger.shared.logWarning("Private keys has been reset", category: .encryption)
            if andMigrateOldSharedKey {
                Persistence.Encryption.privateKeys = [:]
                if let email = Persistence.Authentication.email {
                    _ = self.privateKey(for: email)
                } else {
                    Persistence.Encryption.updateDate = BeamDate.now
                }
            } else {
                Persistence.Encryption.privateKeys = nil
                Persistence.Encryption.updateDate = BeamDate.now
            }
        }
    }

    func replacePrivateKey(for email: String, with base64EncodedString: String) throws {
        try lock.write {
            guard let key = SymmetricKey(base64EncodedString: base64EncodedString) else {
                throw EncryptionManagerError.keyError
            }

            guard var privateKeys = Persistence.Encryption.privateKeys else { return }

            privateKeys.updateValue(key.asString(), forKey: email)

            Persistence.Encryption.privateKeys = privateKeys
            Persistence.Encryption.updateDate = BeamDate.now

            Logger.shared.logWarning("Private key for user \(email) has been replaced", category: .encryption)
        }
    }

    func readPrivateKey(for email: String) -> SymmetricKey? {
        guard !email.isEmpty else {
            fatalError("No Email provided to read PrivateKey, it should never happen")
        }

        return lock.read {
            if let privateKeys = Persistence.Encryption.privateKeys,
               let privateKeyData = privateKeys[email],
               let result = SymmetricKey(base64EncodedString: privateKeyData) {
                return result
            }
            return nil
        }
    }

    @discardableResult
    func privateKey(for email: String) -> SymmetricKey {
        guard !email.isEmpty else {
            fatalError("No Email provided to get PrivateKey, it should never happen")
        }

        return lock.write {
            if let privateKeys = Persistence.Encryption.privateKeys,
               let privateKeyData = privateKeys[email],
               let result = SymmetricKey(base64EncodedString: privateKeyData) {
                return result
            }
            Logger.shared.logWarning("Private key for user \(email) doesn't exist, creating new one",
                                     category: .encryption)
            let key = self.generateKey()
            self.savePrivateKey(for: email, with: key.asString())
            return key
        }
    }

    func existingPrivateKey(for email: String) -> SymmetricKey? {
        guard !email.isEmpty else {
            fatalError("No Email provided to get PrivateKey, it should never happen")
        }

        return lock.write {
            if let privateKeys = Persistence.Encryption.privateKeys,
               let privateKeyData = privateKeys[email],
               let result = SymmetricKey(base64EncodedString: privateKeyData) {
                return result
            }

            return nil
        }
    }

    var accounts: [String] {
        return lock.read {
            guard let pkeys = Persistence.Encryption.privateKeys else { return [] }
            return Array(pkeys.keys)
        }
    }

    private func savePrivateKey(for email: String, with key: String) {
        lock.write {
            if var privateKeys = Persistence.Encryption.privateKeys {
                privateKeys.updateValue(key, forKey: email)
                Persistence.Encryption.privateKeys = privateKeys
            } else {
                Persistence.Encryption.privateKeys = [email: key]
                Persistence.Encryption.creationDate = BeamDate.now

            }
            Persistence.Encryption.updateDate = BeamDate.now
        }
    }

    // MARK: - Local Private Key
    private var cachedLocalPrivateKey: String? // avoid hitting the keychain for each request

    func loadLocalPrivateKey() throws -> String? {
        if cachedLocalPrivateKey != nil {
            return cachedLocalPrivateKey
        }
        let keychain = Keychain(service: Configuration.bundleIdentifier).synchronizable(false)
        cachedLocalPrivateKey = try keychain.getString(Persistence.Encryption.localPrivateKeyName)
        return cachedLocalPrivateKey
    }

    func storeLocalPrivateKey(_ privateKey: String) throws {
        let keychain = Keychain(service: Configuration.bundleIdentifier).synchronizable(false)
        try keychain.set(privateKey, key: Persistence.Encryption.localPrivateKeyName)
        cachedLocalPrivateKey = privateKey
    }

    @discardableResult
    func localPrivateKey(allowCreating: Bool = false) throws -> SymmetricKey {
        try lock.write {
            if let localDataKey = try self.loadLocalPrivateKey(),
               let localPrivateKey = SymmetricKey(base64EncodedString: localDataKey) {
                return localPrivateKey
            }
            guard allowCreating else {
                Logger.shared.logError("Local Private key doesn't exist or has a wrong format, password manager will be disabled", category: .encryption)
                throw EncryptionManagerError.keyError
            }
            Logger.shared.logInfo("Local Private key doesn't exist or has a wrong format, creating new one", category: .encryption)
            let key = self.generateKey()
            try self.storeLocalPrivateKey(key.asString())
            return key
        }
    }

    // MARK: - Encrypt / Decrypt
    func encryptString(_ string: String, _ key: SymmetricKey? = nil, using: Algorithm = .AES_GCM) throws -> String? {
        let key = key ?? privateKey(for: Persistence.emailOrRaiseError())
        switch using {
        case .ChaChaPoly:
            return try ChaChaPoly
                .seal(string.asData, using: key)
                .combined
                .base64EncodedString()
        case .AES_GCM:
            return try AES.GCM
                .seal(string.asData, using: key)
                .combined?
                .base64EncodedString()
        }
    }

    func decryptString(_ encryptedString: String, _ key: SymmetricKey? = nil, using: Algorithm = .AES_GCM) throws -> String? {
        guard let encryptedData = Data(base64Encoded: encryptedString) else {
            Logger.shared.logError("Couldn't change String to Data", category: .encryption)
            throw EncryptionManagerError.stringEncodingError
        }
        let key = key ?? privateKey(for: Persistence.emailOrRaiseError())

        do {
            switch using {
            case .ChaChaPoly:
                let sealedBox = try ChaChaPoly.SealedBox(combined: encryptedData)
                let decryptedData = try ChaChaPoly.open(sealedBox, using: key)
                return decryptedData.asString
            case .AES_GCM:
                let sealedBox = try AES.GCM.SealedBox(combined: encryptedData)
                let decryptedData = try AES.GCM.open(sealedBox, using: key)
                return decryptedData.asString
            }
        } catch CryptoKit.CryptoKitError.authenticationFailure {
            throw EncryptionManagerError.authenticationFailure
        }
    }

    func encryptData(_ data: Data, _ key: SymmetricKey? = nil, using: Algorithm = .AES_GCM) throws -> Data? {
        let key = key ?? privateKey(for: Persistence.emailOrRaiseError())

        switch using {
        case .ChaChaPoly:
            return try ChaChaPoly.seal(data, using: key).combined
        case .AES_GCM:
            return try AES.GCM
                .seal(data, using: key)
                .combined
        }
    }

    func decryptData(_ encryptedData: Data, _ key: SymmetricKey? = nil, using: Algorithm = .AES_GCM) throws -> Data? {
        let key = key ?? privateKey(for: Persistence.emailOrRaiseError())

        switch using {
        case .ChaChaPoly:
            let sealedBox = try ChaChaPoly.SealedBox(combined: encryptedData)
            let decryptedData = try ChaChaPoly.open(sealedBox, using: key)
            return decryptedData
        case .AES_GCM:
            let sealedBox = try AES.GCM.SealedBox(combined: encryptedData)
            let decryptedData = try AES.GCM.open(sealedBox, using: key)
            return decryptedData
        }
    }
}
