package com.peachpasswords.sync

import kotlinx.coroutines.channels.Channel
import kotlinx.coroutines.runBlocking
import org.junit.Assert.assertEquals
import org.junit.Assert.assertNull
import org.junit.Assert.assertTrue
import org.junit.Test

/**
 * Inline vault snapshot protocol parity with the browser
 * `vault-adopt-snapshot.test.ts`. Round-trip, chunking, and validation
 * tests for the ADOPT_SNAPSHOT / ADOPT_SNAPSHOT_APPLIED frames.
 */
class VaultAdoptSnapshotTest {

    private fun memoryPair(): Pair<PairingTransport, PairingTransport> {
        val a2b = Channel<ByteArray>(Channel.UNLIMITED)
        val b2a = Channel<ByteArray>(Channel.UNLIMITED)
        val a = object : PairingTransport {
            override suspend fun send(msg: ByteArray) { a2b.send(msg) }
            override suspend fun receive(): ByteArray = b2a.receive()
            override suspend fun close() {}
        }
        val b = object : PairingTransport {
            override suspend fun send(msg: ByteArray) { b2a.send(msg) }
            override suspend fun receive(): ByteArray = a2b.receive()
            override suspend fun close() {}
        }
        return a to b
    }

    @Test
    fun `round-trips an empty vault`() = runBlocking {
        val (source, receiver) = memoryPair()
        sendVaultSnapshot(source, emptyMap(), emptyMap())
        val (entries, tombstones) = receiveVaultSnapshot(receiver)
        assertTrue(entries.isEmpty())
        assertTrue(tombstones.isEmpty())
    }

    @Test
    fun `round-trips 50 entries and 5 tombstones in a single chunk`() = runBlocking {
        val (source, receiver) = memoryPair()
        val entries = (0 until 50).associate { "id-$it" to "cipher-$it" }
        val tombstones = (0 until 5).associate { "dead-$it" to (1000L + it) }
        sendVaultSnapshot(source, entries, tombstones)
        val (receivedEntries, receivedTombstones) = receiveVaultSnapshot(receiver)
        assertEquals(entries, receivedEntries)
        assertEquals(tombstones, receivedTombstones)
    }

    @Test
    fun `splits a large vault into 3+ chunks and reassembles`() = runBlocking {
        val (source, receiver) = memoryPair()
        var frames = 0
        val countingSource = object : PairingTransport {
            override suspend fun send(msg: ByteArray) { frames++; source.send(msg) }
            override suspend fun receive(): ByteArray = source.receive()
            override suspend fun close() {}
        }
        val big = "x".repeat(200 * 1024)
        val entries = (0 until 5).associate { "id-$it" to big }
        val tombstones = mapOf("dead-0" to 42L)
        sendVaultSnapshot(countingSource, entries, tombstones)
        val (receivedEntries, receivedTombstones) = receiveVaultSnapshot(receiver)
        assertTrue("expected 3+ chunks, got $frames", frames >= 3)
        assertEquals(entries, receivedEntries)
        assertEquals(tombstones, receivedTombstones)
    }

    @Test
    fun `rejects out-of-order chunks`() = runBlocking {
        val (source, receiver) = memoryPair()
        sendFrame(source, AdoptionFrame.Snapshot(
            chunkIndex = 1,
            chunkCount = 2,
            entries = mapOf("a" to "x"),
            tombstones = emptyMap(),
        ))
        val error = runCatching { receiveVaultSnapshot(receiver) }.exceptionOrNull()
        assertTrue(error is IllegalStateException)
        assertTrue(error!!.message!!.contains("OUT_OF_ORDER"))
    }

    @Test
    fun `rejects chunk_count exceeding MAX_SNAPSHOT_CHUNKS`() = runBlocking {
        val (source, receiver) = memoryPair()
        sendFrame(source, AdoptionFrame.Snapshot(
            chunkIndex = 0,
            chunkCount = MAX_SNAPSHOT_CHUNKS + 1,
            entries = emptyMap(),
            tombstones = emptyMap(),
        ))
        val error = runCatching { receiveVaultSnapshot(receiver) }.exceptionOrNull()
        assertTrue(error is IllegalStateException)
        assertTrue(error!!.message!!.contains("TOO_MANY_CHUNKS"))
    }

    @Test
    fun `validates ADOPT_SNAPSHOT frames`() {
        val valid = AdoptionFrame.Snapshot(
            chunkIndex = 0, chunkCount = 1,
            entries = mapOf("a" to "x"), tombstones = emptyMap(),
        )
        assertTrue(isAdoptSnapshot(valid))
        assertTrue(!isAdoptSnapshot(null))
        assertTrue(!isAdoptSnapshot(AdoptionFrame.Complete))
    }

    @Test
    fun `snapshot_capable serializes in ADOPT_REQUEST`() {
        val request = AdoptionFrame.Request(
            sessionRequestId = "req-1",
            deviceName = "test",
            deviceType = "extension",
            snapshotCapable = true,
        )
        val json = AdoptJson.encodeToString(AdoptionFrame.serializer(), request)
        assertTrue(json.contains("\"snapshot_capable\":true"))
        // Absent when null (encodeDefaults = false)
        val noFlag = AdoptionFrame.Request(
            sessionRequestId = "req-2",
            deviceName = "test",
            deviceType = "extension",
        )
        val jsonNoFlag = AdoptJson.encodeToString(AdoptionFrame.serializer(), noFlag)
        assertTrue(!jsonNoFlag.contains("snapshot_capable"))
    }

    @Test
    fun `rejects raw frame exceeding wire ceiling`() = runBlocking {
        // H-09: a hostile source could send an arbitrarily large frame to exhaust
        // receiver memory before JSON parsing. receiveVaultSnapshot checks the raw
        // byte length against SNAPSHOT_FRAME_WIRE_CEILING (512 KiB + 256 KiB = 768 KiB)
        // and throws before attempting deserialization.
        val ceiling = MAX_SNAPSHOT_CHUNK_BYTES + 256 * 1024
        val oversized = ByteArray(ceiling + 1)
        val channel = object : PairingTransport {
            override suspend fun send(msg: ByteArray) {}
            override suspend fun receive(): ByteArray = oversized
            override suspend fun close() {}
        }
        val error = runCatching { receiveVaultSnapshot(channel) }.exceptionOrNull()
        assertTrue(error is IllegalStateException)
        assertTrue(error!!.message!!.contains("CHUNK_TOO_LARGE"))
    }

    @Test
    fun `rejects chunk_count mismatch on subsequent frame`() = runBlocking {
        val (source, receiver) = memoryPair()
        sendFrame(source, AdoptionFrame.Snapshot(
            chunkIndex = 0,
            chunkCount = 2,
            entries = mapOf("a" to "x"),
            tombstones = emptyMap(),
        ))
        sendFrame(source, AdoptionFrame.Snapshot(
            chunkIndex = 1,
            chunkCount = 3,
            entries = mapOf("b" to "y"),
            tombstones = emptyMap(),
        ))
        val error = runCatching { receiveVaultSnapshot(receiver) }.exceptionOrNull()
        assertTrue(error is IllegalStateException)
        assertTrue(error!!.message!!.contains("CHUNK_COUNT_MISMATCH"))
    }

    @Test
    fun `round-trips duress_slot on final chunk`() = runBlocking {
        val (source, receiver) = memoryPair()
        val entries = mapOf("id-0" to "cipher-0")
        val tombstones = mapOf("dead-0" to 42L)
        val duressJson = """{"version":1,"salt":"aa","nonce":"bb","ciphertext":"cc","tag":"dd"}"""
        sendVaultSnapshot(source, entries, tombstones, duressSlot = duressJson)
        val snapshot = receiveVaultSnapshot(receiver)
        assertEquals(entries, snapshot.entries)
        assertEquals(tombstones, snapshot.tombstones)
        assertEquals(duressJson, snapshot.duressSlot)
    }

    @Test
    fun `duress_slot absent returns null`() = runBlocking {
        val (source, receiver) = memoryPair()
        sendVaultSnapshot(source, mapOf("id-0" to "cipher-0"), emptyMap())
        val snapshot = receiveVaultSnapshot(receiver)
        assertNull(snapshot.duressSlot)
    }

    @Test
    fun `duress_slot rides final chunk only in multi-chunk transfer`() = runBlocking {
        val (source, receiver) = memoryPair()
        val big = "x".repeat(200 * 1024)
        val entries = (0 until 5).associate { "id-$it" to big }
        val tombstones = mapOf("dead-0" to 42L)
        val duressJson = """{"version":1,"salt":"ff","nonce":"ee","ciphertext":"dd","tag":"cc"}"""
        sendVaultSnapshot(source, entries, tombstones, duressSlot = duressJson)
        val snapshot = receiveVaultSnapshot(receiver)
        assertEquals(entries, snapshot.entries)
        assertEquals(tombstones, snapshot.tombstones)
        assertEquals(duressJson, snapshot.duressSlot)
    }

    @Test
    fun `old frame without duress_slot deserializes`() {
        val oldJson = """{"type":"ADOPT_SNAPSHOT","chunk_index":0,"chunk_count":1,"entries":{"a":"x"},"tombstones":{}}"""
        val frame = runCatching {
            AdoptJson.decodeFromString(AdoptionFrame.serializer(), oldJson)
        }.getOrNull()
        assertTrue(isAdoptSnapshot(frame))
        assertTrue((frame as AdoptionFrame.Snapshot).duressSlot == null)
    }

}