package no.zyon.shade.protocol

import no.zyon.shade.crypto.CryptoProvider
import no.zyon.shade.types.IdentityKeyPair
import no.zyon.shade.types.KeyPair
import no.zyon.shade.types.OneTimePreKey
import no.zyon.shade.types.PreKeyBundle
import no.zyon.shade.types.PreKeyMessage
import no.zyon.shade.types.SignedPreKey

/**
 * X3DH key agreement. Mirrors @shade/core/x3dh.ts.
 *
 * Identity keys: separate Ed25519 (signing) + X25519 (DH) keypairs stored together.
 */

/** Generate a new identity keypair (Ed25519 + X25519) */
fun generateIdentityKeyPair(crypto: CryptoProvider): IdentityKeyPair {
    val (signPub, signPriv) = crypto.generateEd25519KeyPair()
    val (dhPub, dhPriv) = crypto.generateX25519KeyPair()
    return IdentityKeyPair(
        signingPublicKey = signPub,
        signingPrivateKey = signPriv,
        dhPublicKey = dhPub,
        dhPrivateKey = dhPriv,
    )
}

/** Generate a signed prekey (X25519 keypair + Ed25519 signature over public key) */
fun generateSignedPreKey(
    crypto: CryptoProvider,
    identity: IdentityKeyPair,
    keyId: Int,
): SignedPreKey {
    val (pub, priv) = crypto.generateX25519KeyPair()
    val signature = crypto.sign(identity.signingPrivateKey, pub)
    return SignedPreKey(
        keyId = keyId,
        keyPair = KeyPair(publicKey = pub, privateKey = priv),
        signature = signature,
        timestamp = System.currentTimeMillis(),
    )
}

/** Generate a batch of one-time prekeys */
fun generateOneTimePreKeys(
    crypto: CryptoProvider,
    startId: Int,
    count: Int,
): List<OneTimePreKey> {
    val keys = mutableListOf<OneTimePreKey>()
    for (i in 0 until count) {
        val (pub, priv) = crypto.generateX25519KeyPair()
        keys.add(OneTimePreKey(keyId = startId + i, keyPair = KeyPair(pub, priv)))
    }
    return keys
}

fun createPreKeyBundle(
    registrationId: Int,
    identity: IdentityKeyPair,
    signedPreKey: SignedPreKey,
    oneTimePreKey: OneTimePreKey? = null,
): PreKeyBundle {
    return PreKeyBundle(
        registrationId = registrationId,
        identitySigningKey = identity.signingPublicKey,
        identityDHKey = identity.dhPublicKey,
        signedPreKey = PreKeyBundle.BundleSignedPreKey(
            keyId = signedPreKey.keyId,
            publicKey = signedPreKey.keyPair.publicKey,
            signature = signedPreKey.signature,
        ),
        oneTimePreKey = oneTimePreKey?.let {
            PreKeyBundle.BundleOneTimePreKey(it.keyId, it.keyPair.publicKey)
        },
    )
}

/** Result of processing a prekey bundle (Alice's side) */
data class X3DHInitResult(
    val rootKey: ByteArray,
    val ephemeralPublicKey: ByteArray,
    val signedPreKeyId: Int,
    val preKeyId: Int?,
    val remoteIdentityKey: ByteArray,
    val remoteSignedPreKey: ByteArray,
)

/**
 * Alice processes Bob's prekey bundle to establish a session.
 *
 * Steps:
 * 1. Verify the signed prekey signature
 * 2. Generate an ephemeral X25519 keypair
 * 3. Compute DH1 = DH(Alice identity DH, Bob signed prekey)
 * 4. Compute DH2 = DH(Alice ephemeral, Bob identity DH)
 * 5. Compute DH3 = DH(Alice ephemeral, Bob signed prekey)
 * 6. Compute DH4 = DH(Alice ephemeral, Bob one-time prekey) if available
 * 7. Derive initial root key from concatenated DH outputs
 */
fun processPreKeyBundle(
    crypto: CryptoProvider,
    identity: IdentityKeyPair,
    bundle: PreKeyBundle,
): X3DHInitResult {
    // 1. Verify signed prekey signature
    val valid = crypto.verify(
        bundle.identitySigningKey,
        bundle.signedPreKey.publicKey,
        bundle.signedPreKey.signature,
    )
    if (!valid) throw SecurityException("Signed prekey signature is invalid")

    // 2. Ephemeral keypair
    val (ephPub, ephPriv) = crypto.generateX25519KeyPair()

    // 3-6. DH computations
    val dh1 = crypto.x25519(identity.dhPrivateKey, bundle.signedPreKey.publicKey)
    val dh2 = crypto.x25519(ephPriv, bundle.identityDHKey)
    val dh3 = crypto.x25519(ephPriv, bundle.signedPreKey.publicKey)
    val secrets = mutableListOf(dh1, dh2, dh3)

    var preKeyId: Int? = null
    if (bundle.oneTimePreKey != null) {
        val dh4 = crypto.x25519(ephPriv, bundle.oneTimePreKey.publicKey)
        secrets.add(dh4)
        preKeyId = bundle.oneTimePreKey.keyId
    }

    // 7. Derive root key
    val rootKey = deriveInitialRootKey(crypto, secrets)

    return X3DHInitResult(
        rootKey = rootKey,
        ephemeralPublicKey = ephPub,
        signedPreKeyId = bundle.signedPreKey.keyId,
        preKeyId = preKeyId,
        remoteIdentityKey = bundle.identityDHKey,
        remoteSignedPreKey = bundle.signedPreKey.publicKey,
    )
}

/** Result of processing an incoming PreKeyMessage (Bob's side) */
data class X3DHResponseResult(
    val rootKey: ByteArray,
    val remoteIdentityKey: ByteArray,
    val remoteEphemeralKey: ByteArray,
)

/**
 * Bob processes an incoming PreKeyMessage to establish a session.
 * Mirrors Alice's DH computations from Bob's perspective.
 *
 * Caller is responsible for looking up the signed prekey and (if present)
 * the one-time prekey from storage.
 */
fun processPreKeyMessage(
    crypto: CryptoProvider,
    identity: IdentityKeyPair,
    signedPreKeyPrivate: ByteArray,
    oneTimePreKeyPrivate: ByteArray?,
    message: PreKeyMessage,
): X3DHResponseResult {
    val dh1 = crypto.x25519(signedPreKeyPrivate, message.identityDHKey)
    val dh2 = crypto.x25519(identity.dhPrivateKey, message.ephemeralKey)
    val dh3 = crypto.x25519(signedPreKeyPrivate, message.ephemeralKey)
    val secrets = mutableListOf(dh1, dh2, dh3)

    if (oneTimePreKeyPrivate != null) {
        val dh4 = crypto.x25519(oneTimePreKeyPrivate, message.ephemeralKey)
        secrets.add(dh4)
    }

    val rootKey = deriveInitialRootKey(crypto, secrets)
    return X3DHResponseResult(
        rootKey = rootKey,
        remoteIdentityKey = message.identityDHKey,
        remoteEphemeralKey = message.ephemeralKey,
    )
}
