import Foundation
import CryptoKit
import CommonCrypto

/// ECP1 key encryption format: PBKDF2-HMAC-SHA256 (600k iterations) + AES-256-GCM
/// Wire format: magic(4) + salt(16) + nonce(12) + ciphertext_with_tag(N+16)
enum KeyEncryption {

    /// Encrypt raw key bytes with password using ECP1 format
    static func encrypt(_ rawBytes: Data, password: Data) throws -> Data {
        let salt = Data.randomBytes(16)
        let derivedKey = try pbkdf2(password: password, salt: salt)

        let nonce = Data.randomBytes(12)
        let symmetricKey = SymmetricKey(data: derivedKey)
        let gcmNonce = try AES.GCM.Nonce(data: nonce)

        // AAD = ECP1 magic bytes (matching Python)
        let sealedBox = try AES.GCM.seal(
            rawBytes,
            using: symmetricKey,
            nonce: gcmNonce,
            authenticating: Constants.ecp1Magic
        )

        // ciphertext + tag concatenated (matches Python's AESGCM.encrypt output)
        var result = Data()
        result.append(Constants.ecp1Magic)    // 4 bytes
        result.append(salt)                    // 16 bytes
        result.append(nonce)                   // 12 bytes
        result.append(sealedBox.ciphertext)    // N bytes
        result.append(sealedBox.tag)           // 16 bytes
        return result
    }

    /// Decrypt ECP1-encrypted key bytes with password
    static func decrypt(_ data: Data, password: Data) throws -> Data {
        guard data.count >= 48 else {  // 4 + 16 + 12 + 16 minimum
            throw CryptoError.invalidECP1Format
        }
        guard data.prefix(4) == Constants.ecp1Magic else {
            throw CryptoError.invalidECP1Format
        }

        let salt = data[4..<20]
        let nonce = data[20..<32]
        let ctWithTag = data[32...]

        guard ctWithTag.count >= 16 else {
            throw CryptoError.invalidECP1Format
        }

        let derivedKey = try pbkdf2(password: password, salt: Data(salt))
        let symmetricKey = SymmetricKey(data: derivedKey)
        let gcmNonce = try AES.GCM.Nonce(data: nonce)

        // Split ciphertext and tag
        let ct = ctWithTag.prefix(ctWithTag.count - 16)
        let tag = ctWithTag.suffix(16)

        let sealedBox = try AES.GCM.SealedBox(
            nonce: gcmNonce,
            ciphertext: ct,
            tag: tag
        )

        do {
            return try AES.GCM.open(sealedBox, using: symmetricKey, authenticating: Constants.ecp1Magic)
        } catch {
            throw CryptoError.decryptionFailed("ECP1 decryption failed - wrong password?")
        }
    }

    /// Check if data starts with ECP1 magic
    static func isECP1Format(_ data: Data) -> Bool {
        data.count >= 4 && data.prefix(4) == Constants.ecp1Magic
    }

    // MARK: - PBKDF2

    /// Derive 32-byte key using PBKDF2-HMAC-SHA256 with 600k iterations
    static func pbkdf2(password: Data, salt: Data) throws -> Data {
        var derivedKey = Data(count: 32)
        let status = derivedKey.withUnsafeMutableBytes { derivedKeyPtr in
            password.withUnsafeBytes { passwordPtr in
                salt.withUnsafeBytes { saltPtr in
                    CCKeyDerivationPBKDF(
                        CCPBKDFAlgorithm(kCCPBKDF2),
                        passwordPtr.baseAddress?.assumingMemoryBound(to: Int8.self),
                        password.count,
                        saltPtr.baseAddress?.assumingMemoryBound(to: UInt8.self),
                        salt.count,
                        CCPseudoRandomAlgorithm(kCCPRFHmacAlgSHA256),
                        Constants.pbkdf2Iterations,
                        derivedKeyPtr.baseAddress?.assumingMemoryBound(to: UInt8.self),
                        32
                    )
                }
            }
        }
        guard status == kCCSuccess else {
            throw CryptoError.pbkdf2Failed
        }
        return derivedKey
    }
}