class AndroidKeyStoreClient(
    private val useStrongKeyBox: Boolean = false
) {
    private val keyStoreAlias = "key_store_alias"
    private val KEY_STORE_PROVIDER = "AndroidKeyStore"

    private val keyStore by lazy {
        KeyStore.getInstance(KEY_STORE_PROVIDER).apply {
            load(null)
        }
    }

    private val cipher by lazy {
        Cipher.getInstance("AES/GCM/NoPadding")
    }

    private fun generateSecretKey(): SecretKey {
        val keyStore = keyStore.getEntry(keyStoreAlias, null)
        if (keyStore != null) {
            return (keyStore as KeyStore.SecretKeyEntry).secretKey
        }

        return KeyGenerator.getInstance(KeyProperties.KEY_ALGORITHM_AES, KEY_STORE_PROVIDER)
            .apply {
                init(
                    KeyGenParameterSpec.Builder(
                        keyStoreAlias,
                        KeyProperties.PURPOSE_ENCRYPT or KeyProperties.PURPOSE_DECRYPT
                    ).setBlockModes(KeyProperties.BLOCK_MODE_GCM)
                        .setEncryptionPaddings(KeyProperties.ENCRYPTION_PADDING_NONE)
                        .setIsStrongBoxBacked(useStrongKeyBox)
                        .setKeySize(256)
                        .build()
                )
            }.generateKey()
    }

    fun encrypt(inputByteArray: ByteArray): Result<String> {
        return runCatching {
            val secretKey = generateSecretKey().getOrThrow()
            cipher.init(Cipher.ENCRYPT_MODE, secretKey)
            val encryptedData = cipher.doFinal(inputByteArray)

            cipher.iv.joinToString("|") + ":iv:" + encryptedData.joinToString("|")
        }
    }

    fun decrypt(inputEncryptedString: String): Result<ByteArray> {
        return runCatching {
            val (ivString, encryptedString) = inputEncryptedString.split(":iv:", limit = 2)
            val iv = ivString.split("|").map { it.toByte() }.toByteArray()
            val encryptedData = encryptedString.split("|").map { it.toByte() }.toByteArray()

            val secretKey = generateSecretKey()
            val gcmParameterSpec = GCMParameterSpec(128, iv)

            cipher.init(Cipher.DECRYPT_MODE, secretKey, gcmParameterSpec)
            cipher.doFinal(encryptedData)
        }
    }
}

class TinkClient(
  context: Context
) {
    val keysetName = "key_set"
    val prefFileName = "pref_file"
    val packageName = context.packageName

    var aead: Aead

    init {
        AeadConfig.register()
        aead = buildAead(context)
    }

    private fun buildAead(context: Context): Aead {
        return AndroidKeysetManager.Builder()
            .withKeyTemplate(KeyTemplates.get("AES256_GCM"))
            .withSharedPref(
                context,
                "$packageName.$keysetName",
                "$packageName.$prefFileName"
            )
            .withMasterKeyUri("android-keystore://tink_master_key")
            .build()
            .keysetHandle
            .getPrimitive(RegistryConfiguration.get(), Aead::class.java)
    }


    fun encrypt(inputByteArray: ByteArray): Result<String> {
        return runCatching {
            val encrypted = aead.encrypt(inputByteArray, null)
            Base64.getEncoder().encodeToString(encrypted)
        }
    }

    fun decrypt(inputEncryptedString: String): Result<ByteArray> {
        return runCatching {
            val encrypted = Base64.getDecoder().decode(inputEncryptedString)
            aead.decrypt(encrypted, null)
        }
    }
}

import androidx.benchmark.junit4.BenchmarkRule
import androidx.benchmark.junit4.measureRepeated
import androidx.test.ext.junit.runners.AndroidJUnit4
import androidx.test.platform.app.InstrumentationRegistry
import org.junit.Rule
import org.junit.Test
import org.junit.runner.RunWith

@RunWith(AndroidJUnit4::class)
class ExampleBenchmark {

    @get:Rule
    val benchmarkRule = BenchmarkRule()

    @Test
    fun benchmarkTinkEncrypt() {
        val context = InstrumentationRegistry.getInstrumentation().context
        val client = TinkClient(context)
        val plainText = ByteArray(1024 * 10)

        benchmarkRule.measureRepeated {
            client.encrypt(plainText).getOrThrow()
        }
    }

    @Test
    fun benchmarkStrongBoxEncrypt() {
        val context = InstrumentationRegistry.getInstrumentation().context
        val client = AndroidKeyStoreClient(context, true)
        val plainText = ByteArray(1024 * 10)

        benchmarkRule.measureRepeated {
            client.encrypt(plainText).getOrThrow()
        }
    }

    @Test
    fun benchmarkTeeEncrypt() {
        val context = InstrumentationRegistry.getInstrumentation().context
        val client = AndroidKeyStoreClient(context, false)
        val plainText = ByteArray(1024 * 10)

        benchmarkRule.measureRepeated {
            client.encrypt(plainText).getOrThrow()
        }
    }
}

val secretKey = generateSecretKey()
val kf = SecretKeyFactory.getInstance(KeyProperties.KEY_ALGORITHM_AES, KEY_STORE_PROVIDER)
val ki = kf.getKeySpec(secretKey, KeyInfo::class.java) as KeyInfo
val securityLevelString = when (ki.securityLevel) {
    KeyProperties.SECURITY_LEVEL_STRONGBOX -> "STRONGBOX"
    KeyProperties.SECURITY_LEVEL_TRUSTED_ENVIRONMENT -> "TEE"
    KeyProperties.SECURITY_LEVEL_SOFTWARE -> "SOFTWARE"
    else -> "UNKNOWN"
}

Log.d("KeyStoreSecurityLevel", "Security Level: ${ki.securityLevel}")