import Foundation
import CryptoKit
import CommonCrypto
import P256K

// MARK: - Error Types

enum CryptoError: LocalizedError {
    case invalidEntropy
    case invalidChecksum
    case invalidMnemonic
    case invalidPrivateKey
    case encryptionFailed
    case decryptionFailed
    case keyDerivationFailed
    case compressionFailed
    case invalidData(String)

    var errorDescription: String? {
        switch self {
        case .invalidMnemonic:
            return "Invalid recovery phrase. Please check each word for typos."
        case .decryptionFailed:
            return "Decryption failed. The password may be incorrect."
        case .invalidData(let reason):
            return "Invalid data: \(reason)"
        default:
            return "A cryptographic operation failed."
        }
    }
}

// MARK: - CryptoManager

/// All methods are `nonisolated` — safe to call from any actor/thread.
/// Low-level primitives are in `CryptoPrimitives`.
final class CryptoManager: @unchecked Sendable {

    nonisolated static let shared = CryptoManager()

    // MARK: - Android Parity: PBKDF2 Key Cache

    /// Caches derived PBKDF2 keys to avoid repeated ~50-100ms derivations.
    /// Android: `CryptoManager.pbkdf2KeyCache` (ConcurrentHashMap).
    /// Key format: "algorithm::password", value: derived 32-byte key.
    nonisolated private let pbkdf2CacheLock = NSLock()
    nonisolated(unsafe) private var pbkdf2Cache: [String: Data] = [:]

    // MARK: - Android Parity: Decryption Cache

    /// Caches decrypted results to avoid repeated AES + PBKDF2 for same input.
    /// Android: `CryptoManager.decryptionCache` (ConcurrentHashMap, max 2000).
    /// PERF: Also tracks total memory size — evicts when >50MB to prevent OOM.
    nonisolated private static let decryptionCacheMaxSize = 2000
    nonisolated private static let decryptionCacheMaxBytes = 50 * 1024 * 1024 // 50MB
    nonisolated private let decryptionCacheLock = NSLock()
    nonisolated(unsafe) private var decryptionCache: [String: Data] = [:]
    nonisolated(unsafe) private var decryptionCacheTotalBytes: Int = 0

    private init() {}

    /// Clear all crypto caches on logout (Android parity: `clearCaches()`).
    nonisolated func clearCaches() {
        pbkdf2CacheLock.lock()
        pbkdf2Cache.removeAll()
        pbkdf2CacheLock.unlock()
        decryptionCacheLock.lock()
        decryptionCache.removeAll()
        decryptionCacheTotalBytes = 0
        decryptionCacheLock.unlock()
    }

    // MARK: - BIP39: Mnemonic Generation

    nonisolated func generateMnemonic() throws -> [String] {
        let entropy = try CryptoPrimitives.randomBytes(count: 16)
        return try mnemonicFromEntropy(entropy)
    }

    // MARK: - BIP39: Mnemonic Validation

    nonisolated func validateMnemonic(_ words: [String]) -> Bool {
        guard words.count == 12 else { return false }
        let normalized = words.map { $0.lowercased().trimmingCharacters(in: .whitespaces) }
        guard normalized.allSatisfy({ BIP39.wordSet.contains($0) }) else { return false }
        return (try? entropyFromMnemonic(normalized)) != nil
    }

    // MARK: - BIP39: Mnemonic → Seed (PBKDF2-SHA512)

    nonisolated func mnemonicToSeed(_ words: [String]) -> Data {
        let phrase = words.joined(separator: " ")
        return CryptoPrimitives.pbkdf2(
            password: phrase, salt: "mnemonic", iterations: 2048,
            keyLength: 64, prf: CCPseudoRandomAlgorithm(kCCPRFHmacAlgSHA512)
        )
    }

    // MARK: - Key Pair Derivation (secp256k1)

    nonisolated func deriveKeyPair(from mnemonic: [String]) throws -> (privateKey: Data, publicKey: Data) {
        let seed = mnemonicToSeed(mnemonic)
        let seedHex = seed.hexString
        let privateKey = Data(SHA256.hash(data: Data(seedHex.utf8)))
        let publicKey = try deriveCompressedPublicKey(from: privateKey)
        return (privateKey, publicKey)
    }

    // MARK: - Account Encryption (PBKDF2-SHA256 + rawDeflate + AES-256-CBC)

    /// iOS-only encryption for account keys, device identity, persistence snapshots.
    /// Uses PBKDF2-HMAC-SHA256 + raw deflate (no zlib wrapper).
    nonisolated func encryptWithPassword(_ data: Data, password: String) throws -> String {
        let compressed = try CryptoPrimitives.rawDeflate(data)
        let key = cachedPBKDF2(password: password)
        let iv = try CryptoPrimitives.randomBytes(count: 16)
        let ciphertext = try CryptoPrimitives.aesCBCEncrypt(compressed, key: key, iv: iv)
        return "\(iv.base64EncodedString()):\(ciphertext.base64EncodedString())"
    }

    // MARK: - Desktop-Compatible Encryption (PBKDF2-SHA256 + zlibDeflate + AES-256-CBC)

    /// Desktop parity: CryptoJS v4 PBKDF2 defaults to HMAC-SHA256, pako.deflate() is zlib-wrapped.
    /// Use ONLY for cross-platform data: aesChachaKey, avatar blobs sent to server.
    nonisolated func encryptWithPasswordDesktopCompat(_ data: Data, password: String) throws -> String {
        let compressed = try CryptoPrimitives.zlibDeflate(data)
        let key = cachedPBKDF2(password: password)
        let iv = try CryptoPrimitives.randomBytes(count: 16)
        let ciphertext = try CryptoPrimitives.aesCBCEncrypt(compressed, key: key, iv: iv)
        return "\(iv.base64EncodedString()):\(ciphertext.base64EncodedString())"
    }

    /// - Parameter requireCompression: When `true`, skips the uncompressed fallback.
    ///   Use for attachment blobs which are ALWAYS compressed (zlibDeflate/rawDeflate).
    ///   The uncompressed fallback accepts ANY AES-CBC output as valid, which returns
    ///   garbage when the password is wrong — breaking multi-candidate password loops.
    nonisolated func decryptWithPassword(
        _ encrypted: String,
        password: String,
        requireCompression: Bool = false
    ) throws -> Data {
        // Android parity: decryption cache (max 2000 entries)
        let cacheKey = "\(password):\(encrypted)"
        decryptionCacheLock.lock()
        if let cached = decryptionCache[cacheKey] {
            decryptionCacheLock.unlock()
            return cached
        }
        decryptionCacheLock.unlock()

        let parts = encrypted.components(separatedBy: ":")
        guard parts.count == 2,
              let iv = Data(base64Encoded: parts[0]),
              let ciphertext = Data(base64Encoded: parts[1]) else {
            throw CryptoError.invalidData("Malformed encrypted string")
        }

        let prfOrder: [CCPseudoRandomAlgorithm] = [
            CCPseudoRandomAlgorithm(kCCPRFHmacAlgSHA256),
            CCPseudoRandomAlgorithm(kCCPRFHmacAlgSHA1),
        ]

        // 1) Preferred path: AES-CBC + inflate (handles both rawDeflate and zlibDeflate)
        for prf in prfOrder {
            do {
                let result = try decryptWithPassword(
                    ciphertext: ciphertext, iv: iv, password: password,
                    prf: prf, expectsCompressed: true
                )
                cacheDecryptionResult(cacheKey, result)
                return result
            } catch { }
        }

        // 2) Fallback: AES-CBC without compression (very old/legacy payloads).
        //    Skipped when requireCompression is true — prevents wrong-password garbage
        //    from being accepted as valid data.
        if !requireCompression {
            for prf in prfOrder {
                do {
                    let result = try decryptWithPassword(
                        ciphertext: ciphertext, iv: iv, password: password,
                        prf: prf, expectsCompressed: false
                    )
                    cacheDecryptionResult(cacheKey, result)
                    return result
                } catch { }
            }
        }

        throw CryptoError.decryptionFailed
    }

    /// Store result in decryption cache with eviction (Android: 10% eviction at 2000 entries).
    private nonisolated func cacheDecryptionResult(_ key: String, _ value: Data) {
        decryptionCacheLock.lock()
        // PERF: Evict when count exceeds max OR total memory exceeds 50MB.
        if decryptionCache.count >= Self.decryptionCacheMaxSize
            || decryptionCacheTotalBytes > Self.decryptionCacheMaxBytes {
            let evictCount = max(Self.decryptionCacheMaxSize / 10, 50)
            let keysToRemove = Array(decryptionCache.keys.prefix(evictCount))
            for k in keysToRemove {
                if let removed = decryptionCache.removeValue(forKey: k) {
                    decryptionCacheTotalBytes -= removed.count
                }
            }
        }
        decryptionCache[key] = value
        decryptionCacheTotalBytes += value.count
        decryptionCacheLock.unlock()
    }

    // MARK: - Utilities

    nonisolated func sha256(_ data: Data) -> Data {
        Data(SHA256.hash(data: data))
    }

    nonisolated func generatePrivateKeyHash(privateKeyHex: String) -> String {
        let combined = Data((privateKeyHex + "rosetta").utf8)
        return sha256(combined).hexString
    }

    // MARK: - Android Parity: Cached PBKDF2

    /// Cache PBKDF2 key derivation. Android: `getPbkdf2Key()` in CryptoManager.kt.
    /// Key: "prfAlgorithm::password", Value: 32-byte derived key.
    nonisolated func cachedPBKDF2(password: String, prf: CCPseudoRandomAlgorithm = CCPseudoRandomAlgorithm(kCCPRFHmacAlgSHA256)) -> Data {
        let cacheKey = "\(prf)::\(password)"
        pbkdf2CacheLock.lock()
        if let cached = pbkdf2Cache[cacheKey] {
            pbkdf2CacheLock.unlock()
            return cached
        }
        pbkdf2CacheLock.unlock()

        let key = CryptoPrimitives.pbkdf2(
            password: password, salt: "rosetta", iterations: 1000,
            keyLength: 32, prf: prf
        )

        pbkdf2CacheLock.lock()
        pbkdf2Cache[cacheKey] = key
        pbkdf2CacheLock.unlock()
        return key
    }
}

private extension CryptoManager {
    nonisolated func decryptWithPassword(
        ciphertext: Data,
        iv: Data,
        password: String,
        prf: CCPseudoRandomAlgorithm,
        expectsCompressed: Bool
    ) throws -> Data {
        // Android parity: cache PBKDF2 derived key
        let key = cachedPBKDF2(password: password, prf: prf)
        let decrypted = try CryptoPrimitives.aesCBCDecrypt(ciphertext, key: key, iv: iv)
        if expectsCompressed {
            return try CryptoPrimitives.rawInflate(decrypted)
        }
        return decrypted
    }
}

// MARK: - BIP39 Internal

private extension CryptoManager {

    nonisolated func mnemonicFromEntropy(_ entropy: Data) throws -> [String] {
        guard entropy.count == 16 else { throw CryptoError.invalidEntropy }

        let hashBytes = Data(SHA256.hash(data: entropy))
        let checksumByte = hashBytes[0]

        var bits = [Bool]()
        bits.reserveCapacity(132)
        for byte in entropy {
            for shift in stride(from: 7, through: 0, by: -1) {
                bits.append((byte >> shift) & 1 == 1)
            }
        }
        for shift in stride(from: 7, through: 4, by: -1) {
            bits.append((checksumByte >> shift) & 1 == 1)
        }

        return try (0..<12).map { chunk in
            let index = (0..<11).reduce(0) { acc, bit in
                acc * 2 + (bits[chunk * 11 + bit] ? 1 : 0)
            }
            guard index < BIP39.wordList.count else { throw CryptoError.invalidEntropy }
            return BIP39.wordList[index]
        }
    }

    nonisolated func entropyFromMnemonic(_ words: [String]) throws -> Data {
        guard words.count == 12 else { throw CryptoError.invalidMnemonic }

        var bits = [Bool]()
        bits.reserveCapacity(132)
        for word in words {
            guard let index = BIP39.index(of: word) else { throw CryptoError.invalidMnemonic }
            for shift in stride(from: 10, through: 0, by: -1) {
                bits.append((index >> shift) & 1 == 1)
            }
        }

        var entropy = Data(count: 16)
        for byteIdx in 0..<16 {
            let value: UInt8 = (0..<8).reduce(0) { acc, bit in
                acc * 2 + (bits[byteIdx * 8 + bit] ? 1 : 0)
            }
            entropy[byteIdx] = value
        }

        let hashBytes = Data(SHA256.hash(data: entropy))
        let expectedTopNibble = hashBytes[0] >> 4
        let actualTopNibble: UInt8 = (0..<4).reduce(0) { acc, bit in
            bits[128 + bit] ? acc | (1 << (3 - bit)) : acc
        }
        guard actualTopNibble == expectedTopNibble else { throw CryptoError.invalidChecksum }

        return entropy
    }
}

// MARK: - secp256k1

extension CryptoManager {

    nonisolated func deriveCompressedPublicKey(from privateKey: Data) throws -> Data {
        guard privateKey.count == 32 else { throw CryptoError.invalidPrivateKey }
        let signingKey = try P256K.Signing.PrivateKey(dataRepresentation: privateKey, format: .compressed)
        return signingKey.publicKey.dataRepresentation
    }
}
