// 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
// - https://gist.github.com/monsoir/cc06c298c9c37d629942d87b14a202fc
// - https://developer.apple.com/documentation/foundation/jsonencoder

import CryptoKit
import Foundation
import LocalAuthentication

struct ErrorResponse: Codable {
    var success: Bool
    var error: String
}

struct VersionResponse: Codable {
    var success: Bool
    var version: String
}

struct SupportSecureEnclaveResponse: Codable {
    var success: Bool
    var supported: Bool
}

struct GenerateSecureEnclaveP256KeyPairResponse: Codable {
    var success: Bool
    var publicKeyPointBase64: String
    var publicKeyDerBase64: String
    var dataRepresentationBase64: String
}

struct ComputeSecureEnclaveP256EcsignResponse: Codable {
    var success: Bool
    var signatureDerBase64: String
}

struct ComputeSecureEnclaveP256EcdhResponse: Codable {
    var success: Bool
    var sharedSecret: String
}

func stringify<T: Encodable>(_ value: T) -> String? {
    guard let jsonData = try? JSONEncoder().encode(value) else { return nil }
    let result = String(data: jsonData, encoding: .utf8)
    return result
}

func exitError(_ error: String) {
    if let response = stringify(ErrorResponse(success: false, error: error)) {
        print(response)
    } else {
        print("{\"success\":false,\"error\":\"JSON stringify erorr, raw error: \(error)\"}")
    }
    exit(1)
}

func exitOkWithJson<T: Encodable>(_ value: T) {
    if let response = stringify(value) {
        print(response)
        exit(0)
    } else {
        print("{\"success\":false,\"error\":\"JSON stringify erorr\"}")
        exit(1)
    }
}

func isSupportSecureEnclave() -> SupportSecureEnclaveResponse {
    return SupportSecureEnclaveResponse(success: true, supported: SecureEnclave.isAvailable)
}

func generateSecureEnclaveP256KeyPair(sign: Bool, controlFlag: String) -> GenerateSecureEnclaveP256KeyPairResponse? {
    var error: Unmanaged<CFError>? = nil
    let accessControlCreateFlags: SecAccessControlCreateFlags
    if (controlFlag == "none") {
        accessControlCreateFlags = [.privateKeyUsage]
    } else if (controlFlag == "userPresence") {
        accessControlCreateFlags = [.privateKeyUsage, .userPresence]
    } else if (controlFlag == "devicePasscode") {
        accessControlCreateFlags = [.privateKeyUsage, .devicePasscode]
    } else if (controlFlag == "biometryAny") {
        accessControlCreateFlags = [.privateKeyUsage, .biometryAny]
    }  else if (controlFlag == "biometryCurrentSet") {
        accessControlCreateFlags = [.privateKeyUsage, .biometryCurrentSet]
    } else {
       exitError("invalid control flag: \(controlFlag)")
       return nil
    }
    guard let accessCtrl = SecAccessControlCreateWithFlags(
       nil,
       kSecAttrAccessibleWhenUnlockedThisDeviceOnly,
       accessControlCreateFlags,
       &error
    ) else {
        exitError(error.debugDescription)
        return nil
    }
    do {
        if (sign) {
            let privateKeyReference = try SecureEnclave.P256.Signing.PrivateKey.init(
               accessControl: accessCtrl
            )
            return signingPrivateKeyToResponse(privateKeyReference)
        } else {
            let privateKeyReference = try SecureEnclave.P256.KeyAgreement.PrivateKey.init(
               accessControl: accessCtrl
            )
            return keyAgreementPrivateKeyToResponse(privateKeyReference)
        }
    } catch {
        exitError("\(error)")
        return nil
    }
}

func signingPrivateKeyToResponse(_ privateKeyReference: SecureEnclave.P256.Signing.PrivateKey) -> GenerateSecureEnclaveP256KeyPairResponse {
    let publicKeyPointBase64 = privateKeyReference.publicKey.x963Representation.base64EncodedString()
    let publicKeyDerBase64 = privateKeyReference.publicKey.derRepresentation.base64EncodedString()
    let dataRepresentationBase64 = privateKeyReference.dataRepresentation.base64EncodedString()
    return GenerateSecureEnclaveP256KeyPairResponse(
        success: true, 
        publicKeyPointBase64: publicKeyPointBase64,
        publicKeyDerBase64: publicKeyDerBase64,
        dataRepresentationBase64: dataRepresentationBase64
    )
}

func keyAgreementPrivateKeyToResponse(_ privateKeyReference: SecureEnclave.P256.KeyAgreement.PrivateKey) -> GenerateSecureEnclaveP256KeyPairResponse {
    let publicKeyPointBase64 = privateKeyReference.publicKey.x963Representation.base64EncodedString()
    let publicKeyDerBase64 = privateKeyReference.publicKey.derRepresentation.base64EncodedString()
    let dataRepresentationBase64 = privateKeyReference.dataRepresentation.base64EncodedString()
    return GenerateSecureEnclaveP256KeyPairResponse(
        success: true, 
        publicKeyPointBase64: publicKeyPointBase64,
        publicKeyDerBase64: publicKeyDerBase64,
        dataRepresentationBase64: dataRepresentationBase64
    )
}

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

func recoverSecureEnclaveP256PublicKey(privateKeyDataRepresentation: String, sign: Bool) -> GenerateSecureEnclaveP256KeyPairResponse? {
    guard let privateKeyDataRepresentation = Data(
        base64Encoded: privateKeyDataRepresentation
    ) else {
       exitError("private key base64 decode failed")
       return nil
    }
    do {
        let context = LAContext()
        if (sign)  {
            let privateKeyReference =  try SecureEnclave.P256.Signing.PrivateKey(
                dataRepresentation: privateKeyDataRepresentation,
                authenticationContext: context
            )
            return signingPrivateKeyToResponse(privateKeyReference)
        } else {
            let privateKeyReference =  try SecureEnclave.P256.KeyAgreement.PrivateKey(
                dataRepresentation: privateKeyDataRepresentation,
                authenticationContext: context
            )
            return keyAgreementPrivateKeyToResponse(privateKeyReference)
        }
    } catch {
        exitError("\(error)")
        return nil
    }
}

func computeSecureEnclaveP256Ecsign(privateKeyDataRepresentation: String, content: String) -> ComputeSecureEnclaveP256EcsignResponse? {
    guard let privateKeyDataRepresentation = Data(
        base64Encoded: privateKeyDataRepresentation
    ) else {
        exitError("private key base64 decode failed")
        return nil
    }
    guard let contentData = Data(
        base64Encoded: content
    ) else {
        exitError("content base64 decode failed")
        return nil
    }
    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 ComputeSecureEnclaveP256EcsignResponse(
            success: true,
            signatureDerBase64: signature.derRepresentation.base64EncodedString()
        )
    } catch {
        exitError("\(error)")
        return nil
    }
}

func computeSecureEnclaveP256Ecdh(privateKeyDataRepresentation: String, ephemeraPublicKey: String) -> ComputeSecureEnclaveP256EcdhResponse? {
    guard let privateKeyDataRepresentation = Data(
        base64Encoded: privateKeyDataRepresentation
    ) else {
        exitError("private key base64 decode failed")
        return nil
    }
    guard let ephemeralPublicKeyRepresentation = Data(
        base64Encoded: ephemeraPublicKey
    ) else {
        exitError("ephemeral public key base64 decode failed")
        return nil
    }
    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 ComputeSecureEnclaveP256EcdhResponse(
            success: true,
            sharedSecret: sharedSecret.description
        )
    } catch {
        exitError("\(error)")
        return nil
    }
}

if (CommandLine.arguments.count == 1) {
    exitError("require at least one argument")
}

let command = CommandLine.arguments[1]

if (command == "is_support_secure_enclave") {
    exitOkWithJson(isSupportSecureEnclave())
}

if (command == "generate_p256_ecsign_keypair") {
    if (CommandLine.arguments.count != 3) {
        exitError("require two arguments")
    }
    exitOkWithJson(generateSecureEnclaveP256KeyPair(sign: true, controlFlag: CommandLine.arguments[2]))
}

if (command == "generate_p256_ecdh_keypair") {
    if (CommandLine.arguments.count != 3) {
        exitError("require two arguments")
    }
    exitOkWithJson(generateSecureEnclaveP256KeyPair(sign: false, controlFlag: CommandLine.arguments[2]))
}

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

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

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

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

if (command == "version") {
    exitOkWithJson(VersionResponse(success: true, version: "2.0.0-20250428"))
}

if (command == "help") {
    print("swift-secure-enclave-tool-v2 <command>")
    print("help                                               - print help")
    print("version                                            - print version")
    print("is_support_secure_enclave                          - is Secure Enclave supported")
    print("generate_p256_ecsign_keypair <controlFlag>         - generate Secure Enclave P256 EC sign key pair")
    print("generate_p256_ecdh_keypair <controlFlag>           - 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("> controlFlag              - none, userPresence, devicePasscode, biometryAny, biometryCurrentSet")
    print("> privateKey               - private key representation (dataRepresentationBase64)")
    print("> content                  - content in base64")
    print("> ephemeraPublicKey        - public key der in base64")
    exit(0)
}

exitError("invalid command")