package swiss.qpq.gajumobile.security

import android.content.Context
import kotlinx.serialization.Serializable
import java.security.SecureRandom
import javax.crypto.Cipher
import javax.crypto.KeyGenerator
import javax.crypto.SecretKey
import javax.crypto.spec.GCMParameterSpec
import javax.crypto.spec.SecretKeySpec

@Serializable
data class EncryptedEnvelope(
    val encryptedData: ByteArray,
    val encryptedDek: ByteArray,
    val dataIv: ByteArray,
    val dekIv: ByteArray,
) {
    override fun equals(other: Any?): Boolean {
        if (this === other) return true
        if (javaClass != other?.javaClass) return false

        other as EncryptedEnvelope

        if (!encryptedData.contentEquals(other.encryptedData)) return false
        if (!encryptedDek.contentEquals(other.encryptedDek)) return false
        if (!dataIv.contentEquals(other.dataIv)) return false
        return dekIv.contentEquals(other.dekIv)
    }

    override fun hashCode(): Int {
        var result = encryptedData.contentHashCode()
        result = (31 * result) + encryptedDek.contentHashCode()
        result = (31 * result) + dataIv.contentHashCode()
        result = (31 * result) + dekIv.contentHashCode()
        return result
    }
}

object EncryptionService {
    private const val AES_GCM = "AES/GCM/NoPadding"
    private const val GCM_TAG_LENGTH = 128
    private const val IV_LENGTH = 12

    fun encrypt(context: Context, data: ByteArray): EncryptedEnvelope {
        // 1. Generate a random Data Encryption Key (DEK) in software
        val dek = generateDek()

        // 2. Encrypt the data with the DEK
        val dataCipher = Cipher.getInstance(AES_GCM)
        val dataIv = ByteArray(IV_LENGTH).apply { SecureRandom().nextBytes(this) }
        dataCipher.init(Cipher.ENCRYPT_MODE, dek, GCMParameterSpec(GCM_TAG_LENGTH, dataIv))
        val encryptedData = dataCipher.doFinal(data)

        // 3. Encrypt the DEK with the Master Key (Hardware-backed)
        val masterKey = KeyManager.getMasterKey(context)
        val dekCipher = Cipher.getInstance(AES_GCM)
        // Let Keystore generate the IV for the Master Key encryption
        dekCipher.init(Cipher.ENCRYPT_MODE, masterKey)
        val encryptedDek = dekCipher.doFinal(dek.encoded)
        val dekIv = dekCipher.iv

        return EncryptedEnvelope(
            encryptedData = encryptedData,
            encryptedDek = encryptedDek,
            dataIv = dataIv,
            dekIv = dekIv
        )
    }

    fun decrypt(context: Context, envelope: EncryptedEnvelope): ByteArray {
        // 1. Decrypt the DEK using the Master Key
        val masterKey = KeyManager.getMasterKey(context)
        val dekCipher = Cipher.getInstance(AES_GCM)
        dekCipher.init(Cipher.DECRYPT_MODE, masterKey, GCMParameterSpec(GCM_TAG_LENGTH, envelope.dekIv))
        val decryptedDekBytes = dekCipher.doFinal(envelope.encryptedDek)
        val dek = SecretKeySpec(decryptedDekBytes, "AES")

        // 2. Decrypt the data using the DEK
        val dataCipher = Cipher.getInstance(AES_GCM)
        dataCipher.init(Cipher.DECRYPT_MODE, dek, GCMParameterSpec(GCM_TAG_LENGTH, envelope.dataIv))
        return dataCipher.doFinal(envelope.encryptedData)
    }

    private fun generateDek(): SecretKey {
        val keyGen = KeyGenerator.getInstance("AES")
        keyGen.init(256)
        return keyGen.generateKey()
    }
}