package com.peachpasswords.passkeys

import com.peachpasswords.vault.VaultEntry
import java.math.BigInteger
import java.security.AlgorithmParameters
import java.security.KeyFactory
import java.security.MessageDigest
import java.security.Signature
import java.security.interfaces.ECPrivateKey
import java.security.spec.ECGenParameterSpec
import java.security.spec.ECParameterSpec
import java.security.spec.ECPrivateKeySpec
import java.util.Base64
import org.bouncycastle.crypto.params.Ed25519PrivateKeyParameters
import org.bouncycastle.crypto.signers.Ed25519Signer
import kotlinx.serialization.json.Json
import kotlinx.serialization.json.JsonObject
import kotlinx.serialization.json.JsonPrimitive
import kotlinx.serialization.json.buildJsonObject
import kotlinx.serialization.json.contentOrNull
import kotlinx.serialization.json.intOrNull
import kotlinx.serialization.json.jsonObject
import kotlinx.serialization.json.jsonPrimitive
import kotlinx.serialization.json.put

data class PasskeyAssertionResult(
    val responseJson: String,
    val nextSignCount: Int,
)

object PasskeyAssertion {
    private val json = Json { ignoreUnknownKeys = true }
    private val urlEncoder = Base64.getUrlEncoder().withoutPadding()
    private val urlDecoder = Base64.getUrlDecoder()

    fun build(
        entry: VaultEntry,
        requestJson: String,
        origin: String?,
        clientDataHashOverride: ByteArray? = null,
        // Set to true only when a biometric/device-credential prompt has actually
        // succeeded for this assertion. Defaults to false so existing callers and
        // tests that bypass the prompt still produce a valid flag byte (UV=0).
        userVerified: Boolean = false,
    ): PasskeyAssertionResult? {
        val passkey = entry.passkey ?: return null
        val privateKeyJwk = passkey.privateKeyJwk ?: entry.privateKeyJwk ?: return null
        val request = runCatching { json.parseToJsonElement(requestJson).jsonObject }.getOrNull() ?: return null
        val rpId = request["rpId"]?.jsonPrimitive?.contentOrNull ?: passkey.rpId
        if (rpId != passkey.rpId) return null
        val challenge = request["challenge"]?.jsonPrimitive?.contentOrNull ?: return null

        val nextSignCount = passkey.signCount + 1
        // UP is implicit (user tapped the credential), BE/BS reflect the
        // persisted record (immutable once written), and UV only when the caller
        // actually performed user verification. AT must NOT be set on assertions.
        val flags = assertionFlags(passkey.backupEligible, passkey.backupState, userVerified)
        val authData = buildAuthenticatorData(passkey.rpId, flags, nextSignCount)
        val clientDataJson = buildClientDataJson(
            challenge = challenge,
            origin = origin?.takeIf { it.isNotBlank() } ?: originFromRpId(passkey.rpId),
        )
        val clientDataHash = clientDataHashOverride ?: sha256(clientDataJson)
        val signature = signPasskey(privateKeyJwk, passkey.algorithm, authData + clientDataHash) ?: return null

        val responseJson = json.encodeToString(
            JsonObject.serializer(),
            buildJsonObject {
                put("id", passkey.credentialId)
                put("rawId", passkey.credentialId)
                put("type", "public-key")
                put("response", buildJsonObject {
                    put("clientDataJSON", b64(clientDataJson))
                    put("authenticatorData", b64(authData))
                    put("signature", b64(signature))
                    put("userHandle", passkey.userHandle)
                })
                put("authenticatorAttachment", "platform")
                put("clientExtensionResults", buildJsonObject {})
            },
        )
        return PasskeyAssertionResult(responseJson = responseJson, nextSignCount = nextSignCount)
    }

    private fun buildClientDataJson(challenge: String, origin: String): ByteArray =
        json.encodeToString(
            JsonObject.serializer(),
            buildJsonObject {
                put("type", "webauthn.get")
                put("challenge", challenge)
                put("origin", origin)
                put("crossOrigin", JsonPrimitive(false))
            },
        ).toByteArray(Charsets.UTF_8)

    private fun buildAuthenticatorData(rpId: String, flags: Int, signCount: Int): ByteArray {
        val out = ByteArray(37)
        sha256(rpId.toByteArray(Charsets.UTF_8)).copyInto(out, 0)
        out[32] = flags.toByte()
        out[33] = ((signCount ushr 24) and 0xff).toByte()
        out[34] = ((signCount ushr 16) and 0xff).toByte()
        out[35] = ((signCount ushr 8) and 0xff).toByte()
        out[36] = (signCount and 0xff).toByte()
        return out
    }

    // AT (0x40) is intentionally NOT included — assertions never carry
    // attested credential data; that's a registration-only field.
    private fun assertionFlags(backupEligible: Boolean, backupState: Boolean, userVerified: Boolean): Int {
        var f = FLAG_UP
        if (userVerified) f = f or FLAG_UV
        if (backupEligible) f = f or FLAG_BE
        if (backupState) f = f or FLAG_BS
        return f
    }

    private const val FLAG_UP: Int = 0x01
    private const val FLAG_UV: Int = 0x04
    private const val FLAG_BE: Int = 0x08
    private const val FLAG_BS: Int = 0x10

    private fun signPasskey(privateKeyJwk: String, algorithm: Int, message: ByteArray): ByteArray? =
        when (algorithm) {
            -7 -> signEs256(privateKeyJwk, message)
            -8 -> signEd25519(privateKeyJwk, message)
            else -> null
        }

    private fun signEs256(privateKeyJwk: String, message: ByteArray): ByteArray? {
        val jwk = runCatching { json.parseToJsonElement(privateKeyJwk).jsonObject }.getOrNull() ?: return null
        if (jwk["kty"]?.jsonPrimitive?.contentOrNull != "EC") return null
        if (jwk["crv"]?.jsonPrimitive?.contentOrNull != "P-256") return null
        val d = jwk["d"]?.jsonPrimitive?.contentOrNull?.let(::b64Decode) ?: return null
        val privateKey = ecPrivateKey(d)
        val signer = Signature.getInstance("SHA256withECDSA")
        signer.initSign(privateKey)
        signer.update(message)
        return signer.sign()
    }

    private fun signEd25519(privateKeyJwk: String, message: ByteArray): ByteArray? {
        val jwk = runCatching { json.parseToJsonElement(privateKeyJwk).jsonObject }.getOrNull() ?: return null
        if (jwk["kty"]?.jsonPrimitive?.contentOrNull != "OKP") return null
        if (jwk["crv"]?.jsonPrimitive?.contentOrNull != "Ed25519") return null
        val d = jwk["d"]?.jsonPrimitive?.contentOrNull?.let(::b64Decode) ?: return null
        if (d.size != 32) return null
        val signer = Ed25519Signer()
        signer.init(true, Ed25519PrivateKeyParameters(d, 0))
        signer.update(message, 0, message.size)
        return signer.generateSignature()
    }

    private fun ecPrivateKey(d: ByteArray): ECPrivateKey {
        val keyFactory = KeyFactory.getInstance("EC")
        val spec = ECPrivateKeySpec(BigInteger(1, d), ecP256Spec())
        return keyFactory.generatePrivate(spec) as ECPrivateKey
    }

    private fun ecP256Spec(): ECParameterSpec {
        val params = AlgorithmParameters.getInstance("EC")
        params.init(ECGenParameterSpec("secp256r1"))
        return params.getParameterSpec(ECParameterSpec::class.java)
    }

    private fun originFromRpId(rpId: String): String =
        if (rpId.equals("localhost", ignoreCase = true) || rpId == "127.0.0.1" || rpId == "[::1]") {
            "http://$rpId"
        } else {
            "https://$rpId"
        }

    private fun sha256(data: ByteArray): ByteArray = MessageDigest.getInstance("SHA-256").digest(data)
    private fun b64(data: ByteArray): String = urlEncoder.encodeToString(data)
    private fun b64Decode(value: String): ByteArray {
        val padded = value.padEnd(((value.length + 3) / 4) * 4, '=')
        return urlDecoder.decode(padded)
    }
}

/**
 * H-07: Decide whether to sign a WebAuthn assertion given the relying party's
 * user-verification requirement and the actual ceremony outcome.
 *
 * If the RP requires UV and the authenticator could not verify the user, the
 * assertion must NOT be signed — WebAuthn requires UV=1 only when verification
 * actually happened, and an RP that requires UV will reject UV=0. Fail closed
 * rather than produce a misleading assertion.
 *
 * For "preferred", "discouraged", or absent UV requirements, the assertion is
 * always signed with the real UV bit reflecting what actually happened.
 */
internal fun shouldSignAssertion(uvRequired: Boolean, userVerified: Boolean): Boolean =
    !uvRequired || userVerified