// 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) -> String {
    var error: Unmanaged<CFError>? = nil;
    guard let accessCtrl = SecAccessControlCreateWithFlags(
       nil,
       kSecAttrAccessibleWhenUnlockedThisDeviceOnly,
       [.privateKeyUsage, .biometryCurrentSet],
       &error
    ) else {
        return "err:\(error.debugDescription)"
    }
    do {
        if (sign) {
            let privateKeyReference = try SecureEnclave.P256.Signing.PrivateKey.init(
               accessControl: accessCtrl
            );
            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.init(
               accessControl: accessCtrl
            );
            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 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 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 action = CommandLine.arguments[1];

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

if (action == "generate_secure_enclave_p256_ecsign_keypair") {
    exitWith(generateSecureEnclaveP256KeyPair(sign: true))
}

if (action == "generate_secure_enclave_p256_ecdh_keypair") {
    exitWith(generateSecureEnclaveP256KeyPair(sign: false))
}

if (action == "recover_secure_enclave_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 (action == "recover_secure_enclave_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 (action == "compute_secure_enclave_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 (action == "compute_secure_enclave_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 (action == "version") {
    exitWith("ok:1.0.0-20250118")
}

exitWith("err:invalid action")