package com.peachpasswords.sync

/**
 * Thin JNI boundary over the canonical Rust `snow` state machine.
 *
 * No transcript, key schedule, cipher, or nonce logic belongs in Kotlin.
 */
internal object RustNoiseNative {
    init {
        System.loadLibrary("peach_noise_core")
    }

    external fun nativeCreateNkInitiator(remotePublicKey: ByteArray): Long
    external fun nativeCreateNkResponder(
        localPrivateKey: ByteArray,
        localPublicKey: ByteArray,
    ): Long
    external fun nativeCreateNkInitiatorForTest(
        remotePublicKey: ByteArray,
        ephemeralPrivateKey: ByteArray,
    ): Long
    external fun nativeCreateNkResponderForTest(
        localPrivateKey: ByteArray,
        localPublicKey: ByteArray,
        ephemeralPrivateKey: ByteArray,
    ): Long

    external fun nativeCreateKkInitiator(
        localPrivateKey: ByteArray,
        localPublicKey: ByteArray,
        remotePublicKey: ByteArray,
    ): Long
    external fun nativeCreateKkResponder(
        localPrivateKey: ByteArray,
        localPublicKey: ByteArray,
        remotePublicKey: ByteArray,
    ): Long
    external fun nativeCreateKkInitiatorForTest(
        localPrivateKey: ByteArray,
        localPublicKey: ByteArray,
        remotePublicKey: ByteArray,
        ephemeralPrivateKey: ByteArray,
    ): Long
    external fun nativeCreateKkResponderForTest(
        localPrivateKey: ByteArray,
        localPublicKey: ByteArray,
        remotePublicKey: ByteArray,
        ephemeralPrivateKey: ByteArray,
    ): Long

    external fun nativeWriteHandshake(handle: Long, payload: ByteArray): ByteArray
    external fun nativeReadHandshake(handle: Long, message: ByteArray): ByteArray
    external fun nativeEncrypt(handle: Long, plaintext: ByteArray): ByteArray
    external fun nativeDecrypt(handle: Long, ciphertext: ByteArray): ByteArray
    external fun nativeHandshakeHash(handle: Long): ByteArray
    external fun nativeRekeyOutgoing(handle: Long)
    external fun nativeRekeyIncoming(handle: Long)
    external fun nativeClose(handle: Long)
}

internal class RustNoiseSession(
    private var handle: Long,
) : NoiseNKSession {
    private fun openHandle(): Long {
        check(handle != 0L) { "Noise session is closed" }
        return handle
    }

    fun writeHandshake(payload: ByteArray = ByteArray(0)): ByteArray =
        RustNoiseNative.nativeWriteHandshake(openHandle(), payload)

    fun readHandshake(message: ByteArray): ByteArray =
        RustNoiseNative.nativeReadHandshake(openHandle(), message)

    override fun encrypt(plaintext: ByteArray): ByteArray =
        RustNoiseNative.nativeEncrypt(openHandle(), plaintext)

    override fun decrypt(ciphertext: ByteArray): ByteArray =
        RustNoiseNative.nativeDecrypt(openHandle(), ciphertext)

    override val handshakeHash: ByteArray
        get() = RustNoiseNative.nativeHandshakeHash(openHandle())

    override fun rekeyOutgoing() {
        RustNoiseNative.nativeRekeyOutgoing(openHandle())
    }

    override fun rekeyIncoming() {
        RustNoiseNative.nativeRekeyIncoming(openHandle())
    }

    override fun close() {
        val value = handle
        if (value == 0L) return
        handle = 0L
        RustNoiseNative.nativeClose(value)
    }
}