package com.peachpasswords.vault

import kotlinx.coroutines.runBlocking
import kotlinx.serialization.json.JsonPrimitive
import kotlinx.serialization.json.buildJsonObject
import org.junit.Assert.assertEquals
import org.junit.Assert.assertFalse
import org.junit.Assert.assertNotEquals
import org.junit.Assert.assertNotNull
import org.junit.Assert.assertNull
import org.junit.Assert.assertTrue
import org.junit.Test
import java.util.Base64

/** In-memory [VaultStorage] including a decoy (slot B) slot, for the MF-6 duress model. */
private class DuressFakeStorage : VaultStorage {
    var slotA: StoredVault? = null
    var slotB: String? = 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 { slotB = blobBase64; return true }
    var decoy: StoredVault? = null
    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 }
    var history = VaultHistory()
    override fun loadVaultHistory() = history
    override fun writeVaultHistory(h: VaultHistory): Boolean { history = h; return true }
    override fun clearVault() { slotA = null; slotB = null; decoy = null; ver = 0; settings = VaultSettings(); history = VaultHistory() }
}

/**
 * MF-6: the duress passphrase keys a real second vault (slot B), sealed exactly
 * like slot A. At the lock screen, both slots are tried and the outcome is
 * invisible from outside — no sentinel, no separate hash.
 */
class DuressSlotBTest {
    private val real = "the-real-password-123"
    private val duressPass = "open-the-decoy-456"

    private fun configuredRepo(): VaultRepository = runBlocking {
        val repo = VaultRepository(DuressFakeStorage())
        repo.createVault(real)
        repo.addEntry(VaultEntry(id = "", title = "Secret", type = "login", password = "pw", createdAt = 0, modifiedAt = 0))
        assertTrue("duress passphrase sealed into slot B", repo.setDuressPassphrase(duressPass))
        repo.lock()
        repo
    }

    @Test fun real_password_opens_slot_a_only() = runBlocking {
        val repo = configuredRepo()
        assertNotNull("real password opens slot A", repo.tryUnlockSlotA(real))
        assertNull("real password does NOT open the duress slot", repo.tryUnlockDuress(real))
    }

    @Test fun duress_passphrase_opens_decoy_only() = runBlocking {
        val repo = configuredRepo()
        assertNull("duress passphrase does NOT open slot A", repo.tryUnlockSlotA(duressPass))
        val duress = repo.tryUnlockDuress(duressPass)
        assertNotNull("duress passphrase opens the duress slot", duress)
        assertEquals(UnlockResult.Success, repo.commitDuressUnlock(duress!!))
        assertTrue(repo.isUnlocked.value)
        assertEquals("decoy vault is empty — no real entries leak", 0, repo.getEntries().size)
    }

    @Test fun sibling_duress_tokens_cannot_replace_the_adopted_decoy_session() = runBlocking {
        val repo = configuredRepo()
        val first = repo.tryUnlockDuress(duressPass)!!
        val sibling = repo.tryUnlockDuress(duressPass)!!
        assertEquals(UnlockResult.Success, repo.commitDuressUnlock(first))
        assertEquals(UnlockResult.Error("stale_unlock"), repo.commitDuressUnlock(sibling))
        assertTrue(repo.isUnlocked.value)
    }

    @Test fun wrong_password_opens_neither() = runBlocking {
        val repo = configuredRepo()
        assertNull(repo.tryUnlockSlotA("totally-wrong"))
        assertNull(repo.tryUnlockDuress("totally-wrong"))
    }

    @Test fun duress_passphrase_cannot_equal_real_password() = runBlocking {
        val repo = VaultRepository(DuressFakeStorage())
        repo.createVault(real)
        assertFalse("a duress passphrase matching the real password is rejected", repo.setDuressPassphrase(real))
    }

    @Test fun clear_duress_replaces_with_fake_of_matching_size() = runBlocking {
        // H-04: presence-indistinguishable — clearDuress no longer removes the
        // duress record, it replaces it with a fake blob the duress passphrase
        // can't open. The at-rest slot pair keeps the same shape.
        val storage = DuressFakeStorage()
        val repo = VaultRepository(storage)
        repo.createVault(real)
        repo.addEntry(VaultEntry(id = "", title = "Real", type = "login", password = "pw", createdAt = 0, modifiedAt = 0))
        assertTrue(repo.setDuressPassphrase(duressPass))
        assertEquals(UnlockResult.Success, repo.unlock(real))
        val slotALen = storage.slotA!!.blob.length
        val decoyBefore = storage.decoy!!
        assertEquals(slotALen, decoyBefore.blob.length)
        assertTrue(repo.clearDuress())
        assertNotNull("duress record still present after clear", storage.decoy)
        // The fake has the same blob length as slot A — forensic size tell closed.
        assertEquals(slotALen, storage.decoy!!.blob.length)
        // …and the duress passphrase no longer opens it.
        repo.lock()
        assertNull(repo.tryUnlockDuress(duressPass))
    }

    @Test fun `HOLE 5 the duress vault has no sync identity and cannot sync`() = runBlocking {
        val repo = configuredRepo()

        // Baseline: the real vault does have a sync identity.
        assertNotNull(repo.tryUnlockSlotA(real))
        assertEquals(UnlockResult.Success, repo.unlock(real))
        val realVaultId = repo.syncVaultId()
        assertNotNull("the real vault syncs", realVaultId)
        repo.lock()

        // The decoy session must expose neither a sync key nor a vault id.
        // A second vault_id would be a duress indicator: a coerced user opening
        // the decoy would make sync visibly fail on their other devices.
        val duress = repo.tryUnlockDuress(duressPass)
        assertNotNull(duress)
        assertEquals(UnlockResult.Success, repo.commitDuressUnlock(duress!!))
        assertTrue(repo.isUnlocked.value)
        assertNull("the duress vault must expose no sync key", repo.syncMasterKey())
        assertNull("the duress vault must expose no vault id", repo.syncVaultId())

        // Every sync adapter operation therefore fails closed rather than
        // syncing, manifesting, or revealing an entry count.
        val adapter = com.peachpasswords.sync.VaultSyncAdapter(repo, ourDeviceId = "decoy-device")
        var threw = false
        try {
            adapter.computeVaultId()
        } catch (_: Exception) {
            threw = true
        }
        assertTrue("the decoy vault must not produce a vault id for sync", threw)
    }

    // ---- H-04: duress storage forensics ----------------------------------

    /**
     * H-04: the duress slot's stored blob must match slot A's stored blob
     * length on disk — both when a real decoy is configured (padded at
     * setDuressPassphrase) and when no decoy is configured (a fake of slot A's
     * size is left behind so the duress record's PRESENCE on disk is
     * independent of duressConfigured).
     */
    @Test fun h04_duress_record_size_matches_slot_a() = runBlocking {
        // Empty vault, no duress configured.
        val storage = DuressFakeStorage()
        val repo = VaultRepository(storage)
        repo.createVault(real)
        val slotALen = storage.slotA!!.blob.length
        assertNotNull("duress record exists even with no duress configured (H-04)", storage.decoy)
        assertEquals("fake duress blob length matches slot A blob length", slotALen, storage.decoy!!.blob.length)

        // With a real duress configured — blob length is still matched.
        assertTrue(repo.setDuressPassphrase(duressPass))
        assertEquals("real duress blob length matches slot A blob length", slotALen, storage.decoy!!.blob.length)

        // Slot A grows (entries added): the duress record is RE-PADDED on
        // every persist (the real decoy can't be re-encrypted without its
        // passphrase, so it's left alone — see the setDuressPassphrase
        // ponytail note for the residual tell). A fake duress IS re-padded
        // so it stays indistinguishable from the noise slot.
        repo.addEntry(VaultEntry(id = "", title = "Big", type = "login", password = "pw", createdAt = 0, modifiedAt = 0))
        repo.clearDuress()
        val newSlotALen = storage.slotA!!.blob.length
        assertEquals("fake duress re-padded to slot A after slot A grows", newSlotALen, storage.decoy!!.blob.length)
    }

    /**
     * H-04: a snapshot must capture the duress slot, not just the noise slot.
     * Without this, restoring from history overwrites slot A (the real vault)
     * while the duress slot stays on whatever the previous state was — a
     * forensic tell across restores.
     */
    @Test fun h04_snapshot_contains_duress() = runBlocking {
        val storage = DuressFakeStorage()
        val repo = VaultRepository(storage)
        repo.createVault(real)
        val initialDuress = storage.decoy
        assertNotNull(initialDuress)
        // The first snapshot is taken on every persist — grab it.
        repo.addEntry(VaultEntry(id = "", title = "T", type = "login", password = "pw", createdAt = 0, modifiedAt = 0))
        val snap = storage.history.snapshots.first()
        assertNotNull("snapshot.vault_d captures the duress slot", snap.vault_d)
        assertEquals("snapshotted duress matches the on-disk duress", initialDuress!!.blob, VaultCodec.decodeSlot(snap.vault_d!!)!!.blob)
        // Restoring it writes the captured duress back to disk (atomic).
        assertTrue(repo.restoreVaultSnapshot(snap.id))
        assertNotNull(storage.decoy)
        assertEquals("restored duress matches the snapshot", VaultCodec.encodeSlot(initialDuress), snap.vault_d)
    }

    /**
     * H-04: a Codex restore with a browser-shaped "b" (the browser's duress
     * slot, written there because on the browser slot B IS the duress slot)
     * must land in Android's duress slot, not the noise slot. Writing it to
     * the noise slot would mean the duress record disappears on every Codex
     * restore from a paired browser — another forensic tell.
     */
    @Test fun h04_codex_restore_puts_browser_b_into_duress_slot() = runBlocking {
        // Source vault: a real Android duress exists.
        val srcStorage = DuressFakeStorage()
        val src = VaultRepository(srcStorage)
        src.createVault(real)
        src.addEntry(VaultEntry(id = "", title = "Real", type = "login", password = "pw", createdAt = 0, modifiedAt = 0))
        assertTrue(src.setDuressPassphrase(duressPass))
        val srcDuress = srcStorage.decoy!!

        // Build a browser-shaped Codex: "b" holds the duress slot as a JSON
        // STRING (the browser writes each slot as a string, not an object),
        // matching the browser's A/B naming where B is the duress slot.
        val androidBlob = src.codexVaultBlobBase64()!!
        val wrapper = String(Base64.getUrlDecoder().decode(androidBlob), Charsets.UTF_8)
        // H-04: codex now carries "a" and "b" as siblings — match the first
        // balanced `{…}` after each key, not the wrapper's last brace.
        val slotAJson = balancedBraceAfter(wrapper, "\"a\"")
        val slotDJson = balancedBraceAfter(wrapper, "\"b\"")
        val browserWrapper = buildJsonObject {
            put("a", JsonPrimitive(slotAJson))
            put("b", JsonPrimitive(slotDJson))
        }.toString()
        val browserCodex = Base64.getUrlEncoder().withoutPadding()
            .encodeToString(browserWrapper.toByteArray(Charsets.UTF_8))

        // Restore into a fresh device that already has a vault.
        val dstStorage: DuressFakeStorage = DuressFakeStorage()
        val dst: VaultRepository = VaultRepository(dstStorage)
        dst.createVault(real)

        assertEquals(UnlockResult.Success, dst.restoreFromCodex(browserCodex, real))

        // The duress slot now carries the Codex's "b" — not a fake, not noise.
        assertNotNull("duress slot populated from Codex 'b'", dstStorage.decoy)
        assertEquals("restored duress blob matches the source duress", srcDuress.blob, dstStorage.decoy!!.blob)
        // The noise slot is rewritten to match slot A's size (H-04 size-match
        // invariant) — assert the cross-contamination invariant: noise and
        // duress slots are different.
        assertNotEquals("noise slot is not the duress slot", dstStorage.slotB, dstStorage.decoy?.blob)
    }

    /**
     * Match the first balanced `{…}` whose opening `{` follows the given key
     * literal. Used after H-04 to extract slot A or B from a wrapper that now
     * carries both slots side by side (regex `\{.*\}` would match through to
     * the wrapper's last brace).
     */
    private fun balancedBraceAfter(wrapper: String, keyLiteral: String): String {
        val openMarker = "${Regex.escape(keyLiteral)}\\s*:\\s*\\{".toRegex().find(wrapper)
            ?: error("key $keyLiteral not found")
        val open = openMarker.range.last
        var depth = 1
        var end = open + 1
        while (end < wrapper.length && depth > 0) {
            when (wrapper[end]) {
                '{' -> depth++
                '}' -> depth--
            }
            if (depth == 0) break
            end++
        }
        return wrapper.substring(open, end + 1)
    }
}