package app.closer.crypto

import com.google.crypto.tink.aead.AeadConfig
import com.google.crypto.tink.hybrid.HybridConfig
import com.google.crypto.tink.hybrid.HybridKeyTemplates
import com.google.crypto.tink.KeysetHandle
import org.junit.Assert.assertEquals
import org.junit.Assert.assertTrue
import org.junit.BeforeClass
import org.junit.Test
import java.security.GeneralSecurityException

class ReleaseKeyEncryptorTest {

    companion object {
        @BeforeClass
        @JvmStatic
        fun setup() {
            AeadConfig.register()
            HybridConfig.register()
        }

        private fun eciesKeypair(): KeysetHandle =
            KeysetHandle.generateNew(HybridKeyTemplates.ECIES_P256_HKDF_HMAC_SHA256_AES128_GCM)
    }

    private val sealedEncryptor = SealedAnswerEncryptor()
    private val subject = ReleaseKeyEncryptor(sealedEncryptor)

    private val coupleId = "couple-abc"
    private val questionId = "q-001"
    private val aliceId = "user-alice"
    private val bobId = "user-bob"

    @Test
    fun `wrapForRecipient produces keybox prefix`() {
        val bobKeypair = eciesKeypair()
        val bobPublicKeyB64 = UserKeyManager.publicKeyB64Companion(bobKeypair)
        val oneTimeKey = sealedEncryptor.generateOneTimeKey()

        val keybox = subject.wrapForRecipient(oneTimeKey, bobPublicKeyB64, coupleId, questionId, aliceId, bobId)

        assertTrue(keybox.startsWith(ReleaseKeyEncryptor.KEYBOX_PREFIX))
    }

    @Test
    fun `wrap and unwrap round-trip restores the answer key`() {
        val aliceKeypair = eciesKeypair()
        val bobKeypair = eciesKeypair()

        val aliceOneTimeKey = sealedEncryptor.generateOneTimeKey()
        val bobPublicKeyB64 = UserKeyManager.publicKeyB64Companion(bobKeypair)

        val keybox = subject.wrapForRecipient(
            aliceOneTimeKey, bobPublicKeyB64, coupleId, questionId, aliceId, bobId
        )

        val recoveredKey = subject.unwrapFromSender(keybox, bobKeypair, coupleId, questionId, aliceId, bobId)

        val payload = SealedAnswerEncryptor.AnswerPayload("weekend away", listOf("opt-1"), 9)
        val sealed = sealedEncryptor.seal(payload, aliceOneTimeKey, coupleId, questionId, aliceId)
        val decrypted = sealedEncryptor.open(sealed, recoveredKey, coupleId, questionId, aliceId)

        assertEquals(payload.writtenText, decrypted.writtenText)
        assertEquals(payload.scaleValue, decrypted.scaleValue)
    }

    @Test(expected = GeneralSecurityException::class)
    fun `wrong recipient private key cannot unwrap`() {
        val bobKeypair = eciesKeypair()
        val eveKeypair = eciesKeypair()

        val oneTimeKey = sealedEncryptor.generateOneTimeKey()
        val bobPublicKeyB64 = UserKeyManager.publicKeyB64Companion(bobKeypair)
        val keybox = subject.wrapForRecipient(oneTimeKey, bobPublicKeyB64, coupleId, questionId, aliceId, bobId)

        // Eve tries to unwrap Bob's keybox using her own private key.
        subject.unwrapFromSender(keybox, eveKeypair, coupleId, questionId, aliceId, bobId)
    }

    @Test(expected = GeneralSecurityException::class)
    fun `wrong context info cannot unwrap`() {
        val bobKeypair = eciesKeypair()
        val oneTimeKey = sealedEncryptor.generateOneTimeKey()
        val bobPublicKeyB64 = UserKeyManager.publicKeyB64Companion(bobKeypair)
        val keybox = subject.wrapForRecipient(oneTimeKey, bobPublicKeyB64, coupleId, questionId, aliceId, bobId)

        // Wrong questionId in context info.
        subject.unwrapFromSender(keybox, bobKeypair, coupleId, "other-question", aliceId, bobId)
    }
}
