//
//  CryptoManager.swift
//  Montana Protocol — macOS
//
//  High-level orchestrator for Layer 1 cryptographic operations
//  ML-DSA-65 key generation from BIP-39 seed, signing, address derivation
//

import Foundation
import SwiftUI
import CommonCrypto
import Security
import CryptoKit

@MainActor
class CryptoManager: ObservableObject {
    static let shared = CryptoManager()

    @Published var hasIdentity: Bool = false
    @Published var address: String = ""
    @Published var publicKeyHex: String = ""

    private let keychain = KeychainManager.shared
    private static let addressCacheKey = "montana_cached_address"

    /// Version header for encrypted blobs: "MT" + version (self-describing format)
    /// Prevents downgrade attacks — format is identified from the blob itself, not external state
    private static let encHeader = Data([0x4d, 0x54, 0x01])   // MT1: PBKDF2
    private static let encHeaderV2 = Data([0x4d, 0x54, 0x02]) // MT2: Argon2id + optional SE

    init() {
        // Load from UserDefaults cache — NO Keychain access at startup!
        // This prevents the macOS Keychain password prompt on every launch.
        loadCachedIdentity()
        // PIN existence: UserDefaults is a UI hint ONLY, NOT authoritative
        // Security decisions use isPinActive (Keychain-based) via computed property
        hasPinCache = UserDefaults.standard.bool(forKey: "montana_has_pin")
    }

    // MARK: - Identity

    /// Load identity from UserDefaults cache (no Keychain access at startup)
    private func loadCachedIdentity() {
        if let cachedAddress = UserDefaults.standard.string(forKey: Self.addressCacheKey),
           !cachedAddress.isEmpty {
            hasIdentity = true
            address = cachedAddress
            syncAddressToEngine()
        } else {
            hasIdentity = false
        }
    }

    /// Cache address in UserDefaults
    private func cacheAddress(_ addr: String) {
        UserDefaults.standard.set(addr, forKey: Self.addressCacheKey)
    }

    /// Clear cached address
    private func clearAddressCache() {
        UserDefaults.standard.removeObject(forKey: Self.addressCacheKey)
    }

    /// Sync address to PresenceEngine + auto-start tracking
    private func syncAddressToEngine() {
        guard !address.isEmpty else { return }
        PresenceEngine.shared.address = address
        if !PresenceEngine.shared.isTracking {
            PresenceEngine.shared.autoStart()
        }
    }

    /// Generate new identity: BIP-39 mnemonic → PBKDF2 → ML-DSA-65 keypair
    /// - Returns: 24 mnemonic words (must be shown to user for backup!)
    func generateIdentity() -> [String]? {
        // Reset any previous balance — new address = zero balance
        PresenceEngine.shared.fullReset()

        // SECURITY: Clear ALL old crypto state (including PIN) before storing new plaintext keys
        // This prevents plaintext + active PIN coexistence
        keychain.deleteAll()
        hasPin = false

        // 1. Generate 24-word mnemonic
        guard let words = MontanaSeed.generateMnemonic() else {
            return nil
        }

        // 2. Derive keypair from mnemonic
        guard var keypair = MontanaSeed.keypairFromMnemonic(words) else {
            return nil
        }
        defer {
            MLDSA65.zeroMemory(&keypair.privateKey)
        }

        // 3. Save to Keychain (keys + mnemonic)
        let mnemonicData = words.joined(separator: " ").data(using: .utf8)!
        guard keychain.save(.privateKey, data: keypair.privateKey),
              keychain.save(.publicKey, data: keypair.publicKey),
              keychain.save(.mnemonic, data: mnemonicData) else {
            return nil
        }

        // 4. Update state + cache
        address = MLDSA65.generateAddress(from: keypair.publicKey)
        publicKeyHex = keypair.publicKey.map { String(format: "%02x", $0) }.joined()
        hasIdentity = true
        cacheAddress(address)
        syncAddressToEngine()

        return words
    }

    /// Recover identity from existing 24-word mnemonic
    /// - Parameter words: 24 BIP-39 mnemonic words
    /// - Returns: true if recovery succeeded
    func recoverIdentity(from words: [String]) -> Bool {
        // Reset any previous balance — recovered address syncs from server
        PresenceEngine.shared.fullReset()

        // SECURITY: Clear ALL old crypto state (including PIN) before storing new plaintext keys
        keychain.deleteAll()
        hasPin = false

        guard MontanaSeed.validateMnemonic(words) else {
            return false
        }

        guard var keypair = MontanaSeed.keypairFromMnemonic(words) else {
            return false
        }
        defer {
            MLDSA65.zeroMemory(&keypair.privateKey)
        }

        let mnemonicData = words.joined(separator: " ").data(using: .utf8)!
        guard keychain.save(.privateKey, data: keypair.privateKey),
              keychain.save(.publicKey, data: keypair.publicKey),
              keychain.save(.mnemonic, data: mnemonicData) else {
            return false
        }

        address = MLDSA65.generateAddress(from: keypair.publicKey)
        publicKeyHex = keypair.publicKey.map { String(format: "%02x", $0) }.joined()
        hasIdentity = true
        cacheAddress(address)
        syncAddressToEngine()

        return true
    }

    /// Delete identity from Keychain + clear cache + full balance reset
    func deleteIdentity() {
        keychain.deleteAll()
        clearAddressCache()
        hasIdentity = false
        hasPin = false
        address = ""
        publicKeyHex = ""
        UserDefaults.standard.removeObject(forKey: "montana_has_pin")
        // Clean up legacy UserDefaults lockout keys
        UserDefaults.standard.removeObject(forKey: "montana_pin_fails")
        UserDefaults.standard.removeObject(forKey: "montana_pin_lockout_end")
        UserDefaults.standard.removeObject(forKey: "montana_pin_lockout_level")
        PresenceEngine.shared.fullReset()
    }

    // MARK: - AES-256 Encryption / Key Derivation

    /// [LEGACY] Derive AES-256 key from PIN using PBKDF2 (600k iterations) — MT1 only
    private func deriveAESKeyPBKDF2(from pin: String, salt: Data) -> Data? {
        guard let pinData = pin.data(using: .utf8) else { return nil }
        var derivedKey = [UInt8](repeating: 0, count: 32)

        let status = pinData.withUnsafeBytes { pinPtr in
            salt.withUnsafeBytes { saltPtr in
                CCKeyDerivationPBKDF(
                    CCPBKDFAlgorithm(kCCPBKDF2),
                    pinPtr.baseAddress?.assumingMemoryBound(to: Int8.self),
                    pinData.count,
                    saltPtr.baseAddress?.assumingMemoryBound(to: UInt8.self),
                    salt.count,
                    CCPseudoRandomAlgorithm(kCCPRFHmacAlgSHA256),
                    600_000,
                    &derivedKey, 32
                )
            }
        }

        guard status == kCCSuccess else { return nil }
        return Data(derivedKey)
    }

    /// [NEW] Derive 32-byte key from PIN using Argon2id (64MB, 3 iter, 4 lanes)
    /// GPU/ASIC-resistant KDF — replaces PBKDF2 for MT2 blobs
    private func deriveAESKeyArgon2id(from pin: String, salt: Data) -> Data? {
        guard let pinData = pin.data(using: .utf8), !pinData.isEmpty else { return nil }
        guard salt.count == 16 else { return nil }

        var derivedKey = [UInt8](repeating: 0, count: 32)

        let result = pinData.withUnsafeBytes { pinPtr in
            salt.withUnsafeBytes { saltPtr in
                argon2id_hash_raw(
                    3,                    // t_cost: 3 iterations
                    65536,                // m_cost: 64 MiB (in KiB)
                    4,                    // parallelism: 4 threads
                    pinPtr.baseAddress,   // password
                    pinData.count,        // password length
                    saltPtr.baseAddress,  // salt
                    salt.count,           // salt length
                    &derivedKey,          // output hash
                    32                    // hash length
                )
            }
        }

        guard result == 0 else { return nil } // ARGON2_OK = 0
        return Data(derivedKey)
    }

    /// Derive PIN verification hash via Argon2id + HKDF (domain-separated from encryption key)
    /// SECURITY: This hash is stored in Keychain for verification.
    /// The encryption key uses a DIFFERENT HKDF domain ("montana-pin-encrypt-v2").
    /// Even if this hash leaks, encryption key cannot be derived from it.
    private func hashPinArgon2id(from pin: String, salt: Data) -> Data? {
        guard let raw = deriveAESKeyArgon2id(from: pin, salt: salt) else { return nil }
        let derived = HKDF<SHA256>.deriveKey(
            inputKeyMaterial: SymmetricKey(data: raw),
            salt: Data(),
            info: "montana-pin-verify-v2".data(using: .utf8)!,
            outputByteCount: 32
        )
        return derived.withUnsafeBytes { Data($0) }
    }

    // MARK: - Secure Enclave (Optional Hardware Protection)

    /// Check if Secure Enclave is available on this Mac (Apple Silicon / T2)
    private var isSecureEnclaveAvailable: Bool {
        SecureEnclave.isAvailable
    }

    /// Check if SE was used for the current wallet (Keychain flag)
    var isSecureEnclaveEnabled: Bool {
        guard let flag = keychain.load(.seEnabled) else { return false }
        return flag.count == 1 && flag[0] == 0x01
    }

    /// Create SE P-256 key, generate random 32-byte wrap key, encrypt with ECDH+AES-GCM.
    /// Returns cleartext wrap key (for immediate use) or nil if SE unavailable.
    private func seCreateAndWrapKey() -> Data? {
        guard isSecureEnclaveAvailable else { return nil }

        do {
            // 1. Create persistent SE key (P-256 KeyAgreement)
            let seKey = try SecureEnclave.P256.KeyAgreement.PrivateKey()

            // 2. Generate random 32-byte wrap key
            var wrapKeyBytes = [UInt8](repeating: 0, count: 32)
            guard SecRandomCopyBytes(kSecRandomDefault, 32, &wrapKeyBytes) == errSecSuccess else {
                return nil
            }
            let wrapKey = Data(wrapKeyBytes)

            // 3. Create companion key for ECDH
            let companion = P256.KeyAgreement.PrivateKey()
            let sharedSecret = try companion.sharedSecretFromKeyAgreement(
                with: seKey.publicKey
            )

            // 4. Derive symmetric key via HKDF
            let symmetricKey = sharedSecret.hkdfDerivedSymmetricKey(
                using: SHA256.self,
                salt: "montana-se-wrap-v1".data(using: .utf8)!,
                sharedInfo: Data(),
                outputByteCount: 32
            )

            // 5. AES-GCM seal wrap key
            let sealedBox = try AES.GCM.seal(wrapKey, using: symmetricKey)
            guard let sealedData = sealedBox.combined else { return nil }

            // 6. Build envelope: [1 byte pubLen] + companion public key + sealed box
            let companionPubData = companion.publicKey.x963Representation
            var envelope = Data()
            envelope.append(UInt8(companionPubData.count))
            envelope.append(companionPubData)
            envelope.append(sealedData)

            // 7. Store SE key reference + envelope in Keychain
            keychain.save(.seKeyTag, data: seKey.dataRepresentation)
            keychain.save(.seWrappedKey, data: envelope)
            keychain.save(.seEnabled, data: Data([0x01]))

            return wrapKey
        } catch {
            // SE creation failed (ad-hoc signing, etc.) — graceful fallback
            return nil
        }
    }

    /// Unwrap the 32-byte key using Secure Enclave
    private func seUnwrapKey() -> Data? {
        guard isSecureEnclaveEnabled else { return nil }
        guard let seKeyData = keychain.load(.seKeyTag),
              let envelopeData = keychain.load(.seWrappedKey) else { return nil }

        do {
            // 1. Parse envelope: [1 byte pubLen] + companion public + sealed box
            guard envelopeData.count > 1 else { return nil }
            let pubLen = Int(envelopeData[0])
            guard envelopeData.count >= 1 + pubLen + 12 + 32 + 16 else { return nil }

            let companionPubData = envelopeData.subdata(in: 1..<(1 + pubLen))
            let sealedData = envelopeData.subdata(in: (1 + pubLen)..<envelopeData.count)

            // 2. Reconstruct SE key from stored reference
            let seKey = try SecureEnclave.P256.KeyAgreement.PrivateKey(
                dataRepresentation: seKeyData
            )

            // 3. Reconstruct companion public key
            let companionPub = try P256.KeyAgreement.PublicKey(
                x963Representation: companionPubData
            )

            // 4. ECDH: SE private + companion public → shared secret
            let sharedSecret = try seKey.sharedSecretFromKeyAgreement(with: companionPub)

            // 5. HKDF → symmetric key (same parameters as wrap time)
            let symmetricKey = sharedSecret.hkdfDerivedSymmetricKey(
                using: SHA256.self,
                salt: "montana-se-wrap-v1".data(using: .utf8)!,
                sharedInfo: Data(),
                outputByteCount: 32
            )

            // 6. Unseal wrap key
            let sealedBox = try AES.GCM.SealedBox(combined: sealedData)
            let wrapKey = try AES.GCM.open(sealedBox, using: symmetricKey)

            return wrapKey
        } catch {
            return nil
        }
    }

    // MARK: - Combined Key Derivation v2 (Argon2id + optional SE)

    /// Derive final encryption key: Argon2id(PIN) [+ SE unwrap key via HKDF]
    /// - useSE: nil = auto-detect from Keychain, true/false = override (from blob flags on decrypt)
    /// SECURITY: If SE was used for encryption but unavailable now → FAIL (no silent downgrade)
    private func deriveFinalKeyV2(from pin: String, salt: Data, useSE: Bool? = nil) -> Data? {
        guard let raw = deriveAESKeyArgon2id(from: pin, salt: salt) else { return nil }

        // Domain separation: encryption key ≠ verification hash
        let encMaterial = HKDF<SHA256>.deriveKey(
            inputKeyMaterial: SymmetricKey(data: raw),
            salt: Data(),
            info: "montana-pin-encrypt-v2".data(using: .utf8)!,
            outputByteCount: 32
        )
        let pinKey = encMaterial.withUnsafeBytes { Data($0) }

        let seRequired = useSE ?? isSecureEnclaveEnabled

        if seRequired, let seKey = seUnwrapKey() {
            // Two-factor: PIN + hardware
            let derived = HKDF<SHA256>.deriveKey(
                inputKeyMaterial: SymmetricKey(data: pinKey),
                salt: seKey,
                info: "montana-enc-v2".data(using: .utf8)!,
                outputByteCount: 32
            )
            return derived.withUnsafeBytes { Data($0) }
        } else if seRequired {
            // SE was required but unavailable — FAIL (keys unrecoverable without SE)
            return nil
        } else {
            // PIN-only (no SE)
            let derived = HKDF<SHA256>.deriveKey(
                inputKeyMaterial: SymmetricKey(data: pinKey),
                salt: Data(),
                info: "montana-enc-v2-no-se".data(using: .utf8)!,
                outputByteCount: 32
            )
            return derived.withUnsafeBytes { Data($0) }
        }
    }

    /// MAC key derivation for MT2 (different domain from MT1)
    private func deriveMACKeyV2(from encKey: Data) -> Data {
        let info = "montana-hmac-key-v2".data(using: .utf8)!
        var mac = [UInt8](repeating: 0, count: 32)
        encKey.withUnsafeBytes { keyPtr in
            info.withUnsafeBytes { infoPtr in
                CCHmac(
                    CCHmacAlgorithm(kCCHmacAlgSHA256),
                    keyPtr.baseAddress, encKey.count,
                    infoPtr.baseAddress, info.count,
                    &mac
                )
            }
        }
        return Data(mac)
    }

    /// Derive MAC key from encryption key (domain separation for Encrypt-then-MAC) — MT1
    private func deriveMACKey(from encKey: Data) -> Data {
        let info = "montana-hmac-key".data(using: .utf8)!
        var mac = [UInt8](repeating: 0, count: 32)
        encKey.withUnsafeBytes { keyPtr in
            info.withUnsafeBytes { infoPtr in
                CCHmac(
                    CCHmacAlgorithm(kCCHmacAlgSHA256),
                    keyPtr.baseAddress, encKey.count,
                    infoPtr.baseAddress, info.count,
                    &mac
                )
            }
        }
        return Data(mac)
    }

    /// AES-256-CBC + HMAC-SHA256 encrypt (auto-selects MT2 or MT1 based on available salt)
    private func encryptAES256(_ plaintext: Data, pin: String) -> Data? {
        // Prefer MT2 (Argon2id) if argon2Salt exists
        if let argon2Salt = keychain.load(.argon2Salt) {
            return encryptMT2(plaintext, pin: pin, salt: argon2Salt)
        }
        // Fallback to MT1 (PBKDF2) for legacy
        guard let salt = keychain.load(.pinSalt),
              let key = deriveAESKeyPBKDF2(from: pin, salt: salt) else { return nil }
        return encryptWithKey(plaintext, key: key)
    }

    /// MT2 encrypt: header(3) + flags(1) + IV(16) + ciphertext + HMAC(32)
    private func encryptMT2(_ plaintext: Data, pin: String, salt: Data) -> Data? {
        guard let key = deriveFinalKeyV2(from: pin, salt: salt) else { return nil }
        let macKey = deriveMACKeyV2(from: key)

        var iv = [UInt8](repeating: 0, count: kCCBlockSizeAES128)
        guard SecRandomCopyBytes(kSecRandomDefault, iv.count, &iv) == errSecSuccess else { return nil }

        let bufferSize = plaintext.count + kCCBlockSizeAES128
        var buffer = [UInt8](repeating: 0, count: bufferSize)
        var numBytesEncrypted = 0

        let status = key.withUnsafeBytes { keyPtr in
            plaintext.withUnsafeBytes { dataPtr in
                CCCrypt(
                    CCOperation(kCCEncrypt),
                    CCAlgorithm(kCCAlgorithmAES),
                    CCOptions(kCCOptionPKCS7Padding),
                    keyPtr.baseAddress, kCCKeySizeAES256,
                    iv,
                    dataPtr.baseAddress, plaintext.count,
                    &buffer, bufferSize,
                    &numBytesEncrypted
                )
            }
        }

        guard status == kCCSuccess else { return nil }

        // Flags: bit 0 = SE used
        let flags: UInt8 = isSecureEnclaveEnabled ? 0x01 : 0x00
        let authenticated = Data([flags]) + Data(iv) + Data(buffer.prefix(numBytesEncrypted))

        // HMAC over (flags + IV + ciphertext)
        var hmac = [UInt8](repeating: 0, count: 32)
        macKey.withUnsafeBytes { macPtr in
            authenticated.withUnsafeBytes { dataPtr in
                CCHmac(
                    CCHmacAlgorithm(kCCHmacAlgSHA256),
                    macPtr.baseAddress, macKey.count,
                    dataPtr.baseAddress, authenticated.count,
                    &hmac
                )
            }
        }

        return Self.encHeaderV2 + authenticated + Data(hmac)
    }

    /// Self-describing decrypt: MT2 → Argon2id, MT1 → PBKDF2, no header → legacy CBC
    /// Format decision based on BLOB HEADER, not external state
    private func decryptAES256(_ encrypted: Data, pin: String) -> Data? {
        guard encrypted.count > 3 else { return nil }

        if encrypted.prefix(3) == Self.encHeaderV2 {
            // MT2: Argon2id + optional SE
            return decryptMT2(encrypted, pin: pin)
        }

        if encrypted.prefix(3) == Self.encHeader {
            // MT1: PBKDF2 + HMAC
            guard let salt = keychain.load(.pinSalt) else { return nil }
            guard let key = deriveAESKeyPBKDF2(from: pin, salt: salt) else { return nil }
            let body = encrypted.dropFirst(3)
            return decryptWithHMAC(Data(body), key: key)
        }

        // Legacy: no header → pre-v3.10 CBC-only
        guard keychain.load(.pinSalt) == nil && keychain.load(.argon2Salt) == nil else { return nil }
        let legacySalt = "montana_seed_aes256_v1".data(using: .utf8)!
        guard let key = deriveAESKeyPBKDF2(from: pin, salt: legacySalt) else { return nil }
        return decryptCBC(encrypted, key: key)
    }

    /// Decrypt MT2 blob: header(3) + flags(1) + IV(16) + ciphertext + HMAC(32)
    /// SECURITY: Uses blob flags (not Keychain state) to determine SE requirement
    private func decryptMT2(_ encrypted: Data, pin: String) -> Data? {
        // Minimum: header(3) + flags(1) + IV(16) + 1block(16) + HMAC(32) = 68
        guard encrypted.count >= 68, encrypted.prefix(3) == Self.encHeaderV2 else { return nil }

        // Extract SE flag from blob (byte 3 after header) — authoritative for THIS blob
        let body = encrypted.dropFirst(3)
        let blobUsedSE = (body[body.startIndex] & 0x01) != 0

        guard let salt = keychain.load(.argon2Salt),
              let key = deriveFinalKeyV2(from: pin, salt: salt, useSE: blobUsedSE) else { return nil }

        let macKey = deriveMACKeyV2(from: key)

        let dataForMAC = body.prefix(body.count - 32)
        let storedHMAC = body.suffix(32)

        // Compute + verify HMAC (constant-time)
        var computedHMAC = [UInt8](repeating: 0, count: 32)
        macKey.withUnsafeBytes { macPtr in
            dataForMAC.withUnsafeBytes { dataPtr in
                CCHmac(
                    CCHmacAlgorithm(kCCHmacAlgSHA256),
                    macPtr.baseAddress, macKey.count,
                    dataPtr.baseAddress, dataForMAC.count,
                    &computedHMAC
                )
            }
        }

        var diff: UInt8 = 0
        let hmacBytes = Array(storedHMAC)
        for i in 0..<32 { diff |= computedHMAC[i] ^ hmacBytes[i] }
        guard diff == 0 else { return nil }

        // Skip flags byte (1), then IV + ciphertext
        let ivAndCiphertext = Data(dataForMAC.dropFirst(1))
        return decryptCBC(ivAndCiphertext, key: key)
    }

    /// Encrypt-then-MAC: AES-256-CBC + HMAC-SHA256
    /// Format: IV (16) + ciphertext + HMAC-SHA256(IV+ciphertext) (32)
    private func encryptWithKey(_ plaintext: Data, key: Data) -> Data? {
        let macKey = deriveMACKey(from: key)

        var iv = [UInt8](repeating: 0, count: kCCBlockSizeAES128)
        guard SecRandomCopyBytes(kSecRandomDefault, iv.count, &iv) == errSecSuccess else { return nil }

        let bufferSize = plaintext.count + kCCBlockSizeAES128
        var buffer = [UInt8](repeating: 0, count: bufferSize)
        var numBytesEncrypted = 0

        let status = key.withUnsafeBytes { keyPtr in
            plaintext.withUnsafeBytes { dataPtr in
                CCCrypt(
                    CCOperation(kCCEncrypt),
                    CCAlgorithm(kCCAlgorithmAES),
                    CCOptions(kCCOptionPKCS7Padding),
                    keyPtr.baseAddress, kCCKeySizeAES256,
                    iv,
                    dataPtr.baseAddress, plaintext.count,
                    &buffer, bufferSize,
                    &numBytesEncrypted
                )
            }
        }

        guard status == kCCSuccess else { return nil }
        let ivPlusCiphertext = Data(iv) + Data(buffer.prefix(numBytesEncrypted))

        // Encrypt-then-MAC: HMAC over (IV + ciphertext)
        var hmac = [UInt8](repeating: 0, count: 32)
        macKey.withUnsafeBytes { macPtr in
            ivPlusCiphertext.withUnsafeBytes { dataPtr in
                CCHmac(
                    CCHmacAlgorithm(kCCHmacAlgSHA256),
                    macPtr.baseAddress, macKey.count,
                    dataPtr.baseAddress, ivPlusCiphertext.count,
                    &hmac
                )
            }
        }

        return Self.encHeader + ivPlusCiphertext + Data(hmac)
    }

    /// Strict HMAC-authenticated decrypt: returns nil if HMAC invalid (no fallback)
    /// Format: IV (16) + ciphertext + HMAC-SHA256(IV+ciphertext) (32)
    private func decryptWithHMAC(_ encrypted: Data, key: Data) -> Data? {
        guard encrypted.count > kCCBlockSizeAES128 + 32 else { return nil }

        let macKey = deriveMACKey(from: key)
        let dataForMAC = encrypted.prefix(encrypted.count - 32)
        let storedHMAC = encrypted.suffix(32)

        // Compute expected HMAC
        var computedHMAC = [UInt8](repeating: 0, count: 32)
        macKey.withUnsafeBytes { macPtr in
            dataForMAC.withUnsafeBytes { dataPtr in
                CCHmac(
                    CCHmacAlgorithm(kCCHmacAlgSHA256),
                    macPtr.baseAddress, macKey.count,
                    dataPtr.baseAddress, dataForMAC.count,
                    &computedHMAC
                )
            }
        }

        // Constant-time HMAC comparison — reject if tampered
        var diff: UInt8 = 0
        let hmacBytes = Array(storedHMAC)
        for i in 0..<32 { diff |= computedHMAC[i] ^ hmacBytes[i] }

        guard diff == 0 else { return nil }
        return decryptCBC(dataForMAC, key: key)
    }

    /// Raw AES-256-CBC decryption (internal, used by decryptWithKey)
    private func decryptCBC(_ encrypted: Data, key: Data) -> Data? {
        guard encrypted.count > kCCBlockSizeAES128 else { return nil }

        let iv = Array(encrypted.prefix(kCCBlockSizeAES128))
        let ciphertext = encrypted.suffix(from: kCCBlockSizeAES128)

        let bufferSize = ciphertext.count + kCCBlockSizeAES128
        var buffer = [UInt8](repeating: 0, count: bufferSize)
        var numBytesDecrypted = 0

        let status = key.withUnsafeBytes { keyPtr in
            ciphertext.withUnsafeBytes { dataPtr in
                CCCrypt(
                    CCOperation(kCCDecrypt),
                    CCAlgorithm(kCCAlgorithmAES),
                    CCOptions(kCCOptionPKCS7Padding),
                    keyPtr.baseAddress, kCCKeySizeAES256,
                    iv,
                    dataPtr.baseAddress, ciphertext.count,
                    &buffer, bufferSize,
                    &numBytesDecrypted
                )
            }
        }

        guard status == kCCSuccess else { return nil }
        return Data(buffer.prefix(numBytesDecrypted))
    }

    // MARK: - Mnemonic

    /// Load mnemonic decrypted with PIN (lockout enforced)
    func loadMnemonic(pin: String) -> [String]? {
        // SECURITY: Enforce lockout — prevents brute-force via mnemonic decryption
        guard verifyPin(pin) else { return nil }
        guard let encrypted = keychain.load(.mnemonicEncrypted),
              let decrypted = decryptAES256(encrypted, pin: pin),
              let phrase = String(data: decrypted, encoding: .utf8) else {
            return nil
        }
        let words = phrase.split(separator: " ").map(String.init)
        return words.count == 24 ? words : nil
    }

    /// Check if plaintext (unencrypted) keys exist — used for migration detection
    var hasPlaintextKeys: Bool {
        keychain.load(.privateKey) != nil || keychain.load(.mnemonic) != nil
    }

    // MARK: - PIN

    /// UI cache only — NOT used for security decisions
    @Published var hasPinCache: Bool = false

    /// Published wrapper: use hasPinCache for UI, isPinActive for security
    var hasPin: Bool {
        get { hasPinCache }
        set {
            hasPinCache = newValue
            UserDefaults.standard.set(newValue, forKey: "montana_has_pin")
        }
    }

    /// AUTHORITATIVE PIN check: Keychain-backed (tamper-resistant)
    /// All security decisions MUST use this, never UserDefaults
    var isPinActive: Bool {
        keychain.load(.pinHash) != nil &&
        (keychain.load(.argon2Salt) != nil || keychain.load(.pinSalt) != nil)
    }

    /// Check if wallet has MT1 (PBKDF2) blobs that need migration to MT2 (Argon2id)
    var needsMT2Migration: Bool {
        keychain.load(.pinHash) != nil &&
        keychain.load(.pinSalt) != nil &&
        keychain.load(.argon2Salt) == nil
    }

    // Lockout: 3 wrong → 1min, next → 10min, next → 1hr, next → 24hr
    private static let lockoutDurations: [TimeInterval] = [60, 600, 3600, 86400]

    /// Generate random 16-byte salt and save to Keychain
    private func generateNewSalt() -> Data? {
        var salt = [UInt8](repeating: 0, count: 16)
        guard SecRandomCopyBytes(kSecRandomDefault, salt.count, &salt) == errSecSuccess else { return nil }
        let saltData = Data(salt)
        keychain.save(.pinSalt, data: saltData)
        return saltData
    }

    /// Hash PIN with PBKDF2-HMAC-SHA256 (600k iterations, with random salt)
    private func hashPin(_ pin: String, salt: Data) -> Data? {
        guard let pinData = pin.data(using: .utf8) else { return nil }
        var hash = [UInt8](repeating: 0, count: 32)
        let status = pinData.withUnsafeBytes { pinPtr in
            salt.withUnsafeBytes { saltPtr in
                CCKeyDerivationPBKDF(
                    CCPBKDFAlgorithm(kCCPBKDF2),
                    pinPtr.baseAddress?.assumingMemoryBound(to: Int8.self),
                    pinData.count,
                    saltPtr.baseAddress?.assumingMemoryBound(to: UInt8.self),
                    salt.count,
                    CCPseudoRandomAlgorithm(kCCPRFHmacAlgSHA256),
                    600_000,
                    &hash, 32
                )
            }
        }
        guard status == kCCSuccess else { return nil }
        return Data(hash)
    }

    /// Legacy PIN hash: SHA-256 (no salt) — backward compatibility only
    private func hashPinLegacy(_ pin: String) -> Data? {
        guard let data = pin.data(using: .utf8) else { return nil }
        var hash = [UInt8](repeating: 0, count: 32)
        data.withUnsafeBytes { ptr in
            _ = CC_SHA256(ptr.baseAddress, CC_LONG(data.count), &hash)
        }
        return Data(hash)
    }

    // MARK: - Lockout (Keychain-backed, tamper-resistant)

    private func saveLockoutInt(_ key: KeychainManager.Key, _ value: Int) {
        var v = value
        let data = Data(bytes: &v, count: MemoryLayout<Int>.size)
        keychain.save(key, data: data)
    }

    private func loadLockoutInt(_ key: KeychainManager.Key) -> Int {
        guard let data = keychain.load(key), data.count == MemoryLayout<Int>.size else { return 0 }
        return data.withUnsafeBytes { $0.load(as: Int.self) }
    }

    private func saveLockoutDouble(_ key: KeychainManager.Key, _ value: Double) {
        var v = value
        let data = Data(bytes: &v, count: MemoryLayout<Double>.size)
        keychain.save(key, data: data)
    }

    private func loadLockoutDouble(_ key: KeychainManager.Key) -> Double {
        guard let data = keychain.load(key), data.count == MemoryLayout<Double>.size else { return 0.0 }
        return data.withUnsafeBytes { $0.load(as: Double.self) }
    }

    var failedAttempts: Int {
        loadLockoutInt(.lockoutFails)
    }

    var lockoutEnd: Date? {
        let ts = loadLockoutDouble(.lockoutEnd)
        guard ts > 0 else { return nil }
        let date = Date(timeIntervalSince1970: ts)
        return date > Date() ? date : nil
    }

    var isLockedOut: Bool { lockoutEnd != nil }

    var lockoutSecondsLeft: Int {
        guard let end = lockoutEnd else { return 0 }
        return max(Int(end.timeIntervalSinceNow), 0)
    }

    private func registerFailedAttempt() {
        let fails = loadLockoutInt(.lockoutFails) + 1
        saveLockoutInt(.lockoutFails, fails)

        if fails >= 3 {
            let level = loadLockoutInt(.lockoutLevel)
            let idx = min(level, Self.lockoutDurations.count - 1)
            let duration = Self.lockoutDurations[idx]
            let end = Date().addingTimeInterval(duration)
            saveLockoutDouble(.lockoutEnd, end.timeIntervalSince1970)
            saveLockoutInt(.lockoutLevel, level + 1)
            saveLockoutInt(.lockoutFails, 0)
        }
    }

    /// Reset consecutive failure counter only (on successful PIN entry)
    /// Does NOT reset escalation level — prevents lockout bypass via interleaved success
    private func resetConsecutiveFails() {
        keychain.delete(.lockoutFails)
        keychain.delete(.lockoutEnd)
    }

    /// Full lockout reset including escalation (only on PIN creation/change/delete)
    private func resetLockout() {
        keychain.delete(.lockoutFails)
        keychain.delete(.lockoutEnd)
        keychain.delete(.lockoutLevel)
    }

    /// Create new PIN (8 digits) + encrypt with Argon2id + optional SE
    /// ATOMIC: hasPin only set true after ALL encrypted saves confirmed
    func createPin(_ pin: String) -> Bool {
        guard pin.count == 8, pin.allSatisfy({ $0.isNumber }) else { return false }

        // Phase 1: Generate Argon2id salt (16 bytes)
        var saltBytes = [UInt8](repeating: 0, count: 16)
        guard SecRandomCopyBytes(kSecRandomDefault, 16, &saltBytes) == errSecSuccess else {
            return false
        }
        let argon2Salt = Data(saltBytes)
        guard keychain.save(.argon2Salt, data: argon2Salt) else { return false }

        // Clean stale SE flags before setup (prevents stale state from previous wallet)
        keychain.delete(.seEnabled); keychain.delete(.seKeyTag); keychain.delete(.seWrappedKey)
        // Try Secure Enclave setup (best-effort, nil = SE unavailable)
        let _ = seCreateAndWrapKey()

        // Hash PIN for verification (domain-separated from encryption key)
        guard let hash = hashPinArgon2id(from: pin, salt: argon2Salt) else {
            keychain.delete(.argon2Salt)
            return false
        }

        // Phase 2: Encrypt all plaintext material with MT2
        var encryptedPrivKey: Data?
        var encryptedMnemonic: Data?

        if let privKey = keychain.load(.privateKey) {
            guard let encrypted = encryptMT2(privKey, pin: pin, salt: argon2Salt) else {
                keychain.delete(.argon2Salt)
                keychain.delete(.seEnabled); keychain.delete(.seKeyTag); keychain.delete(.seWrappedKey)
                return false
            }
            encryptedPrivKey = encrypted
        }

        if let mnemonic = keychain.load(.mnemonic) {
            guard let encrypted = encryptMT2(mnemonic, pin: pin, salt: argon2Salt) else {
                keychain.delete(.argon2Salt)
                keychain.delete(.seEnabled); keychain.delete(.seKeyTag); keychain.delete(.seWrappedKey)
                return false
            }
            encryptedMnemonic = encrypted
        }

        // Phase 3: Save pin hash
        guard keychain.save(.pinHash, data: hash) else {
            keychain.delete(.argon2Salt)
            keychain.delete(.seEnabled); keychain.delete(.seKeyTag); keychain.delete(.seWrappedKey)
            return false
        }

        // Phase 4: Save ALL encrypted items (all-or-nothing)
        var allSaved = true
        if let enc = encryptedPrivKey {
            allSaved = allSaved && keychain.save(.privateKeyEncrypted, data: enc)
        }
        if let enc = encryptedMnemonic {
            allSaved = allSaved && keychain.save(.mnemonicEncrypted, data: enc)
        }

        guard allSaved else {
            // Rollback: delete everything, keep plaintext intact
            keychain.delete(.pinHash)
            keychain.delete(.argon2Salt)
            keychain.delete(.privateKeyEncrypted)
            keychain.delete(.mnemonicEncrypted)
            keychain.delete(.seEnabled); keychain.delete(.seKeyTag); keychain.delete(.seWrappedKey)
            return false
        }

        // Phase 5: Delete plaintext + old PBKDF2 salt (encrypted copies confirmed)
        keychain.delete(.privateKey)
        keychain.delete(.mnemonic)
        keychain.delete(.pinSalt) // Remove old PBKDF2 salt if migrating

        // Phase 6: Mark PIN as active
        hasPin = true
        resetLockout()
        return true
    }

    /// Change PIN: verify old → decrypt → re-encrypt with Argon2id MT2 → commit atomically
    func changePin(old: String, new: String) -> Bool {
        guard verifyPin(old) else { return false }
        guard new.count == 8, new.allSatisfy({ $0.isNumber }) else { return false }

        // Phase 1: Decrypt all material with old PIN (auto-detects MT2/MT1/legacy)
        var decryptedPrivKey: Data?
        if let encrypted = keychain.load(.privateKeyEncrypted) {
            guard let decrypted = decryptAES256(encrypted, pin: old) else { return false }
            decryptedPrivKey = decrypted
        }

        var decryptedMnemonic: Data?
        if let encrypted = keychain.load(.mnemonicEncrypted) {
            guard let decrypted = decryptAES256(encrypted, pin: old) else { return false }
            decryptedMnemonic = decrypted
        }

        defer {
            if var dk = decryptedPrivKey { MLDSA65.zeroMemory(&dk); decryptedPrivKey = nil }
            if var dm = decryptedMnemonic { MLDSA65.zeroMemory(&dm); decryptedMnemonic = nil }
        }

        // Phase 2: Generate new Argon2id salt + hash IN MEMORY (NOT saved to Keychain yet)
        // SECURITY: Salt is passed as parameter to encryptMT2, not read from Keychain.
        // This prevents crash-between-steps from bricking the wallet.
        var newSaltBytes = [UInt8](repeating: 0, count: 16)
        guard SecRandomCopyBytes(kSecRandomDefault, 16, &newSaltBytes) == errSecSuccess else { return false }
        let newSalt = Data(newSaltBytes)

        // Hash new PIN with Argon2id (domain-separated from encryption key)
        guard let newHash = hashPinArgon2id(from: new, salt: newSalt) else { return false }

        // Phase 3: Re-encrypt with new PIN as MT2 (salt passed as param, NOT from Keychain)
        var reEncPrivKey: Data?
        if let privKey = decryptedPrivKey {
            guard let enc = encryptMT2(privKey, pin: new, salt: newSalt) else { return false }
            reEncPrivKey = enc
        }

        var reEncMnemonic: Data?
        if let mnemonic = decryptedMnemonic {
            guard let enc = encryptMT2(mnemonic, pin: new, salt: newSalt) else { return false }
            reEncMnemonic = enc
        }

        // Phase 4: Save encrypted blobs
        var ok = true
        if let enc = reEncPrivKey {
            ok = ok && keychain.save(.privateKeyEncrypted, data: enc)
        }
        if let enc = reEncMnemonic {
            ok = ok && keychain.save(.mnemonicEncrypted, data: enc)
        }
        guard ok else { return false }

        // Phase 5: COMMIT — save salt + hash + cleanup (atomic point)
        // Only now do we modify Keychain state — all encrypted blobs already saved
        keychain.save(.argon2Salt, data: newSalt)
        keychain.save(.pinHash, data: newHash)
        keychain.delete(.pinSalt) // Remove old PBKDF2 salt (fully migrated to Argon2id)

        hasPin = true
        resetLockout()
        return true
    }

    /// Delete PIN + delete all encrypted data (keys & seed become unrecoverable)
    func deletePin() {
        keychain.delete(.pinHash)
        keychain.delete(.pinSalt)
        keychain.delete(.argon2Salt)
        keychain.delete(.privateKeyEncrypted)
        keychain.delete(.mnemonicEncrypted)
        // Clean up SE keys
        keychain.delete(.seEnabled)
        keychain.delete(.seKeyTag)
        keychain.delete(.seWrappedKey)
        hasPin = false
        resetLockout()
    }

    /// Constant-time comparison (prevents timing attacks)
    private func constantTimeEqual(_ a: Data, _ b: Data) -> Bool {
        guard a.count == b.count else { return false }
        var result: UInt8 = 0
        for i in 0..<a.count {
            result |= a[i] ^ b[i]
        }
        return result == 0
    }

    /// Verify PIN with lockout (timing-safe)
    /// Supports: Argon2id (v2), PBKDF2 (v1), SHA-256 (legacy)
    func verifyPin(_ pin: String) -> Bool {
        guard !isLockedOut else { return false }
        guard let stored = keychain.load(.pinHash) else { return false }

        let input: Data?
        if let argon2Salt = keychain.load(.argon2Salt) {
            // v2: Argon2id + HKDF domain separation (64MB, 3 iter, 4 lanes)
            input = hashPinArgon2id(from: pin, salt: argon2Salt)
        } else if let pbkdf2Salt = keychain.load(.pinSalt) {
            // v1: PBKDF2-HMAC-SHA256 (600k iterations)
            input = hashPin(pin, salt: pbkdf2Salt)
        } else {
            // Legacy: SHA-256 (no salt)
            input = hashPinLegacy(pin)
        }

        guard let input = input else { return false }

        if constantTimeEqual(stored, input) {
            resetConsecutiveFails()
            return true
        } else {
            registerFailedAttempt()
            return false
        }
    }

    // MARK: - Signing

    /// Sign a message (requires PIN verification + lockout enforcement)
    func signMessage(_ message: Data, pin: String) -> Data? {
        // SECURITY: Enforce lockout for signing — prevents PIN brute-force via sign API
        guard verifyPin(pin) else { return nil }
        guard var privateKey = loadPrivateKey(pin: pin) else {
            return nil
        }
        defer {
            MLDSA65.zeroMemory(&privateKey)
        }
        return MLDSA65.sign(message: message, privateKey: privateKey)
    }

    /// Sign a string message (convenience, requires PIN + lockout)
    func signString(_ message: String, pin: String) -> Data? {
        guard let data = message.data(using: .utf8) else { return nil }
        return signMessage(data, pin: pin)
    }

    /// Load private key decrypted with PIN
    /// No plaintext fallback when PIN is active (Keychain-authoritative check)
    private func loadPrivateKey(pin: String) -> Data? {
        // Try encrypted first (PIN-protected)
        if let encrypted = keychain.load(.privateKeyEncrypted) {
            return decryptAES256(encrypted, pin: pin)
        }
        // Plaintext ONLY if no PIN exists (Keychain-authoritative, not UserDefaults)
        guard !isPinActive else { return nil }
        return keychain.load(.privateKey)
    }

    // MARK: - Verification

    /// Verify a signature against a public key
    func verifySignature(message: Data, signature: Data, publicKey: Data) -> Bool {
        return MLDSA65.verify(message: message, signature: signature, publicKey: publicKey)
    }

    /// Validate a Montana address (format + checksum)
    func validateAddress(_ address: String) -> Bool {
        return MLDSA65.validateAddress(address)
    }

    // MARK: - Migration (MT1 → MT2)

    /// Migrate from MT1 (6-digit, PBKDF2) to MT2 (8-digit, Argon2id + optional SE)
    /// Called after user enters old 6-digit PIN + new 8-digit PIN
    /// SECURITY: Plaintext keys NEVER touch Keychain — decrypt in memory, re-encrypt directly
    func migrateMT1toMT2(oldPin: String, newPin: String) -> Bool {
        guard verifyPin(oldPin) else { return false }
        guard newPin.count == 8, newPin.allSatisfy({ $0.isNumber }) else { return false }

        // Phase 1: Decrypt with old PIN (MT1/legacy format) — keys stay in MEMORY only
        var decryptedPrivKey: Data?
        if let encrypted = keychain.load(.privateKeyEncrypted) {
            guard let decrypted = decryptAES256(encrypted, pin: oldPin) else { return false }
            decryptedPrivKey = decrypted
        }

        var decryptedMnemonic: Data?
        if let encrypted = keychain.load(.mnemonicEncrypted) {
            guard let decrypted = decryptAES256(encrypted, pin: oldPin) else { return false }
            decryptedMnemonic = decrypted
        }

        defer {
            if var dk = decryptedPrivKey { MLDSA65.zeroMemory(&dk) }
            if var dm = decryptedMnemonic { MLDSA65.zeroMemory(&dm) }
        }

        // Phase 2: Generate new Argon2id salt + SE setup
        var saltBytes = [UInt8](repeating: 0, count: 16)
        guard SecRandomCopyBytes(kSecRandomDefault, 16, &saltBytes) == errSecSuccess else { return false }
        let argon2Salt = Data(saltBytes)

        // Clean stale SE flags + setup SE (best-effort)
        keychain.delete(.seEnabled); keychain.delete(.seKeyTag); keychain.delete(.seWrappedKey)
        let _ = seCreateAndWrapKey()

        // Hash new PIN (domain-separated)
        guard let hash = hashPinArgon2id(from: newPin, salt: argon2Salt) else {
            keychain.delete(.seEnabled); keychain.delete(.seKeyTag); keychain.delete(.seWrappedKey)
            return false
        }

        // Phase 3: Re-encrypt directly with new PIN as MT2 (no plaintext in Keychain!)
        var encPrivKey: Data?
        if let privKey = decryptedPrivKey {
            guard let enc = encryptMT2(privKey, pin: newPin, salt: argon2Salt) else {
                keychain.delete(.seEnabled); keychain.delete(.seKeyTag); keychain.delete(.seWrappedKey)
                return false
            }
            encPrivKey = enc
        }
        var encMnemonic: Data?
        if let mnemonic = decryptedMnemonic {
            guard let enc = encryptMT2(mnemonic, pin: newPin, salt: argon2Salt) else {
                keychain.delete(.seEnabled); keychain.delete(.seKeyTag); keychain.delete(.seWrappedKey)
                return false
            }
            encMnemonic = enc
        }

        // Phase 4: Save new encrypted blobs (old blobs still intact as backup)
        var ok = true
        if let enc = encPrivKey {
            ok = ok && keychain.save(.privateKeyEncrypted, data: enc)
        }
        if let enc = encMnemonic {
            ok = ok && keychain.save(.mnemonicEncrypted, data: enc)
        }
        guard ok else { return false }

        // Phase 5: COMMIT — save salt + hash, delete old state
        // At this point: new encrypted blobs saved, old salt/hash still intact
        // If crash here: old PIN still works (old salt/hash present, blobs overwritten with MT2
        // which will fail to decrypt with old PIN → user must recover from seed)
        keychain.save(.argon2Salt, data: argon2Salt)
        keychain.save(.pinHash, data: hash)
        keychain.delete(.pinSalt)
        keychain.delete(.privateKey)  // Delete any plaintext leftovers
        keychain.delete(.mnemonic)

        hasPin = true
        resetLockout()
        return true
    }
}