// 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 GenerateSecureEnclaveP256KeyPairRequest {
    var controlFlag: String
}

func parseGenerateSecureEnclaveP256KeyPairRequest() -> GenerateSecureEnclaveP256KeyPairRequest? {
    var controlFlagOpt: String?
    let len = CommandLine.arguments.count;
    if CommandLine.arguments.count > 2 {
        var i = 2
        while i < len {
            let k = CommandLine.arguments[i];
            if (k == "--control-flag") {
                controlFlagOpt = CommandLine.arguments[i + 1]
                i += 2
            } else {
                i += 1
            }
        }
    }
    guard let controlFlag = controlFlagOpt else {
        exitError("parameter --control-flag required.")
        return nil
    }
    return GenerateSecureEnclaveP256KeyPairRequest(
        controlFlag: controlFlag
    )
}

struct ComputeP256EcSignRequest {
    var dataRepresentationBase64: String
    var messageBase64: String
}

func parseComputeP256EcSignRequest() -> ComputeP256EcSignRequest? {
    var dataRepresentationBase64Opt: String?
    var messageBase64Opt: String?
    let len = CommandLine.arguments.count;
    if CommandLine.arguments.count > 2 {
        var i = 2
        while i < len {
            let k = CommandLine.arguments[i];
            if (k == "--data-representation-base64" || k == "--private-key") {
                dataRepresentationBase64Opt = CommandLine.arguments[i + 1]
                i += 2
            } else if (k == "--message-base64") {
                messageBase64Opt = CommandLine.arguments[i + 1]
                i += 2
            } else {
                i += 1
            }
        }
    }
    guard let dataRepresentationBase64 = dataRepresentationBase64Opt else {
        exitError("parameter --data-representation-base64 or --private-key required.")
        return nil
    }
    guard let messageBase64 = messageBase64Opt else {
        exitError("parameter --message-base64 required.")
        return nil
    }
    return ComputeP256EcSignRequest(
        dataRepresentationBase64: dataRepresentationBase64,
        messageBase64: messageBase64
    )
}

func parseExternalSignRequest() -> ComputeP256EcSignRequest? {
    var algOpt: String?
    var dataRepresentationBase64Opt: String?
    var messageBase64Opt: String?
    let len = CommandLine.arguments.count;
    if CommandLine.arguments.count > 2 {
        var i = 2
        while i < len {
            let k = CommandLine.arguments[i];
            if (k == "--parameter") {
                dataRepresentationBase64Opt = CommandLine.arguments[i + 1]
                i += 2
            } else if (k == "--message-base64") {
                messageBase64Opt = CommandLine.arguments[i + 1]
                i += 2
            } else if (k == "--alg") {
                algOpt = CommandLine.arguments[i + 1]
                i += 2
            } else {
                i += 1
            }
        }
    }
    guard let alg = algOpt else {
        exitError("parameter --alg required.")
        return nil
    }
    if alg != "ES256" {
        exitError("parameter --alg MUST be ES256.")
        return nil
    }
    guard let dataRepresentationBase64 = dataRepresentationBase64Opt else {
        exitError("parameter --parameter required.")
        return nil
    }
    guard let messageBase64 = messageBase64Opt else {
        exitError("parameter --message-base64 required.")
        return nil
    }
    return ComputeP256EcSignRequest(
        dataRepresentationBase64: dataRepresentationBase64,
        messageBase64: messageBase64
    )
}

struct ComputeP256EcdhRequest {
    var dataRepresentationBase64: String
    var ephemeraPublicKeyBase64: String
}

func parseComputeP256EcdhRequest() -> ComputeP256EcdhRequest? {
    var dataRepresentationBase64Opt: String?
    var ephemeraPublicKeyBase64Opt: String?
    let len = CommandLine.arguments.count;
    if CommandLine.arguments.count > 2 {
        var i = 2
        while i < len {
            let k = CommandLine.arguments[i];
            if (k == "--data-representation-base64" || k == "--private-key") {
                dataRepresentationBase64Opt = CommandLine.arguments[i + 1]
                i += 2
            } else if (k == "--ephemera-public-key-base64" || k == "--ephemera-public-key") {
                ephemeraPublicKeyBase64Opt = CommandLine.arguments[i + 1]
                i += 2
            } else {
                i += 1
            }
        }
    }
    guard let dataRepresentationBase64 = dataRepresentationBase64Opt else {
        exitError("parameter --data-representation-base64 or --private-key required.")
        return nil
    }
    guard let ephemeraPublicKeyBase64 = ephemeraPublicKeyBase64Opt else {
        exitError("parameter --ephemera-public-key-base64 required.")
        return nil
    }
    return ComputeP256EcdhRequest(
        dataRepresentationBase64: dataRepresentationBase64,
        ephemeraPublicKeyBase64: ephemeraPublicKeyBase64
    )
}

struct RecoverSecureEnclaveP256PublicKeyRequest {
    var dataRepresentationBase64: String
}

func parseRecoverSecureEnclaveP256PublicKeyRequest() -> RecoverSecureEnclaveP256PublicKeyRequest? {
    var dataRepresentationBase64Opt: String?
    let len = CommandLine.arguments.count;
    if CommandLine.arguments.count > 2 {
        var i = 2
        while i < len {
            let k = CommandLine.arguments[i];
            if (k == "--data-representation-base64" || k == "--private-key") {
                dataRepresentationBase64Opt = CommandLine.arguments[i + 1]
                i += 2
            } else {
                i += 1
            }
        }
    }
    guard let dataRepresentationBase64 = dataRepresentationBase64Opt else {
        exitError("parameter --data-representation-base64 or --private-key required.")
        return nil
    }
    return RecoverSecureEnclaveP256PublicKeyRequest(
        dataRepresentationBase64: dataRepresentationBase64
    )
}

func parseExternalPublicKeyRequest() -> RecoverSecureEnclaveP256PublicKeyRequest? {
    var dataRepresentationBase64Opt: String?
    let len = CommandLine.arguments.count;
    if CommandLine.arguments.count > 2 {
        var i = 2
        while i < len {
            let k = CommandLine.arguments[i];
            if (k == "--parameter") {
                dataRepresentationBase64Opt = CommandLine.arguments[i + 1]
                i += 2
            } else {
                i += 1
            }
        }
    }
    guard let dataRepresentationBase64 = dataRepresentationBase64Opt else {
        exitError("parameter --parameter required.")
        return nil
    }
    return RecoverSecureEnclaveP256PublicKeyRequest(
        dataRepresentationBase64: dataRepresentationBase64
    )
}

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 public_key_point_base64: String
    var public_key_base64: String
    var data_representation_base64: String
}

struct ExternalPublicKeyResponse: Codable {
    var success: Bool
    var public_key_base64: String
}

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

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

struct ExternalSpecResponse: Codable {
    var success: Bool
    var agent: String
    var specification: 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, request: GenerateSecureEnclaveP256KeyPairRequest) -> GenerateSecureEnclaveP256KeyPairResponse? {
    var error: Unmanaged<CFError>? = nil
    let accessControlCreateFlags: SecAccessControlCreateFlags
    let controlFlag = request.controlFlag
    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, 
        public_key_point_base64: publicKeyPointBase64,
        public_key_base64: publicKeyDerBase64,
        data_representation_base64: 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, 
        public_key_point_base64: publicKeyPointBase64,
        public_key_base64: publicKeyDerBase64,
        data_representation_base64: dataRepresentationBase64
    )
}

func recoverSecureEnclaveP256PublicKey(request: RecoverSecureEnclaveP256PublicKeyRequest, sign: Bool) -> GenerateSecureEnclaveP256KeyPairResponse? {
    guard let privateKeyDataRepresentation = Data(
        base64Encoded: request.dataRepresentationBase64
    ) 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(request: ComputeP256EcSignRequest) -> ComputeSecureEnclaveP256EcsignResponse? {
    guard let privateKeyDataRepresentation = Data(
        base64Encoded: request.dataRepresentationBase64
    ) else {
        exitError("private key base64 decode failed")
        return nil
    }
    guard let contentData = Data(
        base64Encoded: request.messageBase64
    ) 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,
            signature_base64: signature.derRepresentation.base64EncodedString()
        )
    } catch {
        exitError("\(error)")
        return nil
    }
}

func computeSecureEnclaveP256Ecdh(request: ComputeP256EcdhRequest) -> ComputeSecureEnclaveP256EcdhResponse? {
    guard let privateKeyDataRepresentation = Data(
        base64Encoded: request.dataRepresentationBase64
    ) else {
        exitError("private key base64 decode failed")
        return nil
    }
    guard let ephemeralPublicKeyRepresentation = Data(
        base64Encoded: request.ephemeraPublicKeyBase64
    ) 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,
            shared_secret: sharedSecret.description
        )
    } catch {
        exitError("\(error)")
        return nil
    }
}

func externalSpec() -> ExternalSpecResponse {
    return ExternalSpecResponse(
        success: true,
        agent: "swift-secure-enclave-external-provider/2.0.0-alpha",
        specification: "External/1.0.0-alpha"
    )
}

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") {
    let request = parseGenerateSecureEnclaveP256KeyPairRequest()!;
    exitOkWithJson(generateSecureEnclaveP256KeyPair(sign: true, request: request))
}

if (command == "generate_p256_ecdh_keypair") {
    let request = parseGenerateSecureEnclaveP256KeyPairRequest()!;
    exitOkWithJson(generateSecureEnclaveP256KeyPair(sign: false, request: request))
}

if (command == "recover_p256_ecsign_public_key") {
    let request = parseRecoverSecureEnclaveP256PublicKeyRequest()!
    let response =  recoverSecureEnclaveP256PublicKey(request: request, sign: true)
    exitOkWithJson(response)
}

if (command == "recover_p256_ecdh_public_key") {
    let request = parseRecoverSecureEnclaveP256PublicKeyRequest()!
    let response =  recoverSecureEnclaveP256PublicKey(request: request, sign: false)
    exitOkWithJson(response)
}

if (command == "compute_p256_ecsign") {
    let request = parseComputeP256EcSignRequest()!;
    let response =  computeSecureEnclaveP256Ecsign(request: request)
    exitOkWithJson(response)
}

if (command == "compute_p256_ecdh") {
    let request = parseComputeP256EcdhRequest()!;
    let response =  computeSecureEnclaveP256Ecdh(request: request)
    exitOkWithJson(response)
}

if (command == "external_spec") {
    exitOkWithJson(externalSpec())
}

if (command == "external_public_key") {
    let request = parseExternalPublicKeyRequest()!
    let response =  recoverSecureEnclaveP256PublicKey(request: request, sign: true)!
    exitOkWithJson(ExternalPublicKeyResponse(success: true, public_key_base64: response.public_key_base64))
}

if (command == "external_sign") {
    let request = parseExternalSignRequest()!;
    let response =  computeSecureEnclaveP256Ecsign(request: request)
    exitOkWithJson(response)
}

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

if (command == "help" || command == "-h" || command == "--help") {
    print("swift-secure-enclave-tool-v2 <command> [parameters]")
    print("help                                                           - print help")
    print("version                                                        - print version")
    print("is_support_secure_enclave                                      - is Secure Enclave supported")
    print("generate_p256_ecsign_keypair --control-flag <>                 - generate Secure Enclave P256 EC sign key pair")
    print("generate_p256_ecdh_keypair --control-flag <>                   - generate Secure Enclave P256 EC DH key pair")
    print("recover_p256_ecsign_public_key --private-key <>                - recover Secure Enclave P256 EC sign key pair")
    print("recover_p256_ecdh_public_key --private-key <>                  - recover Secure Enclave P256 EC DH key pair")
    print("compute_p256_ecsign --private-key <> --message-base64 <>       - compure Secure Enclave P256 EC sign")
    print("compute_p256_ecdh --private-key <> --ephemera-public-key <>    - compure Secure Enclave P256 EC DH")
    print("external_spec                                                  - external specification")
    print("external_public_key --parameter <>                             - external public key")
    print("external_sign --parameter <> --alg ES256  --message-base64 <>  - external sign")
    print()
    print("options:")
    print("> --control-flag           - none, userPresence, devicePasscode, biometryAny, biometryCurrentSet")
    print("> --private-key            - private key representation (dataRepresentationBase64)")
    print("> --message-base64         - content in base64")
    print("> --ephemera-public-key    - public key der in base64")
    exit(0)
}

exitError("invalid command")