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 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
        if (passkey.algorithm != -7) 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 = signEs256(privateKeyJwk, 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 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 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)
    }
}