package com.peachpasswords.vault

import kotlinx.coroutines.runBlocking
import org.junit.Assert.assertEquals
import org.junit.Assert.assertTrue
import org.junit.Test

/** Exposes the raw stored slots so the padding lengths (MN-1) can be inspected. */
private class PaddingFakeStorage : VaultStorage {
    var slotA: StoredVault? = null
    var slotBNoise: String? = null
    var decoy: StoredVault? = null
    private var ver: Long = 0
    private var settings = VaultSettings()
    override fun loadSlotA() = slotA
    override fun hasVault() = slotA != null
    override fun writeSlotA(slot: StoredVault): Boolean { slotA = slot; return true }
    override fun writeSlotB(blobBase64: String): Boolean { slotBNoise = blobBase64; return true }
    override fun loadDecoySlot() = decoy
    override fun writeDecoySlot(slot: StoredVault): Boolean { decoy = slot; return true }
    override fun hasDecoySlot() = decoy != null
    override fun clearDecoySlot() { decoy = null }
    override fun version() = ver
    override fun stampLocalMutation(): Long { ver = maxOf(System.currentTimeMillis(), ver + 1); return ver }
    override fun loadSettings() = settings
    override fun writeSettings(s: VaultSettings): Boolean { settings = s; return true }
    override fun clearVault(): Boolean { slotA = null; slotBNoise = null; ver = 0; settings = VaultSettings(); return true }
}

class SlotPaddingTest {
    @Test fun payload_round_trips_through_bucket_padding() {
        val payload = VaultPayload(
            entries = listOf(VaultEntry(id = "1", title = "GitHub", type = "login", password = "pw", createdAt = 0, modifiedAt = 0)),
            rootKeyB64 = "AAAA",
        )
        val padded = VaultCodec.encode(payload)
        assertEquals("padded to the smallest bucket", VaultCrypto.PADDING_BUCKETS.first(), padded.size)
        val decoded = VaultCodec.decode(padded)
        assertEquals("padding is transparent to decode", "GitHub", decoded.entries.single().title)
        assertEquals("AAAA", decoded.rootKeyB64)
    }

    @Test fun legacy_unpadded_blob_still_decodes() {
        // A blob written before MN-1 has no trailing NUL padding — must still decode.
        val raw = "{\"entries\":[],\"rootKeyB64\":\"X\"}".toByteArray(Charsets.UTF_8)
        assertEquals("X", VaultCodec.decode(raw).rootKeyB64)
    }

    @Test fun above_top_bucket_vault_codec_round_trip_preserves_payload() {
        val notes = "n".repeat(300_000)
        val payload = VaultPayload(
            entries = listOf(
                VaultEntry(
                    id = "large",
                    title = "Large note",
                    type = "note",
                    notes = notes,
                    createdAt = 0,
                    modifiedAt = 0,
                ),
            ),
            rootKeyB64 = "root",
        )

        val encoded = VaultCodec.encode(payload)
        // Above the top bucket the shared rule doubles from 262144 (PB-26-128).
        assertEquals(524288, encoded.size)
        assertEquals(notes, VaultCodec.decode(encoded).entries.single().notes)
    }

    @Test fun real_and_duress_slots_share_a_bucket_for_small_vaults() = runBlocking {
        val storage = PaddingFakeStorage()
        val repo = VaultRepository(storage)
        repo.createVault("the-real-password-123")
        repo.addEntry(VaultEntry(id = "", title = "Bank", type = "login", password = "pw", createdAt = 0, modifiedAt = 0))
        assertTrue(repo.setDuressPassphrase("open-the-decoy-456"))

        val slotALen = storage.slotA!!.blob.length
        val decoyLen = storage.decoy!!.blob.length
        assertEquals("a small real vault and its decoy quantise to the same slot length", slotALen, decoyLen)
        // And the always-present noise slot matches slot A's length too (±base64 rounding).
        assertTrue("noise slot matches slot A length", kotlin.math.abs(storage.slotBNoise!!.length - slotALen) <= 4)
    }

    /** Mirrors VaultStore.writeAligned: write one record, re-align both. */
    private class AlignedSlotRecords {
        var a: String? = null
        var d: String? = null
        fun writeSlotA(slot: StoredVault) {
            val (aligned, sibling) = VaultCodec.alignSlotRecords(VaultCodec.encodeSlot(slot), d)
            a = aligned
            if (sibling != null) d = sibling
        }
        fun writeDecoy(slot: StoredVault) {
            val (aligned, sibling) = VaultCodec.alignSlotRecords(VaultCodec.encodeSlot(slot), a)
            d = aligned
            if (sibling != null) a = sibling
        }
    }

    private fun slotWithBlob(blobLength: Int) = StoredVault(
        version = 1,
        salt = "c2FsdHNhbHRzYWx0c2FsdA==",
        blob = "A".repeat(blobLength),
    )

    /**
     * PAI-2026-004: a real decoy's blob is frozen at the size it had when
     * duress was configured, because re-sealing it would need the duress
     * passphrase. Walking slot A up through every padding bucket used to leave
     * the decoy record visibly shorter, which told an at-rest observer both
     * that duress was configured and how far the real vault had grown.
     */
    @Test fun slot_records_stay_aligned_as_the_vault_crosses_every_bucket() {
        val slots = AlignedSlotRecords()
        slots.writeSlotA(slotWithBlob(VaultCrypto.PADDING_BUCKETS.first()))
        // Duress is configured once, at the smallest bucket, and frozen there.
        val frozenDecoy = slotWithBlob(VaultCrypto.PADDING_BUCKETS.first())
        slots.writeDecoy(frozenDecoy)
        val configuredAt = slots.d!!.length

        for (bucket in VaultCrypto.PADDING_BUCKETS) {
            slots.writeSlotA(slotWithBlob(bucket))
            assertEquals("bucket $bucket: the two records are the same length", slots.a!!.length, slots.d!!.length)
            assertEquals("bucket $bucket: padding is transparent to decode", bucket, VaultCodec.decodeSlot(slots.a!!)!!.blob.length)
            assertEquals("bucket $bucket: the decoy still decodes", frozenDecoy, VaultCodec.decodeSlot(slots.d!!))
        }
        assertTrue("alignment actually did work at the top bucket", slots.a!!.length > configuredAt)
    }

    @Test fun slot_records_stay_aligned_when_the_vault_shrinks_below_the_decoy() {
        val slots = AlignedSlotRecords()
        slots.writeSlotA(slotWithBlob(VaultCrypto.PADDING_BUCKETS.first()))
        slots.writeDecoy(slotWithBlob(VaultCrypto.PADDING_BUCKETS.last()))
        slots.writeSlotA(slotWithBlob(VaultCrypto.PADDING_BUCKETS.first()))
        assertEquals(slots.a!!.length, slots.d!!.length)
    }

    /**
     * The target is the longer TRUE record, not the longer stored one, so a
     * vault that grows and then shrinks is not pinned at a high-water mark.
     */
    @Test fun a_fake_decoy_lets_the_pair_follow_the_vault_back_down() {
        val slots = AlignedSlotRecords()
        slots.writeSlotA(slotWithBlob(VaultCrypto.PADDING_BUCKETS.first()))
        slots.writeDecoy(slotWithBlob(VaultCrypto.PADDING_BUCKETS.first()))
        val small = slots.a!!.length

        slots.writeSlotA(slotWithBlob(VaultCrypto.PADDING_BUCKETS.last()))
        assertTrue("grows with the vault", slots.a!!.length > small)
        assertEquals(slots.a!!.length, slots.d!!.length)

        // Unconfigured: a fresh fake is generated at the new, smaller size.
        slots.writeSlotA(slotWithBlob(VaultCrypto.PADDING_BUCKETS.first()))
        slots.writeDecoy(slotWithBlob(VaultCrypto.PADDING_BUCKETS.first()))
        assertEquals("back to the small record length", small, slots.a!!.length)
        assertEquals(slots.a!!.length, slots.d!!.length)
    }

    @Test fun a_slot_record_with_no_sibling_is_never_padded() {
        val record = VaultCodec.encodeSlot(slotWithBlob(64))
        val (aligned, sibling) = VaultCodec.alignSlotRecords(record, null)
        assertEquals(record, aligned)
        assertEquals(null, sibling)
    }

    /**
     * PB-26-002: restoreVaultTriad must align slot A and the decoy to equal
     * lengths, the same way writeAligned does, so a post-restore at-rest
     * observer cannot tell which slot is the real vault by length.
     */
    @Test fun restored_triad_aligns_slot_records_to_equal_lengths() {
        val slotA = slotWithBlob(500)
        val decoy = slotWithBlob(50)
        val encodedA = VaultCodec.encodeSlot(slotA)
        val encodedDecoy = VaultCodec.encodeSlot(decoy)
        assertTrue("test setup: slot A is longer than decoy", encodedA.length > encodedDecoy.length)

        val (alignedA, alignedDecoy) = VaultCodec.alignSlotRecords(encodedA, encodedDecoy)
        assertEquals("restored records must be equal length", alignedA.length, alignedDecoy!!.length)
        assertEquals("slot A content is preserved through padding", slotA, VaultCodec.decodeSlot(alignedA))
        assertEquals("decoy content is preserved through padding", decoy, VaultCodec.decodeSlot(alignedDecoy))
    }
}