package com.peachpasswords.codex

import org.junit.Assert.assertArrayEquals
import org.junit.Assert.assertEquals
import org.junit.Assert.assertTrue
import org.junit.Test
import kotlin.random.Random

class CodexCodecTest {

    @Test fun rs_recovers_dropped_shards() {
        val data = 6
        val parity = computeParityShards(data) // 2
        val size = 40
        val rng = Random(7)
        val original = Array(data) { ByteArray(size) { rng.nextInt().toByte() } }
        val shards = Array(data + parity) { i -> if (i < data) original[i].copyOf() else ByteArray(size) }
        ReedSolomon.encode(shards, data, parity)

        // Drop `parity` shards (the max tolerable) — here drop 2 data shards.
        val present = BooleanArray(data + parity) { true }
        shards[1] = ByteArray(0); present[1] = false
        shards[4] = ByteArray(0); present[4] = false
        ReedSolomon.reconstruct(shards, data, parity, present)

        for (i in 0 until data) assertArrayEquals("shard $i", original[i], shards[i])
    }

    @Test fun gf_mul_div_are_inverse() {
        for (a in 1..255) for (b in 1..255) {
            assertEquals(a, Galois.div(Galois.mul(a, b), b))
        }
    }

    @Test fun full_codec_round_trip() {
        val payload = CodexPayload(
            version = 2, verifyCode = "PX-AB12-CD34",
            vaultBlob = "x".repeat(5000), timestamp = 111, entryCount = 7,
            vaultFingerprint = "DEADBEEF".repeat(4), generatedAt = 222,
        )
        val bytes = CodexCodec.json.encodeToString(CodexPayload.serializer(), payload).toByteArray()
        val compressed = CodexCodec.deflate(bytes)
        assertArrayEquals(bytes, CodexCodec.inflate(compressed))

        val (envelopes, layout) = CodexCodec.encodeToEnvelopes(compressed, payload.verifyCode)
        assertTrue(envelopes.all { CodexCodec.verifyChunkIntegrity(it) })
        assertEquals(layout.dataChunks + layout.parityChunks, envelopes.size)

        // Lose the maximum tolerable codes in stripe 0, then rebuild.
        val parity = layout.parityChunks / layout.stripeCount
        val survivors = envelopes.toMutableList()
        repeat(parity) { survivors.removeAt(0) } // drop first `parity` data shards
        val recovered = CodexCodec.reconstructCompressed(survivors)
        val parsed = CodexCodec.json.decodeFromString(CodexPayload.serializer(), String(CodexCodec.inflate(recovered)))
        assertEquals(payload, parsed)
    }

    @Test fun corrupted_chunk_fails_integrity() {
        val (envelopes, _) = CodexCodec.encodeToEnvelopes(CodexCodec.deflate("hello world".toByteArray()), "PX-0000-0000")
        val tampered = envelopes.first().copy(d = CodexCodec.base64Url("tampered".toByteArray()))
        assertTrue(!CodexCodec.verifyChunkIntegrity(tampered))
    }

    @Test fun inflate_exactly_at_the_cap_round_trips() {
        val data = ByteArray(CodexCodec.MAX_INFLATED_BYTES.toInt()) // zeros compress to ~17 KB
        assertArrayEquals(data, CodexCodec.inflate(CodexCodec.deflate(data)))
    }

    /**
     * PAI-2026-002: the per-chunk hash covers `d` only, so a hostile code
     * passes verifyChunkIntegrity with a tiny payload and any `t`/`S`/`i` it
     * likes. Every case below used to reach Array(t) before being rejected.
     */
    private fun hostileChunk(
        total: Int = 4,
        index: Int = 0,
        stripe: Int = 0,
        stripeCount: Int = 1,
        data: ByteArray = ByteArray(CHUNK_SIZE),
        verifyCode: String = "PX-AB12-CD34",
    ): ChunkEnvelope = ChunkEnvelope(
        v = 3, c = verifyCode, s = stripe, capS = stripeCount, i = index, t = total,
        p = index >= total - 2,
        d = CodexCodec.base64Url(data),
        h = CodexCodec.sha256Hex(data).take(32),
    )

    @Test fun hostile_shard_total_is_rejected_before_allocation() {
        val hostile = hostileChunk(total = Int.MAX_VALUE)
        assertTrue("geometry is rejected", !CodexCodec.isValidEnvelope(hostile))
        val error = try {
            CodexCodec.reconstructCompressed(listOf(hostile))
            null
        } catch (e: CodexException) {
            e
        }
        assertEquals(CodexErrorCode.INSUFFICIENT_CHUNKS, error?.code)
    }

    @Test fun out_of_range_shard_and_stripe_indices_are_rejected() {
        assertTrue(!CodexCodec.isValidEnvelope(hostileChunk(total = 4, index = 4)))
        assertTrue(!CodexCodec.isValidEnvelope(hostileChunk(total = 4, index = -1)))
        assertTrue(!CodexCodec.isValidEnvelope(hostileChunk(stripe = 2, stripeCount = 2)))
        assertTrue(!CodexCodec.isValidEnvelope(hostileChunk(stripeCount = 0)))
        assertTrue(!CodexCodec.isValidEnvelope(hostileChunk(stripeCount = 100_000)))
        assertTrue(!CodexCodec.isValidEnvelope(hostileChunk(total = 256)))
    }

    @Test fun short_and_oversized_chunk_data_are_rejected() {
        val short = hostileChunk(data = ByteArray(CHUNK_SIZE - 1))
        val oversized = hostileChunk(data = ByteArray(CHUNK_SIZE + 1))
        assertTrue(!CodexCodec.isValidEnvelope(short))
        assertTrue(!CodexCodec.isValidEnvelope(oversized))
    }

    @Test fun malformed_verify_code_is_rejected() {
        assertTrue(!CodexCodec.isValidEnvelope(hostileChunk(verifyCode = "PX-AB12-CD3")))
        assertTrue(CodexCodec.isValidEnvelope(hostileChunk()))
    }

    @Test fun mixed_verify_codes_and_shard_totals_are_rejected() {
        val (envelopes, _) = CodexCodec.encodeToEnvelopes(
            CodexCodec.deflate("hello codex".toByteArray()),
            "PX-AB12-CD34",
        )
        val other = CodexCodec.encodeToEnvelopes(
            CodexCodec.deflate("hello codex".toByteArray()),
            "PX-ZZ99-YY88",
        ).first.first()
        var error = try {
            CodexCodec.reconstructCompressed(envelopes + other); null
        } catch (e: CodexException) {
            e
        }
        assertEquals(CodexErrorCode.INVALID_ENVELOPE, error?.code)

        val mixedTotal = envelopes.first().copy(t = envelopes.first().t + 1, i = 0)
        error = try {
            CodexCodec.reconstructCompressed(listOf(envelopes.first(), mixedTotal)); null
        } catch (e: CodexException) {
            e
        }
        assertEquals(CodexErrorCode.INVALID_ENVELOPE, error?.code)
    }

    @Test fun hostile_envelopes_are_skipped_and_real_codes_still_restore() {
        val (envelopes, _) = CodexCodec.encodeToEnvelopes(
            CodexCodec.deflate("survives hostile neighbours".toByteArray()),
            "PX-AB12-CD34",
        )
        val recovered = CodexCodec.reconstructCompressed(
            listOf(hostileChunk(total = Int.MAX_VALUE)) + envelopes +
                listOf(hostileChunk(verifyCode = "nope")),
        )
        assertEquals("survives hostile neighbours", String(CodexCodec.inflate(recovered)))
    }

    @Test fun inflate_refuses_decompression_bombs() {
        val bomb = CodexCodec.deflate(ByteArray(CodexCodec.MAX_INFLATED_BYTES.toInt() + 1))
        assertTrue("bomb compresses tiny", bomb.size < 64 * 1024)
        val error = try {
            CodexCodec.inflate(bomb)
            null
        } catch (e: CodexException) {
            e
        }
        assertEquals(CodexErrorCode.PAYLOAD_TOO_LARGE, error?.code)
    }
}