// Reference:
// - https://developer.apple.com/documentation/swift/commandline/arguments
// - https://git.hatter.ink/hatter/card-cli/src/branch/master/swift-lib/src/lib.swift
// - https://developer.apple.com/documentation/security/secaccesscontrolcreateflags

import CryptoKit
import LocalAuthentication

func isSupportSecureEnclave() -> Bool {
    return SecureEnclave.isAvailable
}

func generateSecureEnclaveP256KeyPair(sign: Bool, requireBio: Bool) -> String {
    var error: Unmanaged<CFError>? = nil
    let accessControlCreateFlags: SecAccessControlCreateFlags
    if (requireBio) {
        accessControlCreateFlags = [.privateKeyUsage, .biometryCurrentSet]
    } else {
        accessControlCreateFlags = [.privateKeyUsage]
    }
    guard let accessCtrl = SecAccessControlCreateWithFlags(
       nil,
       kSecAttrAccessibleWhenUnlockedThisDeviceOnly,
       accessControlCreateFlags,
       &error
    ) else {
        return "err:\(error.debugDescription)"
    }
    do {
        if (sign) {
            let privateKeyReference = try SecureEnclave.P256.Signing.PrivateKey.init(
               accessControl: accessCtrl
            )
            let publicKeyPointBase64 = privateKeyReference.publicKey.x963Representation.base64EncodedString()
            let publicKeyDerBase64 = privateKeyReference.publicKey.derRepresentation.base64EncodedString()
            let dataRepresentationBase64 = privateKeyReference.dataRepresentation.base64EncodedString()
            return "ok:\(publicKeyPointBase64),\(publicKeyDerBase64),\(dataRepresentationBase64)"
        } else {
            let privateKeyReference = try SecureEnclave.P256.KeyAgreement.PrivateKey.init(
               accessControl: accessCtrl
            )
            let publicKeyPointBase64 = privateKeyReference.publicKey.x963Representation.base64EncodedString()
            let publicKeyDerBase64 = privateKeyReference.publicKey.derRepresentation.base64EncodedString()
            let dataRepresentationBase64 = privateKeyReference.dataRepresentation.base64EncodedString()
            return "ok:\(publicKeyPointBase64),\(publicKeyDerBase64),\(dataRepresentationBase64)"
        }
    } catch {
        return "err:\(error)"
    }
}

func recoverSecureEnclaveP256PublicKeyEcsign(privateKeyDataRepresentation: String) -> String {
    return recoverSecureEnclaveP256PublicKey(privateKeyDataRepresentation: privateKeyDataRepresentation, sign: true)
}

func recoverSecureEnclaveP256PublicKey(privateKeyDataRepresentation: String, sign: Bool) -> String {
    guard let privateKeyDataRepresentation = Data(
        base64Encoded: privateKeyDataRepresentation
    ) else {
       return "err:private key base64 decode failed"
    }
    do {
        let context = LAContext()
        if (sign)  {
            let privateKeyReference =  try SecureEnclave.P256.Signing.PrivateKey(
                dataRepresentation: privateKeyDataRepresentation,
                authenticationContext: context
            )
            let publicKeyBase64 = privateKeyReference.publicKey.x963Representation.base64EncodedString()
            let publicKeyPem = privateKeyReference.publicKey.derRepresentation.base64EncodedString()
            let dataRepresentationBase64 = privateKeyReference.dataRepresentation.base64EncodedString()
            return "ok:\(publicKeyBase64),\(publicKeyPem),\(dataRepresentationBase64)"
        } else {
            let privateKeyReference =  try SecureEnclave.P256.KeyAgreement.PrivateKey(
                dataRepresentation: privateKeyDataRepresentation,
                authenticationContext: context
            )
            let publicKeyBase64 = privateKeyReference.publicKey.x963Representation.base64EncodedString()
            let publicKeyPem = privateKeyReference.publicKey.derRepresentation.base64EncodedString()
            let dataRepresentationBase64 = privateKeyReference.dataRepresentation.base64EncodedString()
            return "ok:\(publicKeyBase64),\(publicKeyPem),\(dataRepresentationBase64)"
        }
    } catch {
        return "err:\(error)"
    }
}

func computeSecureEnclaveP256Ecsign(privateKeyDataRepresentation: String, content: String) -> String {
    guard let privateKeyDataRepresentation = Data(
        base64Encoded: privateKeyDataRepresentation
    ) else {
       return "err:private key base64 decode failed"
    }
    guard let contentData = Data(
        base64Encoded: content
    ) else {
       return "err:content base64 decode failed"
    }
    do {
        let context = LAContext()
        let p =  try SecureEnclave.P256.Signing.PrivateKey(
            dataRepresentation: privateKeyDataRepresentation,
            authenticationContext: context
        )

        let digest = SHA256.hash(data: contentData)
        let signature = try p.signature(for: digest)

        return "ok:\(signature.derRepresentation.base64EncodedString())"
    } catch {
        return "err:\(error)"
    }
}

func computeSecureEnclaveP256Ecdh(privateKeyDataRepresentation: String, ephemeraPublicKey: String) -> String {
    guard let privateKeyDataRepresentation = Data(
        base64Encoded: privateKeyDataRepresentation
    ) else {
       return "err:private key base64 decode failed"
    }
    guard let ephemeralPublicKeyRepresentation = Data(
        base64Encoded: ephemeraPublicKey
    ) else {
       return "err:ephemeral public key base64 decode failed"
    }
    do {
        let context = LAContext()
        let p =  try SecureEnclave.P256.KeyAgreement.PrivateKey(
            dataRepresentation: privateKeyDataRepresentation,
            authenticationContext: context
        )

        let ephemeralPublicKey = try P256.KeyAgreement.PublicKey.init(derRepresentation: ephemeralPublicKeyRepresentation)

        let sharedSecret = try p.sharedSecretFromKeyAgreement(
            with: ephemeralPublicKey)

        return "ok:\(sharedSecret.description)"
    } catch {
        return "err:\(error)"
    }
}

func readArgumentAsBool(index: Int, defaultValue: Bool) -> Bool {
    if CommandLine.arguments.count >= index + 1 {
        let val = CommandLine.arguments[index]
        if (val == "true" || val == "yes" || val == "on" || val == "1") {
            return true
        }
        return false
    }
    return defaultValue
}

func exitWith(_ response: String) {
    print(response)
    if (response.hasPrefix("ok:")) {
        exit(0)
    } else {
        exit(1)
    }
}


if (CommandLine.arguments.count == 1) {
    exitWith("err:require one argument")
}

let command = CommandLine.arguments[1]

if (command == "is_support_secure_enclave") {
    exitWith("ok:\(isSupportSecureEnclave())")
}

if (command == "generate_p256_ecsign_keypair") {
    let requireBio = readArgumentAsBool(index: 2, defaultValue: true)
    exitWith(generateSecureEnclaveP256KeyPair(sign: true, requireBio: requireBio))
}

if (command == "generate_p256_ecdh_keypair") {
    let requireBio = readArgumentAsBool(index: 2, defaultValue: true)
    exitWith(generateSecureEnclaveP256KeyPair(sign: false, requireBio: requireBio))
}

if (command == "recover_p256_ecsign_public_key") {
    if (CommandLine.arguments.count != 3) {
        exitWith("err:require two arguments")
    }
    let response =  recoverSecureEnclaveP256PublicKey(
        privateKeyDataRepresentation: CommandLine.arguments[2], sign: true)
    exitWith(response)
}

if (command == "recover_p256_ecdh_public_key") {
    if (CommandLine.arguments.count != 3) {
        exitWith("err:require two arguments")
    }
    let response =  recoverSecureEnclaveP256PublicKey(
        privateKeyDataRepresentation: CommandLine.arguments[2], sign: false)
    exitWith(response)
}

if (command == "compute_p256_ecsign") {
    if (CommandLine.arguments.count != 4) {
        exitWith("err:require three arguments")
    }
    let response =  computeSecureEnclaveP256Ecsign(
        privateKeyDataRepresentation: CommandLine.arguments[2],
        content: CommandLine.arguments[3]
    )
    exitWith(response)
}

if (command == "compute_p256_ecdh") {
    if (CommandLine.arguments.count != 4) {
        exitWith("err:require three arguments")
    }
    let response =  computeSecureEnclaveP256Ecdh(
        privateKeyDataRepresentation: CommandLine.arguments[2],
        ephemeraPublicKey: CommandLine.arguments[3]
    )
    exitWith(response)
}

if (command == "version") {
    exitWith("ok:1.0.0-20250118")
}

if (command == "help") {
    print("swift-secure-enclave-tool <command>")
    print("help                                               - print help")
    print("version                                            - print version")
    print("is_support_secure_enclave                          - is Secure Enclave supported")
    print("generate_p256_ecsign_keypair [requireBio]          - generate Secure Enclave P256 EC sign key pair")
    print("generate_p256_ecdh_keypair [requireBio]            - generate Secure Enclave P256 EC DH key pair")
    print("recover_p256_ecsign_public_key <privateKey>        - recover Secure Enclave P256 EC sign key pair")
    print("recover_p256_ecdh_public_key <privateKey>          - recover Secure Enclave P256 EC DH key pair")
    print("compute_p256_ecsign <privateKey> <content>         - compure Secure Enclave P256 EC sign")
    print("compute_p256_ecdh <privateKey> <ephemeraPublicKey> - compure Secure Enclave P256 EC DH")
    print()
    print("options:")
    print("> requireBio               - true or false (default true)")
    print("> privateKey               - private key representation (dataRepresentationBase64)")
    print("> content                  - content in base64")
    print("> ephemeraPublicKey        - public key der in base64")
    print()
    print("generate secure enclave key pair format:")
    print("> ok:publicKeyPointBase64,publicKeyDerBase64,dataRepresentationBase64")
    exit(0)
}

exitWith("err:invalid command")