package social.grain.support

import android.security.keystore.KeyGenParameterSpec
import java.io.InputStream
import java.io.OutputStream
import java.security.Key
import java.security.KeyStore
import java.security.KeyStoreSpi
import java.security.Provider
import java.security.SecureRandom
import java.security.Security
import java.security.cert.Certificate
import java.security.spec.AlgorithmParameterSpec
import java.util.Collections
import java.util.Date
import java.util.Enumeration
import javax.crypto.KeyGenerator
import javax.crypto.KeyGeneratorSpi
import javax.crypto.SecretKey
import javax.crypto.spec.SecretKeySpec

/**
 * A stand-in for `AndroidKeyStore`, which no JVM sandbox provides.
 *
 * [social.grain.data.auth.SecureStore] seals every credential under a
 * non-exportable Keystore key, and that provider simply doesn't exist off a
 * device — so without this, anything that touches a session can only be tested
 * on an emulator. The keys here are ordinary in-memory AES keys: the point is
 * to exercise the encrypt/decrypt/rotate paths that sit *above* the Keystore,
 * not to reproduce hardware backing, which a test could not verify anyway.
 *
 * Held for the life of the JVM, matching the device: a second `SecureStore`
 * built in the same test finds the master key the first one created, and a
 * value written by one is readable by the other.
 */
object FakeAndroidKeyStore {

    internal val keys = Collections.synchronizedMap(mutableMapOf<String, SecretKey>())

    fun install() {
        if (Security.getProvider(NAME) == null) Security.addProvider(FakeKeyStoreProvider())
    }

    /** Forget every key, as clearing the app's data would. */
    fun reset() = keys.clear()

    internal const val NAME = "AndroidKeyStore"
}

/** The JCA provider entry points `SecureStore` asks for by name. */
class FakeKeyStoreProvider : Provider(FakeAndroidKeyStore.NAME, 1.0, "In-memory test keystore") {
    init {
        put("KeyStore.${FakeAndroidKeyStore.NAME}", FakeKeyStoreSpi::class.java.name)
        put("KeyGenerator.AES", FakeKeyGeneratorSpi::class.java.name)
    }
}

/** Just enough of a keystore to hold secret keys under an alias. */
class FakeKeyStoreSpi : KeyStoreSpi() {

    override fun engineGetKey(alias: String, password: CharArray?): Key? =
        FakeAndroidKeyStore.keys[alias]

    override fun engineGetEntry(
        alias: String,
        protParam: KeyStore.ProtectionParameter?,
    ): KeyStore.Entry? =
        FakeAndroidKeyStore.keys[alias]?.let { KeyStore.SecretKeyEntry(it) }

    override fun engineSetEntry(
        alias: String,
        entry: KeyStore.Entry,
        protParam: KeyStore.ProtectionParameter?,
    ) {
        (entry as? KeyStore.SecretKeyEntry)?.let { FakeAndroidKeyStore.keys[alias] = it.secretKey }
    }

    override fun engineSetKeyEntry(
        alias: String,
        key: Key,
        password: CharArray?,
        chain: Array<out Certificate>?,
    ) {
        (key as? SecretKey)?.let { FakeAndroidKeyStore.keys[alias] = it }
    }

    override fun engineSetKeyEntry(alias: String, key: ByteArray?, chain: Array<out Certificate>?) =
        throw UnsupportedOperationException()

    override fun engineDeleteEntry(alias: String) {
        FakeAndroidKeyStore.keys.remove(alias)
    }

    override fun engineContainsAlias(alias: String): Boolean =
        FakeAndroidKeyStore.keys.containsKey(alias)

    override fun engineAliases(): Enumeration<String> =
        Collections.enumeration(FakeAndroidKeyStore.keys.keys.toList())

    override fun engineSize(): Int = FakeAndroidKeyStore.keys.size

    override fun engineIsKeyEntry(alias: String): Boolean = engineContainsAlias(alias)

    override fun engineIsCertificateEntry(alias: String): Boolean = false

    override fun engineGetCertificate(alias: String): Certificate? = null

    override fun engineGetCertificateChain(alias: String): Array<Certificate>? = null

    override fun engineGetCertificateAlias(cert: Certificate?): String? = null

    override fun engineSetCertificateEntry(alias: String, cert: Certificate?) =
        throw UnsupportedOperationException()

    override fun engineGetCreationDate(alias: String): Date = Date(0)

    override fun engineLoad(stream: InputStream?, password: CharArray?) = Unit

    override fun engineStore(stream: OutputStream?, password: CharArray?) = Unit
}

/**
 * Generates an ordinary AES key and files it under the alias the
 * [KeyGenParameterSpec] names — which is what the real provider does, and what
 * makes a key created once findable on the next launch.
 */
class FakeKeyGeneratorSpi : KeyGeneratorSpi() {

    private var alias: String? = null
    private var sizeBits = 256

    override fun engineInit(random: SecureRandom?) = Unit

    override fun engineInit(params: AlgorithmParameterSpec?, random: SecureRandom?) {
        val spec = params as? KeyGenParameterSpec ?: return
        alias = spec.keystoreAlias
        sizeBits = spec.keySize.takeIf { it > 0 } ?: 256
    }

    override fun engineInit(keysize: Int, random: SecureRandom?) {
        sizeBits = keysize
    }

    override fun engineGenerateKey(): SecretKey {
        val key = KeyGenerator.getInstance("AES").apply { init(sizeBits) }.generateKey()
            .let { SecretKeySpec(it.encoded, "AES") }
        alias?.let { FakeAndroidKeyStore.keys[it] = key }
        return key
    }
}