package com.peachpasswords.vault

import kotlinx.coroutines.runBlocking
import kotlinx.serialization.json.Json
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

/**
 * End-to-end Codex round-trip tests for the duress (slot B / decoy) slot.
 * Proves that a browser-shaped Codex carrying a real duress "b" slot — sealed
 * under the duress passphrase, not the main password — restores correctly on
 * Android: the real password still opens slot A, the duress passphrase opens
 * slot B, and a wrong decoy password fails closed.
 */
class DuressCodexRoundTripTest {

    private class MemStorage : 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 loadSlotB() = slotB
        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 clearVaultHistory() { history = VaultHistory() }
        override fun clearVault(): Boolean { slotA = null; slotB = null; decoy = null; ver = 0; settings = VaultSettings(); history = VaultHistory(); return true }
    }

    private val mainPassword = "correct-horse-battery-staple"
    private val duressPassword = "open-the-decoy-vault-456"
    private val wrongPassword = "totally-wrong-password"

    private fun hex(bytes: ByteArray): String =
        bytes.joinToString("") { "%02x".format(it) }

    /**
     * Convert an Android StoredVault into a browser-shaped BrowserSlotA JSON
     * string (hex salt/nonce/ciphertext/tag). This simulates what the browser
     * writes into a Codex — the encrypted content is preserved as-is, only the
     * encoding changes from Android's base64(salt)+base64(nonce‖ct‖tag) to the
     * browser's separate hex fields.
     */
    private fun androidSlotToBrowserJson(slot: StoredVault): String {
        val salt = Base64.getDecoder().decode(slot.salt)
        val blob = Base64.getDecoder().decode(slot.blob)
        val nonce = blob.copyOfRange(0, VaultCrypto.NONCE_SIZE)
        val ctAndTag = blob.copyOfRange(VaultCrypto.NONCE_SIZE, blob.size)
        // AES-GCM tag is the last 16 bytes of the ciphertext+tag blob.
        val ct = ctAndTag.copyOfRange(0, ctAndTag.size - 16)
        val tag = ctAndTag.copyOfRange(ctAndTag.size - 16, ctAndTag.size)
        val browserSlot = BrowserSlotA(
            version = slot.version,
            salt = hex(salt),
            nonce = hex(nonce),
            ciphertext = hex(ct),
            tag = hex(tag),
        )
        return Json.encodeToString(BrowserSlotA.serializer(), browserSlot)
    }

    /**
     * Build a browser-shaped Codex from an Android source vault. Both "a" and
     * "b" are written as JSON strings (the browser's format), with the slot
     * content in browser hex-field shape.
     */
    private fun buildBrowserCodex(src: VaultRepository): String {
        val androidBlob = src.codexVaultBlobBase64()!!
        val wrapper = String(Base64.getUrlDecoder().decode(androidBlob), Charsets.UTF_8)
        val parsed = Json.parseToJsonElement(wrapper).asJsonObject()
        val slotAJson = parsed["a"].toString()
        val slotBJon = parsed["b"].toString()
        val slotA = VaultCodec.decodeSlot(slotAJson)!!
        val slotB = VaultCodec.decodeSlot(slotBJon)!!
        val browserA = androidSlotToBrowserJson(slotA)
        val browserB = androidSlotToBrowserJson(slotB)
        val browserWrapper = buildJsonObject {
            put("a", JsonPrimitive(browserA))
            put("b", JsonPrimitive(browserB))
        }.toString()
        return Base64.getUrlEncoder().withoutPadding()
            .encodeToString(browserWrapper.toByteArray(Charsets.UTF_8))
    }

    private fun kotlinx.serialization.json.JsonElement.asJsonObject() =
        this as kotlinx.serialization.json.JsonObject

    /**
     * Full round-trip: build a vault with duress on "device A", generate a
     * browser-shaped Codex, restore on "device B", and verify all unlock paths.
     */
    @Test fun browser_codex_with_duress_restores_and_unlocks_both_slots() = runBlocking {
        // ── Device A: create vault, add entries, configure duress ──────────
        val srcStorage = MemStorage()
        val src = VaultRepository(srcStorage)
        src.createVault(mainPassword)
        src.addEntry(VaultEntry(id = "", title = "RealEntry", type = "login", password = "secret", createdAt = 0, modifiedAt = 0))
        assertTrue("duress passphrase configured", src.setDuressPassphrase(duressPassword))
        // Verify source duress works before export.
        src.lock()
        val srcDuress = src.tryUnlockDuress(duressPassword)
        assertNotNull("source duress unlock works", srcDuress)
        src.discardDuressUnlock(srcDuress!!)
        assertEquals(UnlockResult.Success, src.unlock(mainPassword))

        // ── Build browser-shaped Codex ─────────────────────────────────────
        val browserCodex = buildBrowserCodex(src)

        // ── Device B: restore from browser Codex ──────────────────────────
        val dstStorage = MemStorage()
        val dst = VaultRepository(dstStorage)

        // (a) Real password unlocks slot A and entries are present.
        assertEquals("real password restores slot A",
            UnlockResult.Success, dst.restoreFromCodex(browserCodex, mainPassword))
        assertTrue("entries present after restore", dst.getEntries().any { it.title == "RealEntry" })

        // Lock so we can test duress unlock independently.
        dst.lock()

        // (b) Duress password unlocks slot B.
        val duress = dst.tryUnlockDuress(duressPassword)
        assertNotNull("duress passphrase unlocks the restored duress slot", duress)
        assertEquals(UnlockResult.Success, dst.commitDuressUnlock(duress!!))
        assertTrue("decoy session is active", dst.isUnlocked.value)
        assertEquals("decoy vault has no real entries", 0, dst.getEntries().size)

        // (c) Wrong decoy password fails.
        dst.lock()
        assertNull("wrong password does NOT unlock duress slot",
            dst.tryUnlockDuress(wrongPassword))
        assertNull("wrong password does NOT unlock slot A either",
            dst.tryUnlockSlotA(wrongPassword))
    }

    /**
     * When a browser Codex has no "b" slot (null), a size-matched fake duress
     * must be written so the at-rest record stays forensically indistinguishable.
     */
    @Test fun no_duress_codex_produces_size_matched_fake() = runBlocking {
        // Source: vault without duress configured.
        val srcStorage = MemStorage()
        val src = VaultRepository(srcStorage)
        src.createVault(mainPassword)
        src.addEntry(VaultEntry(id = "", title = "OnlyReal", type = "login", password = "pw", createdAt = 0, modifiedAt = 0))

        // Build a browser Codex with "b" as null.
        val androidBlob = src.codexVaultBlobBase64()!!
        val wrapper = String(Base64.getUrlDecoder().decode(androidBlob), Charsets.UTF_8)
        val parsed = Json.parseToJsonElement(wrapper).asJsonObject()
        val slotAJson = parsed["a"].toString()
        val slotA = VaultCodec.decodeSlot(slotAJson)!!
        val browserA = androidSlotToBrowserJson(slotA)
        val browserWrapper = buildJsonObject {
            put("a", JsonPrimitive(browserA))
            put("b", kotlinx.serialization.json.JsonNull)
        }.toString()
        val codex = Base64.getUrlEncoder().withoutPadding()
            .encodeToString(browserWrapper.toByteArray(Charsets.UTF_8))

        // Restore.
        val dstStorage = MemStorage()
        val dst = VaultRepository(dstStorage)
        assertEquals(UnlockResult.Success, dst.restoreFromCodex(codex, mainPassword))

        // Fake duress slot exists and is size-matched to slot A.
        assertNotNull("fake duress slot written", dstStorage.decoy)
        val slotABlobLen = Base64.getDecoder().decode(dstStorage.slotA!!.blob).size
        val duressBlobLen = Base64.getDecoder().decode(dstStorage.decoy!!.blob).size
        assertEquals("fake duress blob is size-matched to slot A", slotABlobLen, duressBlobLen)

        // Duress unlock fails (it's a fake).
        dst.lock()
        assertNull("fake duress cannot be unlocked", dst.tryUnlockDuress(mainPassword))
        assertNull("fake duress cannot be unlocked with any password", dst.tryUnlockDuress("anything"))
    }

    /**
     * Android-to-Android Codex round-trip with duress: the "b" slot is in
     * Android format (StoredVault JSON) and should pass through VaultCodec.decodeSlot
     * without needing the browser passthrough path.
     */
    @Test fun android_codex_with_duress_round_trips() = runBlocking {
        val srcStorage = MemStorage()
        val src = VaultRepository(srcStorage)
        src.createVault(mainPassword)
        src.addEntry(VaultEntry(id = "", title = "AndroidOnly", type = "login", password = "pw", createdAt = 0, modifiedAt = 0))
        assertTrue(src.setDuressPassphrase(duressPassword))

        val codex = src.codexVaultBlobBase64()!!

        val dstStorage = MemStorage()
        val dst = VaultRepository(dstStorage)
        assertEquals(UnlockResult.Success, dst.restoreFromCodex(codex, mainPassword))
        assertTrue(dst.getEntries().any { it.title == "AndroidOnly" })

        dst.lock()
        val duress = dst.tryUnlockDuress(duressPassword)
        assertNotNull("Android-to-Android duress round-trip works", duress)
        assertEquals(UnlockResult.Success, dst.commitDuressUnlock(duress!!))
        assertEquals(0, dst.getEntries().size)
    }
}