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,
    ): 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
        val flags = 0x05 or (if (passkey.backupEligible) 0x08 else 0) or (if (passkey.backupState) 0x10 else 0)
        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
    }

    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)
    }
}