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.
 *
 * Mirrors the invariants in peach-browser `sync-padding.test.ts` and
 * exercises the wire format end-to-end the browser's transport.test.ts does
 * not.
 */
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 sessions (initiator and responder) against each other, sharing
     * an in-memory channel pair so the test owns every ciphertext on the wire
     * — no socket, no real network.
     */
    private fun sessions(): Pair<NoiseNKSession, NoiseNKSession> = runBlocking {
        val initiator = keys()
        val responder = keys()
        val aToB = Channel<ByteArray>(Channel.UNLIMITED)
        val bToA = Channel<ByteArray>(Channel.UNLIMITED)
        val initiatorHandshake = async {
            noiseKKInitiator(
                initiator.privateKey,
                initiator.publicKey,
                responder.publicKey,
                { aToB.send(it) },
                { bToA.receive() },
            )
        }
        val responderHandshake = async {
            noiseKKResponder(
                responder.privateKey,
                responder.publicKey,
                initiator.publicKey,
                { bToA.send(it) },
                { aToB.receive() },
            )
        }
        initiatorHandshake.await() to responderHandshake.await()
    }

    /**
     * Capture every ciphertext the chunking helper would write onto the wire
     * and the noise session would produce for it. We instrument the session
     * itself by routing [sendCiphertext] and [receiveCiphertext] through an
     * in-memory channel and observing the pre-encryption plaintexts as the
     * test reads them back.
     */
    private class Wire(
        val outbound: Channel<ByteArray>,
        val ciphertextBack: Channel<ByteArray>,
    )

    private fun bridge(): Pair<Wire, Wire> {
        val aToB = Channel<ByteArray>(Channel.UNLIMITED)
        val bToA = Channel<ByteArray>(Channel.UNLIMITED)
        return Wire(aToB, bToA) to Wire(bToA, aToB)
    }

    @Test
    fun `a payload just over 64 KiB round-trips across the chunking layer`() = runBlocking {
        val (initiator, responder) = sessions()
        val (aWire, bWire) = bridge()

        // 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 -> aWire.outbound.send(ct) }
        }
        val received = async {
            receiveSyncFrame(
                { aWire.ciphertextBack.receive() },
                responder::decrypt,
            )
        }
        sender.await()
        assertArrayEquals(payload, received.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 (aWire, bWire) = bridge()

        // 256-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); aWire.outbound.send(ct) }
        }
        val received = async {
            receiveSyncFrame(
                { aWire.ciphertextBack.receive() },
                responder::decrypt,
            )
        }
        sender.await()
        // Single-chunk path: no marker on the wire, exactly one ciphertext.
        assertEquals(1, sentCiphertexts.size)
        // 256 plaintext + 16 Noise tag = 272 ciphertext.
        assertEquals(256 + 16, sentCiphertexts[0].size)
        assertArrayEquals(payload, received.await())
        initiator.close()
        responder.close()
    }

    @Test
    fun `a chunked payload marks every non-final chunk with 0x01`() = runBlocking {
        val (initiator, responder) = sessions()
        val (aWire, bWire) = bridge()

        // 130 KiB payload → padded frame 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 -> aWire.outbound.send(ct) }
        }
        val received = async {
            receiveSyncFrame(
                { aWire.ciphertextBack.receive() },
                responder::decrypt,
            )
        }
        sender.await()
        assertArrayEquals(payload, received.await())
        initiator.close()
        responder.close()

        // Multi-chunk: at least 2 chunks, every non-final is exactly
        // MAX_NOISE_PLAINTEXT_BYTES 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],
            )
            assertEquals(
                "every non-final chunk must hit the Noise cap exactly",
                MAX_NOISE_PLAINTEXT_BYTES, it.size,
            )
        }
        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,
        )
    }
}