//
// Copyright 2020-2024 New Vector Ltd.
// Copyright (c) 2021 BWI GmbH
//
// SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-Element-Commercial
// Please see LICENSE files in the repository root for full details.
//

import UIKit
import KeychainAccess
import CommonCrypto
import MatrixSDK

@objcMembers
class EncryptionKeyManager: NSObject, MXKeyProviderDelegate {
    static let shared = EncryptionKeyManager()

    private static let keychainService: String = BuildSettings.baseBundleIdentifier + ".encryption-manager-service"
    private static let contactsIv = "iv"
    private static let contactsAesKey = "key"
    private static let accountIv = "acountIv"
    private static let accountAesKey = "acountKey"
    private static let cryptoOlmPickleKey = "cryptoOlmPickleKey"
    private static let roomLastMessageIv = "roomLastMessageIv"
    private static let roomLastMessageAesKey = "roomLastMessageAesKey"
    private static let cryptoSDKStoreKey = "cryptoSDKStoreKey"

    private var initialized = false
    
    private let keychainVault = KeychainVault(Keychain(service: keychainService, accessGroup: BuildSettings.keychainAccessGroup))
    private let securedFileStore: SecureFileStorage = SecureFileStorage.shared
    private let vault: KeyValueVault

    private override init() {
        vault = BWIBuildSettings.shared.forcedPinProtection ? securedFileStore : keychainVault
    }

    func initKeys() {
        generateIvIfNotExists(forKey: EncryptionKeyManager.accountIv, inStore: keychainVault)
        generateAesKeyIfNotExists(forKey: EncryptionKeyManager.accountAesKey, inStore: keychainVault)
        generateKeyIfNotExists(forKey: EncryptionKeyManager.cryptoOlmPickleKey, inStore: keychainVault, size: 32)

        assert(keychainVault.objectExists(withKey: EncryptionKeyManager.accountIv), "[EncryptionKeyManager] initKeys: Failed to generate IV for acount")
        assert(keychainVault.objectExists(withKey: EncryptionKeyManager.accountAesKey), "[EncryptionKeyManager] initKeys: Failed to generate AES Key for acount")
        assert(keychainVault.objectExists(withKey: EncryptionKeyManager.cryptoOlmPickleKey), "[EncryptionKeyManager] initKeys: Failed to generate Key for olm pickle key")

        generateIvIfNotExists(forKey: EncryptionKeyManager.roomLastMessageIv, inStore: keychainVault)
        generateAesKeyIfNotExists(forKey: EncryptionKeyManager.roomLastMessageAesKey, inStore: keychainVault)
        generateKeyIfNotExists(forKey: EncryptionKeyManager.cryptoSDKStoreKey, inStore: keychainVault, size: 32)

        assert(keychainVault.objectExists(withKey: EncryptionKeyManager.roomLastMessageIv), "[EncryptionKeyManager] initKeys: Failed to generate IV for room last message")
        assert(keychainVault.objectExists(withKey: EncryptionKeyManager.roomLastMessageAesKey), "[EncryptionKeyManager] initKeys: Failed to generate AES Key for room last message encryption")
        assert(keychainVault.objectExists(withKey: EncryptionKeyManager.cryptoSDKStoreKey), "[EncryptionKeyManager] initKeys: Failed to generate Key for crypto sdk store")

        guard !BWIBuildSettings.shared.forcedPinProtection || !SecureFileStorage.shared.locked else {
            MXLog.debug("[EncryptionKeyManager] initKeys: cannot init keys as store is not ready")
            return
        }

        generateIvIfNotExists(forKey: EncryptionKeyManager.contactsIv, inStore: vault)
        generateAesKeyIfNotExists(forKey: EncryptionKeyManager.contactsAesKey, inStore: vault)

        self.initialized = true
        
        assert(vault.objectExists(withKey: EncryptionKeyManager.contactsIv), "[EncryptionKeyManager] initKeys: Failed to generate IV for contacts")
        assert(vault.objectExists(withKey: EncryptionKeyManager.contactsAesKey), "[EncryptionKeyManager] initKeys: Failed to generate AES Key for contacts")
    }

    // MARK: - MXKeyProviderDelegate

    func isEncryptionAvailableForData(ofType dataType: String) -> Bool {
        return dataType == MXKContactManagerDataType
            || dataType == MXKAccountManagerDataType
            || dataType == MXCryptoOlmPickleKeyDataType
            || dataType == MXRoomLastMessageDataType
            || dataType == MXCryptoSDKStoreKeyDataType
    }

    func hasKeyForData(ofType dataType: String) -> Bool {
        switch dataType {
        case MXKContactManagerDataType:
            return vault.objectExists(withKey: EncryptionKeyManager.contactsIv) && vault.objectExists(withKey: EncryptionKeyManager.contactsAesKey)
        case MXKAccountManagerDataType:
            return keychainVault.objectExists(withKey: EncryptionKeyManager.accountIv) && keychainVault.objectExists(withKey: EncryptionKeyManager.accountAesKey)
        case MXCryptoOlmPickleKeyDataType:
            return keychainVault.objectExists(withKey: EncryptionKeyManager.cryptoOlmPickleKey)
        case MXRoomLastMessageDataType:
            return keychainVault.objectExists(withKey: EncryptionKeyManager.roomLastMessageIv) &&
            keychainVault.objectExists(withKey: EncryptionKeyManager.roomLastMessageAesKey)
        case MXCryptoSDKStoreKeyDataType:
            return keychainVault.objectExists(withKey: EncryptionKeyManager.cryptoSDKStoreKey)
        default:
            MXLog.warning("[EncryptionKeyManager] hasKeyForData: No key for \(dataType)")
            return false
        }
    }

    func keyDataForData(ofType dataType: String) -> MXKeyData? {
        switch dataType {
        case MXKContactManagerDataType:
            if let ivKey = try? vault.data(forKey: EncryptionKeyManager.contactsIv),
               let aesKey = try? vault.data(forKey: EncryptionKeyManager.contactsAesKey) {
                return MXAesKeyData(iv: ivKey, key: aesKey)
            }
        case MXKAccountManagerDataType:
            if let ivKey = try? keychainVault.data(forKey: EncryptionKeyManager.accountIv),
               let aesKey = try? keychainVault.data(forKey: EncryptionKeyManager.accountAesKey) {
                return MXAesKeyData(iv: ivKey, key: aesKey)
            }
        case MXCryptoOlmPickleKeyDataType:
            if let key = try? keychainVault.data(forKey: EncryptionKeyManager.cryptoOlmPickleKey) {
                return MXRawDataKey(key: key)
            }
        case MXRoomLastMessageDataType:
            if let ivKey = try? keychainVault.data(forKey: EncryptionKeyManager.roomLastMessageIv),
               let aesKey = try? keychainVault.data(forKey: EncryptionKeyManager.roomLastMessageAesKey) {
                return MXAesKeyData(iv: ivKey, key: aesKey)
            }
        case MXCryptoSDKStoreKeyDataType:
            if let key = try? keychainVault.data(forKey: EncryptionKeyManager.cryptoSDKStoreKey) {
                return MXRawDataKey(key: key)
            }
        default:
            MXLog.failure("[EncryptionKeyManager] keyDataForData: Attempting to get data for unknown type", dataType)
            return nil
        }
        return nil
    }
    
    func isInitFinished() -> Bool {
        return self.initialized
    }

    // MARK: - Private methods

    private func generateIvIfNotExists(forKey key: String, inStore store: KeyValueVault) {
        guard !store.objectExists(withKey: key) else {
            return
        }

        do {
            try store.set(MXAes.iv(), forKey: key)
        } catch {
            MXLog.debug("[EncryptionKeyManager] initKeys: Failed to generate IV: \(error.localizedDescription)")
        }
    }

    private func generateAesKeyIfNotExists(forKey key: String, inStore store: KeyValueVault) {
        generateKeyIfNotExists(forKey: key, inStore: store, size: kCCKeySizeAES256)
    }

    private func generateKeyIfNotExists(forKey key: String, inStore store: KeyValueVault, size: Int) {
        guard !store.objectExists(withKey: key) else {
            return
        }

        do {
            var aesKeyBytes = [UInt8](repeating: 0, count: size)
              _ = SecRandomCopyBytes(kSecRandomDefault, size, &aesKeyBytes)
            try store.set(Data(bytes: aesKeyBytes, count: size), forKey: key)
        } catch {
            MXLog.debug("[EncryptionKeyManager] initKeys: Failed to generate Key[\(key)]: \(error.localizedDescription)")
        }
    }
}