package com.peachpasswords.sync

import org.json.JSONObject
import java.security.MessageDigest
import java.security.SecureRandom
import javax.crypto.Mac
import javax.crypto.spec.SecretKeySpec

private val PAIRING_AUTH_DOMAIN = "peach-sync-v1-initiator-auth".encodeToByteArray()

/**
 * E1 post-NK proof:
 * HMAC-SHA256(X25519(initiatorPriv, responderPub),
 *   handshakeHash || domain || initiatorPub || responderPub || nonce)
 */
internal fun computeInitiatorAuthProof(
    initiatorPrivateKey: ByteArray,
    initiatorPublicKey: ByteArray,
    responderPublicKey: ByteArray,
    nonce: ByteArray,
    handshakeHash: ByteArray,
): ByteArray {
    require(initiatorPrivateKey.size == 32) { "initiator private key must be 32 bytes" }
    require(initiatorPublicKey.size == 32) { "initiator public key must be 32 bytes" }
    require(responderPublicKey.size == 32) { "responder public key must be 32 bytes" }
    require(nonce.size == 32) { "auth nonce must be 32 bytes" }
    require(handshakeHash.size == 32) { "handshake hash must be 32 bytes" }

    return computeAuthProofForPeer(
        localPrivateKey = initiatorPrivateKey,
        remotePublicKey = responderPublicKey,
        initiatorPublicKey = initiatorPublicKey,
        responderPublicKey = responderPublicKey,
        nonce = nonce,
        handshakeHash = handshakeHash,
    )
}

private fun computeAuthProofForPeer(
    localPrivateKey: ByteArray,
    remotePublicKey: ByteArray,
    initiatorPublicKey: ByteArray,
    responderPublicKey: ByteArray,
    nonce: ByteArray,
    handshakeHash: ByteArray,
): ByteArray {
    val dh = x25519SharedKey(localPrivateKey, remotePublicKey)
    return Mac.getInstance("HmacSHA256").run {
        init(SecretKeySpec(dh, "HmacSHA256"))
        update(handshakeHash)
        update(PAIRING_AUTH_DOMAIN)
        update(initiatorPublicKey)
        update(responderPublicKey)
        update(nonce)
        doFinal()
    }
}

/**
 * Browser/Android initiator half, useful for native-to-native first pairing.
 * The challenge's responder key must match the key already authenticated by NK.
 */
internal suspend fun performInitiatorPairingAuth(
    session: NoiseNKSession,
    ourPrivateKey: ByteArray,
    ourPublicKey: ByteArray,
    responderPublicKey: ByteArray,
    sendEncryptedJson: suspend (ByteArray) -> Unit,
    receiveEncryptedJson: suspend () -> ByteArray,
) {
    val challenge = JSONObject(String(receiveEncryptedJson(), Charsets.UTF_8))
    require(challenge.optString("type") == "CHALLENGE") { "expected CHALLENGE frame" }
    val nonce = hexToPubkey(challenge.getString("nonce_hex"))
    val challengedResponder = hexToPubkey(challenge.getString("responder_public_key_hex"))
    require(MessageDigest.isEqual(challengedResponder, responderPublicKey)) {
        "responder key in challenge does not match Noise handshake"
    }
    val proof = computeInitiatorAuthProof(
        ourPrivateKey,
        ourPublicKey,
        responderPublicKey,
        nonce,
        session.handshakeHash,
    )
    sendEncryptedJson(
        JSONObject()
            .put("type", "PROOF")
            .put("initiator_public_key_hex", pubkeyToHex(ourPublicKey))
            .put("proof_hex", pubkeyToHex(proof))
            .toString()
            .toByteArray(Charsets.UTF_8),
    )
}

/**
 * Android responder half. Returns the initiator key whose private half was
 * proven, so the caller can bind it to the following PAIR_IDENTITY frame.
 */
internal suspend fun performResponderPairingAuth(
    session: NoiseNKSession,
    ourPrivateKey: ByteArray,
    ourPublicKey: ByteArray,
    sendEncryptedJson: suspend (ByteArray) -> Unit,
    receiveEncryptedJson: suspend () -> ByteArray,
    random: SecureRandom = SecureRandom(),
): ByteArray {
    val nonce = ByteArray(32).also(random::nextBytes)
    sendEncryptedJson(
        JSONObject()
            .put("type", "CHALLENGE")
            .put("nonce_hex", pubkeyToHex(nonce))
            .put("responder_public_key_hex", pubkeyToHex(ourPublicKey))
            .toString()
            .toByteArray(Charsets.UTF_8),
    )

    val frame = JSONObject(String(receiveEncryptedJson(), Charsets.UTF_8))
    require(frame.optString("type") == "PROOF") { "expected PROOF frame" }
    val initiatorPublicKey = hexToPubkey(frame.getString("initiator_public_key_hex"))
    val claimedProof = hexToPubkey(frame.getString("proof_hex"))
    val expectedProof = computeAuthProofForPeer(
        localPrivateKey = ourPrivateKey,
        remotePublicKey = initiatorPublicKey,
        initiatorPublicKey = initiatorPublicKey,
        responderPublicKey = ourPublicKey,
        nonce = nonce,
        handshakeHash = session.handshakeHash,
    )
    require(MessageDigest.isEqual(expectedProof, claimedProof)) { "auth proof did not verify" }
    return initiatorPublicKey
}