package com.peachpasswords.sync

/**
 * E7b: split a too-large padded sync frame across multiple Noise transport
 * messages and reassemble it on the other side. Mirror of peach-browser
 * `src/shared/sync-transport.ts`. Constants and the chunking protocol MUST
 * match the browser byte-for-byte.
 *
 * Why this exists: [SyncPadding.pad] can produce frames up to 24 MiB, but
 * Noise caps a single transport message at [MAX_NOISE_PLAINTEXT_BYTES]
 * plaintext (65519 — see crypto-core/src/lib.rs: 65535 ciphertext minus the
 * 16-byte Poly1305 auth tag). Until this layer existed, a payload above ~64 KiB
 * tripped the core's "Noise plaintext exceeds transport limit" on the sender
 * while the receiver sat waiting for a frame that never arrived.
 *
 * Wire shape for a multi-chunk frame:
 *
 *   Noise msg 1: [0x01][chunk 1 bytes ...]   <- 0x01 = [FRAME_CONTINUATION_MARKER]
 *   Noise msg 2: [0x01][chunk 2 bytes ...]
 *   ...
 *   Noise msg N: [chunk N bytes ...]         <- no marker on the last chunk
 *
 * Single-chunk frames go out as one Noise message with no marker, so the
 * common small-frame case is unchanged on the wire. The receiver recognises
 * the difference by the first byte of the decrypted plaintext: a marker
 * means "more chunks coming, do not unpad yet".
 *
 * The bucket ladder itself lives in [SyncPadding]. If [SyncPadding.bucket] or
 * this limit ever drifts, both sides change in lockstep — see the
 * cross-reference comment in [SyncPadding].
 */

/** Noise plaintext cap: 65535 ciphertext − 16-byte auth tag = 65519. */
const val MAX_NOISE_PLAINTEXT_BYTES: Int = 65519

/**
 * First byte of every non-last chunk of a chunked padded frame. The byte
 * `0x01` is unambiguous against a real padded frame: a 4-byte big-endian
 * length header whose payload is < 4 bytes can only start at byte 0 with a
 * zero, never `0x01`.
 */
const val FRAME_CONTINUATION_MARKER: Byte = 0x01

/** One Noise plaintext, rounded up — every control frame fits in a single chunk. */
private const val INITIAL_FRAME_BUFFER_BYTES: Int = 1 shl 16

/**
 * Pad [plaintext], optionally chunk it across the Noise plaintext limit,
 * encrypt each chunk, and hand each ciphertext to [sendCiphertext].
 *
 * A padded frame that fits one Noise message goes out as one ciphertext,
 * no marker, no per-chunk overhead. Otherwise the padded frame is sliced
 * into 64 KiB chunks; the first N−1 carry the [FRAME_CONTINUATION_MARKER],
 * the last does not, and the receiver reassembles by stripping the byte.
 */
internal suspend fun sendSyncFrame(
    plaintext: ByteArray,
    encrypt: (ByteArray) -> ByteArray,
    sendCiphertext: suspend (ByteArray) -> Unit,
) {
    val padded = SyncPadding.pad(plaintext)
    if (padded.size <= MAX_NOISE_PLAINTEXT_BYTES) {
        sendCiphertext(encrypt(padded))
        return
    }
    val total = padded.size
    var offset = 0
    while (offset < total) {
        val remaining = total - offset
        // A non-last chunk must leave room for the one-byte marker under the
        // Noise plaintext cap; the last chunk carries no marker and can use
        // the full cap. This slicing is one byte shy of the browser's loop,
        // which slices at the cap and then prepends the marker — that path
        // produces [marker][MAX bytes] = MAX+1 bytes, which the crypto core
        // rejects. Wire shape ([0x01][data]) is unchanged across platforms.
        val cap = if (remaining <= MAX_NOISE_PLAINTEXT_BYTES) {
            MAX_NOISE_PLAINTEXT_BYTES
        } else {
            MAX_NOISE_PLAINTEXT_BYTES - 1
        }
        val end = minOf(offset + cap, total)
        val chunkBytes = end - offset
        val chunk = padded.copyOfRange(offset, end)
        val isLast = end == total
        val pt: ByteArray = if (isLast) {
            chunk
        } else {
            ByteArray(chunkBytes + 1).also {
                it[0] = FRAME_CONTINUATION_MARKER
                chunk.copyInto(it, 1)
            }
        }
        sendCiphertext(encrypt(pt))
        offset += chunkBytes
    }
}

/**
 * Inverse of [sendSyncFrame]: read Noise ciphertexts from
 * [receiveCiphertext], reassemble chunked padded frames, and return the
 * original payload. The first byte of every decrypted message tells the
 * loop whether more chunks are coming; the last chunk is unmarker and
 * triggers unpad.
 */
internal suspend fun receiveSyncFrame(
    receiveCiphertext: suspend () -> ByteArray,
    decrypt: (ByteArray) -> ByteArray,
): ByteArray {
    // The buffer doubles on demand up to MAX_SYNC_FRAME_BYTES rather than
    // reserving that cap for every frame: most frames are small control
    // messages, and a fixed 24 MiB allocation per frame is itself the
    // availability problem (PAI-2026-003). Doubling keeps the old O(n²)
    // copy cost away — at most ~18 copies to reach the cap. Cumulative size
    // and chunk count are still checked BEFORE writing, so an authenticated
    // peer can't grow memory unboundedly or pin the loop on tiny marked
    // chunks.
    var buffer = ByteArray(INITIAL_FRAME_BUFFER_BYTES)
    var size = 0
    var chunkCount = 0
    var chunkedFrame = false
    while (true) {
        chunkCount++
        val pt = decrypt(receiveCiphertext())
        val marked = pt.isNotEmpty() && pt[0] == FRAME_CONTINUATION_MARKER
        val dataStart = if (marked) 1 else 0
        val dataSize = pt.size - dataStart
        // MAX_SYNC_FRAME_BYTES doubles as the chunk-count ceiling: even a
        // maximally-fragmented adversarial stream (one byte per chunk) hits a
        // finite loop bound instead of running until the size cap.
        if (size + dataSize > MAX_SYNC_FRAME_BYTES || chunkCount > MAX_SYNC_FRAME_BYTES) {
            buffer.fill(0, 0, size)
            pt.fill(0)
            if (size + dataSize > MAX_SYNC_FRAME_BYTES) {
                requireSyncFrameSize(size + dataSize)
            }
            throw SyncFrameTooLargeException(
                "sync frame chunk count exceeded $MAX_SYNC_FRAME_BYTES",
            )
        }
        if (size + dataSize > buffer.size) {
            val grown = (buffer.size * 2).coerceAtMost(MAX_SYNC_FRAME_BYTES)
            buffer = buffer.copyOf(maxOf(grown, size + dataSize))
        }
        if (!chunkedFrame) {
            // First chunk. Marker ⇒ chunked frame; no marker ⇒ the entire
            // single-chunk frame, unchanged from before the chunking layer
            // existed. The chunkedFrame latch then enforces the rest of the
            // protocol: every subsequent chunk must be marked until one isn't.
            chunkedFrame = true
            if (!marked) {
                return SyncPadding.unpad(pt)
            }
        } else if (!marked) {
            // Final chunk of a chunked frame. Unpad and zero the buffer so
            // partial plaintext doesn't linger.
            pt.copyInto(buffer, size, dataStart, pt.size)
            val written = size + dataSize
            val out = SyncPadding.unpad(buffer.copyOfRange(0, written))
            buffer.fill(0, 0, written)
            return out
        }
        if (dataSize == 0) {
            // PB-26-007: a marker-only continuation chunk carries no data and
            // only burns iterations. Require positive progress per chunk.
            buffer.fill(0, 0, size)
            pt.fill(0)
            throw SyncFrameTooLargeException("sync frame continuation chunk made no progress")
        }
        pt.copyInto(buffer, size, dataStart, pt.size)
        size += dataSize
    }
}