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
import org.bouncycastle.crypto.params.Ed25519PrivateKeyParameters
import org.bouncycastle.crypto.params.Ed25519PublicKeyParameters
import org.bouncycastle.crypto.signers.Ed25519Signer

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 builds_ed25519_assertion() {
        val seed = ByteArray(32) { (it + 1).toByte() }
        val privateKey = Ed25519PrivateKeyParameters(seed, 0)
        val publicKey = privateKey.generatePublicKey().encoded
        val entry = passkeyEntry(null).copy(
            passkey = passkeyEntry(null).passkey!!.copy(
                privateKeyJwk = """{"kty":"OKP","crv":"Ed25519","d":"${b64(seed)}","x":"${b64(publicKey)}"}""",
                algorithm = -8,
            ),
        )

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

        assertNotNull(result)
        result!!
        val payload = Json.parseToJsonElement(result.responseJson).jsonObject["response"]!!.jsonObject
        val authData = b64Decode(payload["authenticatorData"]!!.jsonPrimitive.content)
        val clientData = b64Decode(payload["clientDataJSON"]!!.jsonPrimitive.content)
        val message = authData + MessageDigest.getInstance("SHA-256").digest(clientData)
        val signature = b64Decode(payload["signature"]!!.jsonPrimitive.content)
        assertEquals(64, signature.size)
        val verifier = Ed25519Signer()
        verifier.init(false, Ed25519PublicKeyParameters(publicKey, 0))
        verifier.update(message, 0, message.size)
        assertTrue(verifier.verifySignature(signature))
    }

    @Test fun returns_null_for_missing_private_key() {
        val entry = passkeyEntry(null)
        assertEquals(null, PasskeyAssertion.build(entry, """{"challenge":"c","rpId":"example.com"}""", "https://example.com"))
    }

    @Test fun assertion_flag_byte_reflects_backup_and_uv_state() {
        // Backup-eligible + UV performed → UP | UV | BE = 0x0d. AT (0x40) must
        // NOT be set on an assertion (it's registration-only). BS depends on
        // passkey.backupState which is false here.
        val keyPair = KeyPairGenerator.getInstance("EC").apply {
            initialize(ECGenParameterSpec("secp256r1"))
        }.generateKeyPair()
        val privKey = keyPair.private as ECPrivateKey
        val backupEntry = passkeyEntry(privKey).copy(
            passkey = passkeyEntry(privKey).passkey!!.copy(
                backupEligible = true,
                backupState = false,
            ),
        )

        val verified = PasskeyAssertion.build(
            entry = backupEntry,
            requestJson = """{"challenge":"Y2hhbGxlbmdl","rpId":"example.com"}""",
            origin = "https://example.com",
            userVerified = true,
        )
        assertNotNull(verified)
        assertEquals(0x0d, extractFlags(verified!!.responseJson))

        // Same entry, no UV performed: UP | BE only = 0x09.
        val unverified = PasskeyAssertion.build(
            entry = backupEntry,
            requestJson = """{"challenge":"Y2hhbGxlbmdl","rpId":"example.com"}""",
            origin = "https://example.com",
            userVerified = false,
        )
        assertEquals(0x09, extractFlags(unverified!!.responseJson))

        // No backup eligibility, no UV: just UP = 0x01.
        val plainEntry = passkeyEntry(privKey).copy(
            passkey = passkeyEntry(privKey).passkey!!.copy(
                backupEligible = false,
                backupState = false,
            ),
        )
        val plain = PasskeyAssertion.build(
            entry = plainEntry,
            requestJson = """{"challenge":"Y2hhbGxlbmdl","rpId":"example.com"}""",
            origin = "https://example.com",
            userVerified = false,
        )
        assertEquals(0x01, extractFlags(plain!!.responseJson))
        // AT (0x40) is never set on assertions regardless of the other bits.
        assertEquals(0, 0x40 and extractFlags(plain.responseJson))
    }

    private fun extractFlags(responseJson: String): Int {
        val payload = Json.parseToJsonElement(responseJson).jsonObject["response"]!!.jsonObject
        val authData = b64Decode(payload["authenticatorData"]!!.jsonPrimitive.content)
        return authData[32].toInt() and 0xff
    }

    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, '='))
}