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.assertFalse
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 browser_client_data_hash_does_not_emit_conflicting_client_data_json() {
        val keyPair = KeyPairGenerator.getInstance("EC").apply {
            initialize(ECGenParameterSpec("secp256r1"))
        }.generateKeyPair()
        val clientDataHash = MessageDigest.getInstance("SHA-256")
            .digest("framework-owned-client-data".toByteArray())

        val result = PasskeyAssertion.build(
            entry = passkeyEntry(keyPair.private as ECPrivateKey),
            requestJson = """{"challenge":"Y2hhbGxlbmdl","rpId":"example.com"}""",
            origin = "https://example.com",
            clientDataHashOverride = clientDataHash,
        )

        assertNotNull(result)
        val payload = Json.parseToJsonElement(result!!.responseJson).jsonObject["response"]!!.jsonObject
        assertFalse(payload.containsKey("clientDataJSON"))
        val authData = b64Decode(payload["authenticatorData"]!!.jsonPrimitive.content)
        val verifier = Signature.getInstance("SHA256withECDSA")
        verifier.initVerify(keyPair.public as ECPublicKey)
        verifier.update(authData + clientDataHash)
        assertTrue(verifier.verify(b64Decode(payload["signature"]!!.jsonPrimitive.content)))
    }

    @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 rejects_missing_or_blank_origin() {
        val keyPair = KeyPairGenerator.getInstance("EC").apply {
            initialize(ECGenParameterSpec("secp256r1"))
        }.generateKeyPair()
        val entry = passkeyEntry(keyPair.private as ECPrivateKey)
        assertEquals(null, PasskeyAssertion.build(entry, """{"challenge":"c","rpId":"example.com"}""", null))
        assertEquals(null, PasskeyAssertion.build(entry, """{"challenge":"c","rpId":"example.com"}""", " "))
        assertEquals(
            PasskeyAssertionFailure.MISSING_ORIGIN,
            (PasskeyAssertion.buildDetailed(
                entry,
                """{"challenge":"c","rpId":"example.com"}""",
                null,
            ) as PasskeyAssertionBuild.Failure).reason,
        )
    }

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

        val outcome = PasskeyAssertion.buildDetailed(
            entry,
            """{"challenge":"c","rpId":"example.com"}""",
            "https://unrelated.test",
        )

        assertEquals(
            PasskeyAssertionFailure.CALLER_NOT_ALLOWED,
            (outcome as PasskeyAssertionBuild.Failure).reason,
        )
    }

    @Test fun bound_credential_refuses_every_other_caller() {
        val keyPair = KeyPairGenerator.getInstance("EC").apply {
            initialize(ECGenParameterSpec("secp256r1"))
        }.generateKeyPair()
        val appA = "android:apk-key-hash:AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA"
        val appB = "android:apk-key-hash:BBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBB"
        val bound = passkeyEntry(keyPair.private as ECPrivateKey).copy(
            passkey = passkeyEntry(keyPair.private as ECPrivateKey).passkey!!.copy(callerOrigin = appA),
        )
        val request = """{"challenge":"c","rpId":"example.com"}"""
        assertNotNull(PasskeyAssertion.build(bound, request, appA))
        assertEquals(null, PasskeyAssertion.build(bound, request, appB))
        // A web caller cannot use a credential bound to a native app either.
        assertEquals(null, PasskeyAssertion.build(bound, request, "https://example.com"))
    }

    @Test fun legacy_credential_stays_hidden_from_native_app_callers() {
        val keyPair = KeyPairGenerator.getInstance("EC").apply {
            initialize(ECGenParameterSpec("secp256r1"))
        }.generateKeyPair()
        val legacy = passkeyEntry(keyPair.private as ECPrivateKey)
        assertEquals(
            null,
            PasskeyAssertion.build(
                legacy,
                """{"challenge":"c","rpId":"example.com"}""",
                "android:apk-key-hash:AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA",
            ),
        )
        // ...but a verified web caller inside the rpId domain still works.
        assertNotNull(PasskeyAssertion.build(legacy, """{"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, '='))
}