#!/bin/bash

operator="${KEY_MANAGER_CLI_BIN:-./key_manager}"
algorithm=$1
encoding="pem"
version=$2

set -e

export BAO_ADDR=http://127.0.0.1:8200/
export BAO_TOKEN=

validate_input() {
    if [ $# -ne 2 ]; then
        echo "Usage: $0 <algorithm> <version>" >&2
        echo "Supported algorithms: rsa_2048, rsa_3072, rsa_4096, sm2, ed25519" >&2
        exit 1
    fi

    if ! [[ "$version" =~ ^[0-9]+$ ]] || [ "$version" -lt 1 ]; then
        echo "Version must be an integer greater than 0" >&2
        exit 1
    fi

    case "$algorithm" in
        rsa_2048|rsa_3072|rsa_4096|sm2|ed25519) ;;
        *)
            echo "Unsupported algorithm: $algorithm" >&2
            echo "Supported algorithms: rsa_2048, rsa_3072, rsa_4096, sm2, ed25519" >&2
            exit 1
            ;;
    esac
}

generate_key() {
    local algorithm=$1
    local key_content

    case "$algorithm" in
        rsa_2048)
            openssl genpkey -algorithm RSA-PSS \
                -pkeyopt rsa_keygen_bits:2048 \
                -pkeyopt rsa_pss_keygen_md:sha256 \
                -pkeyopt rsa_pss_keygen_mgf1_md:sha256 \
                -pkeyopt rsa_pss_keygen_saltlen:32 \
                2>/dev/null
            ;;
        rsa_3072)
            openssl genpkey -algorithm RSA-PSS \
                -pkeyopt rsa_keygen_bits:3072 \
                -pkeyopt rsa_pss_keygen_md:sha256 \
                -pkeyopt rsa_pss_keygen_mgf1_md:sha256 \
                -pkeyopt rsa_pss_keygen_saltlen:32 \
                2>/dev/null
            ;;
        rsa_4096)
            openssl genpkey -algorithm RSA-PSS \
                -pkeyopt rsa_keygen_bits:4096 \
                -pkeyopt rsa_pss_keygen_md:sha256 \
                -pkeyopt rsa_pss_keygen_mgf1_md:sha256 \
                -pkeyopt rsa_pss_keygen_saltlen:32 \
                2>/dev/null
            ;;
        sm2)
            openssl genpkey -algorithm EC \
                -pkeyopt ec_paramgen_curve:SM2 \
                -pkeyopt ec_param_enc:named_curve \
                2>/dev/null
            ;;
        ed25519)
            openssl genpkey -algorithm ED25519 2>/dev/null
            ;;
        *)
            echo "Unsupported algorithm: $algorithm" >&2
            exit 1
            ;;
    esac
}

if [ $# -ne 2 ]; then
    echo "Usage: $0 <algorithm> <version>" >&2
    echo "Supported algorithms: rsa_2048, rsa_3072, rsa_4096, sm2, ed25519" >&2
    exit 1
fi

validate_input "$@"

for ((i=1; i<=version; i++)); do
    # Generate a new key for each signing role and import it into Key Manager.
    nsk=$(generate_key "$algorithm")
    nsk_file=$(mktemp /tmp/nsk_XXXXXX.pem)
    echo "$nsk" > "$nsk_file"

    fsk=$(generate_key "$algorithm")
    fsk_file=$(mktemp /tmp/fsk_XXXXXX.pem)
    echo "$fsk" > "$fsk_file"

    tsk=$(generate_key "$algorithm")
    tsk_file=$(mktemp /tmp/tsk_XXXXXX.pem)
    echo "$tsk" > "$tsk_file"
    "${operator}" put --key_name NSK --algorithm "${algorithm}" --encoding "${encoding}" --key_file "${nsk_file}"
    "${operator}" put --key_name FSK --algorithm "${algorithm}" --encoding "${encoding}" --key_file "${fsk_file}"
    "${operator}" put --key_name TSK --algorithm "${algorithm}" --encoding "${encoding}" --key_file "${tsk_file}"
    rm -f "$nsk_file" "$fsk_file" "$tsk_file"
done

echo "Test key data generated successfully."