package com.peachpasswords.passkeys

import com.peachpasswords.vault.PasskeyData
import com.peachpasswords.vault.VaultEntry
import java.math.BigInteger
import java.security.MessageDigest
import java.security.KeyPairGenerator
import java.security.Signature
import java.security.interfaces.ECPrivateKey
import java.security.interfaces.ECPublicKey
import java.security.spec.ECGenParameterSpec
import java.util.Base64
import kotlinx.serialization.json.Json
import kotlinx.serialization.json.contentOrNull
import kotlinx.serialization.json.jsonObject
import kotlinx.serialization.json.jsonPrimitive
import org.junit.Assert.assertEquals
import org.junit.Assert.assertNotNull
import org.junit.Assert.assertTrue
import org.junit.Test

class PasskeyAssertionTest {
    @Test fun builds_es256_assertion_and_increments_counter() {
        val keyPair = KeyPairGenerator.getInstance("EC").apply {
            initialize(ECGenParameterSpec("secp256r1"))
        }.generateKeyPair()
        val entry = passkeyEntry(keyPair.private as ECPrivateKey)

        val result = PasskeyAssertion.build(
            entry = entry,
            requestJson = """{"challenge":"Y2hhbGxlbmdl","rpId":"example.com"}""",
            origin = "https://login.example.com",
        )

        assertNotNull(result)
        result!!
        assertEquals(42, result.nextSignCount)
        val response = Json.parseToJsonElement(result.responseJson).jsonObject
        assertEquals("cred-1", response["id"]?.jsonPrimitive?.contentOrNull)
        val payload = response["response"]!!.jsonObject
        val clientDataBytes = b64Decode(payload["clientDataJSON"]!!.jsonPrimitive.content)
        val clientData = String(clientDataBytes, Charsets.UTF_8)
        val client = Json.parseToJsonElement(clientData).jsonObject
        assertEquals("webauthn.get", client["type"]?.jsonPrimitive?.contentOrNull)
        assertEquals("https://login.example.com", client["origin"]?.jsonPrimitive?.contentOrNull)
        val authData = b64Decode(payload["authenticatorData"]!!.jsonPrimitive.content)
        assertEquals(37, authData.size)
        assertEquals(42, ((authData[33].toInt() and 0xff) shl 24) or ((authData[34].toInt() and 0xff) shl 16) or ((authData[35].toInt() and 0xff) shl 8) or (authData[36].toInt() and 0xff))

        val verifier = Signature.getInstance("SHA256withECDSA")
        verifier.initVerify(keyPair.public as ECPublicKey)
        verifier.update(authData + MessageDigest.getInstance("SHA-256").digest(clientDataBytes))
        assertTrue(verifier.verify(b64Decode(payload["signature"]!!.jsonPrimitive.content)))
    }

    @Test fun returns_null_for_non_es256_or_missing_private_key() {
        val entry = passkeyEntry(null)
        assertEquals(null, PasskeyAssertion.build(entry, """{"challenge":"c","rpId":"example.com"}""", "https://example.com"))
        val eddsa = entry.copy(passkey = entry.passkey!!.copy(privateKeyJwk = """{"kty":"OKP","crv":"Ed25519","d":"x","x":"y"}""", algorithm = -8))
        assertEquals(null, PasskeyAssertion.build(eddsa, """{"challenge":"c","rpId":"example.com"}""", "https://example.com"))
    }

    private fun passkeyEntry(privateKey: ECPrivateKey?) = VaultEntry(
        id = "entry-1",
        type = "passkey",
        title = "Example",
        passkey = PasskeyData(
            credentialId = "cred-1",
            publicKeyCose = "cose",
            privateKeyJwk = privateKey?.let(::privateJwk),
            rpId = "example.com",
            userHandle = "user-handle",
            userName = "june",
            userDisplayName = "June",
            algorithm = -7,
            signCount = 41,
        ),
        createdAt = 1,
        modifiedAt = 1,
    )

    private fun privateJwk(key: ECPrivateKey): String =
        """{"kty":"EC","crv":"P-256","d":"${b64(key.s.toFixed32())}"}"""

    private fun BigInteger.toFixed32(): ByteArray {
        val raw = toByteArray()
        val unsigned = if (raw.size > 1 && raw[0] == 0.toByte()) raw.copyOfRange(1, raw.size) else raw
        return ByteArray(32).also { unsigned.copyInto(it, 32 - unsigned.size) }
    }

    private fun b64(bytes: ByteArray): String = Base64.getUrlEncoder().withoutPadding().encodeToString(bytes)
    private fun b64Decode(value: String): ByteArray =
        Base64.getUrlDecoder().decode(value.padEnd(((value.length + 3) / 4) * 4, '='))
}