package com.peachpasswords.sync

import kotlinx.coroutines.async
import kotlinx.coroutines.channels.Channel
import kotlinx.coroutines.runBlocking
import org.bouncycastle.crypto.params.X25519PrivateKeyParameters
import org.junit.Assert.assertArrayEquals
import org.junit.Assert.assertEquals
import org.junit.Assert.assertFalse
import org.junit.Assert.assertTrue
import org.junit.Test
import java.security.SecureRandom

/**
 * E7b end-to-end chunking — the deliverable for the wire shape. Initiator and
 * responder each build a real [noiseKKInitiator]/[noiseKKResponder] session,
 * pipe ciphertexts through an in-memory channel, and run [sendSyncFrame] /
 * [receiveSyncFrame] across a payload large enough that it MUST chunk.
 *
 * Pins the protocol that the browser, an Android phone-to-phone pair, and
 * future Android↔iOS ports all depend on: 0x01 marker on every non-last
 * chunk, no marker on the last, reassembled bytes equal original, and a
 * single-chunk frame goes out as one message unchanged.
 */
class SyncFrameTransportTest {
    private data class KeyPair(val publicKey: ByteArray, val privateKey: ByteArray)

    private fun keys(): KeyPair {
        val privateKey = X25519PrivateKeyParameters(SecureRandom())
        return KeyPair(privateKey.generatePublicKey().encoded, privateKey.encoded)
    }

    /**
     * Build two Noise sessions against each other. The two halves share an
     * in-memory channel pair so the test owns every ciphertext on the wire —
     * no socket, no real network, and no state from one test leaking into
     * the next.
     */
    private fun sessions(): Pair<NoiseNKSession, NoiseNKSession> = runBlocking {
        val initiator = keys()
        val responder = keys()
        val initiatorToResponder = Channel<ByteArray>(Channel.UNLIMITED)
        val responderToInitiator = Channel<ByteArray>(Channel.UNLIMITED)
        val initiatorHandshake = async {
            noiseKKInitiator(
                initiator.privateKey,
                initiator.publicKey,
                responder.publicKey,
                { initiatorToResponder.send(it) },
                { responderToInitiator.receive() },
            )
        }
        val responderHandshake = async {
            noiseKKResponder(
                responder.privateKey,
                responder.publicKey,
                initiator.publicKey,
                { responderToInitiator.send(it) },
                { initiatorToResponder.receive() },
            )
        }
        initiatorHandshake.await() to responderHandshake.await()
    }

    @Test
    fun `a payload just over 64 KiB round-trips across the chunking layer`() = runBlocking {
        val (initiator, responder) = sessions()
        val senderToReceiver = Channel<ByteArray>(Channel.UNLIMITED)

        // 70 KiB lands on the bucket ladder at 128 KiB — comfortably above the
        // 65519 byte Noise plaintext limit, forcing at least 3 chunks.
        val payload = ByteArray(70 * 1024) { (it * 17 and 0xFF).toByte() }

        val sender = async {
            sendSyncFrame(payload, initiator::encrypt) { ct ->
                senderToReceiver.send(ct)
            }
        }
        val receiver = async {
            receiveSyncFrame({ senderToReceiver.receive() }, responder::decrypt)
        }
        sender.await()
        assertArrayEquals(payload, receiver.await())
        initiator.close()
        responder.close()
    }

    @Test
    fun `a payload under the Noise limit goes out as exactly one ciphertext`() = runBlocking {
        val (initiator, responder) = sessions()
        val senderToReceiver = Channel<ByteArray>(Channel.UNLIMITED)

        // 5-byte payload pads to the 256 minimum bucket — well under 65519.
        val payload = "hello".toByteArray()

        val sentCiphertexts = mutableListOf<ByteArray>()
        val sender = async {
            sendSyncFrame(payload, initiator::encrypt) { ct ->
                sentCiphertexts.add(ct)
                senderToReceiver.send(ct)
            }
        }
        val receiver = async {
            receiveSyncFrame({ senderToReceiver.receive() }, responder::decrypt)
        }
        sender.await()
        // Single-chunk path: no marker on the wire, exactly one ciphertext.
        assertEquals(1, sentCiphertexts.size)
        // 256 plaintext + 16-byte Noise tag = 272 bytes ciphertext.
        assertEquals(256 + 16, sentCiphertexts[0].size)
        assertArrayEquals(payload, receiver.await())
        initiator.close()
        responder.close()
    }

    @Test
    fun `a chunked payload marks every non-final chunk with 0x01`() = runBlocking {
        val (initiator, responder) = sessions()
        val senderToReceiver = Channel<ByteArray>(Channel.UNLIMITED)

        // 130 KiB payload → padded frame is well above 192 KiB; definite
        // multi-chunk.
        val payload = ByteArray(130 * 1024) { (it and 0xFF).toByte() }

        // Wrap encrypt so we can also observe what plaintexts the chunking
        // layer hands to the Noise session. This is what `0x01` is being
        // asserted against on the wire.
        val observedPlaintexts = mutableListOf<ByteArray>()
        val observingEncrypt: (ByteArray) -> ByteArray = { pt ->
            observedPlaintexts.add(pt.copyOf())
            initiator.encrypt(pt)
        }
        val sender = async {
            sendSyncFrame(payload, observingEncrypt) { ct ->
                senderToReceiver.send(ct)
            }
        }
        val receiver = async {
            receiveSyncFrame({ senderToReceiver.receive() }, responder::decrypt)
        }
        sender.await()
        assertArrayEquals(payload, receiver.await())
        initiator.close()
        responder.close()

        // Multi-chunk: at least 2 chunks, every non-final is at most
        // MAX_NOISE_PLAINTEXT_BYTES (its data leaves one byte for the marker)
        // and starts with the marker; the final chunk is at most the cap and
        // does not carry the marker.
        assertTrue("expected multi-chunk split", observedPlaintexts.size >= 2)
        observedPlaintexts.dropLast(1).forEach {
            assertEquals(
                "non-final chunk must start with the continuation marker",
                FRAME_CONTINUATION_MARKER, it[0],
            )
            assertTrue(
                "every non-final chunk must fit under the Noise cap with the marker",
                it.size <= MAX_NOISE_PLAINTEXT_BYTES,
            )
        }
        val last = observedPlaintexts.last()
        assertTrue("last chunk must be ≤ the Noise cap", last.size <= MAX_NOISE_PLAINTEXT_BYTES)
        assertFalse(
            "final chunk must NOT carry the continuation marker",
            last.isNotEmpty() && last[0] == FRAME_CONTINUATION_MARKER,
        )
    }
}