package com.peachpasswords.vault

import kotlinx.coroutines.runBlocking
import kotlinx.serialization.json.JsonNull
import kotlinx.serialization.json.JsonPrimitive
import kotlinx.serialization.json.buildJsonObject
import org.junit.Assert.assertEquals
import org.junit.Assert.assertTrue
import org.junit.Test
import java.util.Base64
import java.util.concurrent.CountDownLatch
import java.util.concurrent.TimeUnit
import java.util.concurrent.atomic.AtomicReference

/**
 * C-01: a Codex restore must not touch the device's vault until the candidate
 * has been decrypted and decoded in memory, and must roll back if any write in
 * the commit fails. An Android-shaped Codex parses structurally under ANY
 * password, so the old write-then-unlock order wiped the local vault on a typo.
 */
class CodexRestoreAtomicityTest {

    private class FaultStorage : VaultStorage {
        var slotA: StoredVault? = null
        var slotB: String? = null
        var failSlotA = false
        var failSlotB = false
        var failNextSlotB = false
        var failedSlotBEntered: CountDownLatch? = null
        var releaseFailedSlotB: CountDownLatch? = 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 {
            if (failSlotA) return false
            slotA = slot
            return true
        }
        override fun loadSlotB() = slotB
        override fun writeSlotB(blobBase64: String): Boolean {
            if (failNextSlotB) {
                failNextSlotB = false
                failedSlotBEntered?.countDown()
                check(releaseFailedSlotB?.await(5, TimeUnit.SECONDS) != false) {
                    "test writer did not release failed slot-B write"
                }
                return false
            }
            if (failSlotB) return false
            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 }
        override fun clearVault(): Boolean { slotA = null; slotB = null; ver = 0; settings = VaultSettings(); return true }
    }

    private fun sourceCodexBlob(): String = runBlocking {
        val src = VaultRepository(FaultStorage())
        src.createVault("source pass")
        src.addEntry(VaultEntry(id = "", title = "Incoming", type = "login", password = "p", createdAt = 0, modifiedAt = 0))
        src.codexVaultBlobBase64()!!
    }

    /** Re-wrap an Android blob the way the browser writes it: slot as a JSON string. */
    private fun asBrowserShapedWrapper(androidBlob: String): String {
        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 `"a":` — earlier regex matched through to the
        // wrapper's last brace and broke when the duress slot was added.
        val openChar = Regex("\"a\"\\s*:\\s*\\{").find(wrapper)!!.range.last
        val open = openChar + 1
        var depth = 1
        var end = open
        while (end < wrapper.length && depth > 0) {
            when (wrapper[end]) {
                '{' -> depth++
                '}' -> depth--
            }
            if (depth == 0) break
            end++
        }
        val slotA = wrapper.substring(openChar, end + 1)
        val browserWrapper = buildJsonObject {
            put("a", JsonPrimitive(slotA))
            put("b", JsonNull)
        }.toString()
        return Base64.getUrlEncoder().withoutPadding().encodeToString(browserWrapper.toByteArray(Charsets.UTF_8))
    }

    /** An existing local vault holding one entry, plus its storage. */
    private fun existingVault(): Pair<VaultRepository, FaultStorage> = runBlocking {
        val storage = FaultStorage()
        val repo = VaultRepository(storage)
        repo.createVault("local pass")
        repo.addEntry(VaultEntry(id = "", title = "Local", type = "login", password = "p", createdAt = 0, modifiedAt = 0))
        repo.lock()
        repo to storage
    }

    private fun assertLocalVaultIntact(repo: VaultRepository, storage: FaultStorage, before: StoredVault?) =
        runBlocking {
            assertEquals("slot A must be untouched", before, storage.slotA)
            assertEquals(UnlockResult.Success, repo.unlock("local pass"))
            assertTrue(repo.getEntries().any { it.title == "Local" })
        }

    @Test fun wrong_password_on_android_shaped_codex_preserves_local_vault() {
        val blob = sourceCodexBlob()
        val (repo, storage) = existingVault()
        val before = storage.slotA
        runBlocking { assertEquals(UnlockResult.WrongPassword, repo.restoreFromCodex(blob, "not the pass")) }
        assertLocalVaultIntact(repo, storage, before)
    }

    @Test fun wrong_password_on_browser_shaped_codex_preserves_local_vault() {
        val blob = asBrowserShapedWrapper(sourceCodexBlob())
        val (repo, storage) = existingVault()
        val before = storage.slotA
        runBlocking { assertEquals(UnlockResult.WrongPassword, repo.restoreFromCodex(blob, "not the pass")) }
        assertLocalVaultIntact(repo, storage, before)
    }

    @Test fun slot_a_write_failure_preserves_local_vault() {
        val blob = sourceCodexBlob()
        val (repo, storage) = existingVault()
        val before = storage.slotA
        storage.failSlotA = true
        runBlocking {
            assertEquals(UnlockResult.Error("storage_write_failed"), repo.restoreFromCodex(blob, "source pass"))
        }
        storage.failSlotA = false
        assertLocalVaultIntact(repo, storage, before)
    }

    @Test fun later_write_failure_rolls_back_to_local_vault() {
        val blob = sourceCodexBlob()
        val (repo, storage) = existingVault()
        val before = storage.slotA
        storage.failSlotB = true
        runBlocking {
            assertEquals(UnlockResult.Error("storage_write_failed"), repo.restoreFromCodex(blob, "source pass"))
        }
        storage.failSlotB = false
        assertLocalVaultIntact(repo, storage, before)
    }

    @Test
    fun partial_codex_write_holds_epoch_fence_through_rollback() = runBlocking {
        val blob = sourceCodexBlob()
        val (repo, storage) = existingVault()
        val priorA = storage.slotA
        val priorB = storage.slotB
        val priorD = storage.decoy
        val failedSlotBEntered = CountDownLatch(1)
        val releaseFailedSlotB = CountDownLatch(1)
        val rollbackComplete = CountDownLatch(1)
        val releaseAfterRollback = CountDownLatch(1)
        storage.failNextSlotB = true
        storage.failedSlotBEntered = failedSlotBEntered
        storage.releaseFailedSlotB = releaseFailedSlotB
        repo.afterCodexRollback = {
            rollbackComplete.countDown()
            check(releaseAfterRollback.await(5, TimeUnit.SECONDS)) {
                "test did not release post-rollback fence"
            }
        }

        val restoreResult = AtomicReference<UnlockResult>()
        val restoreDone = CountDownLatch(1)
        val restoreThread = Thread {
            restoreResult.set(runBlocking { repo.restoreFromCodex(blob, "source pass") })
            restoreDone.countDown()
        }
        restoreThread.start()
        assertTrue(
            "Codex transaction did not reach the partial slot-B write",
            failedSlotBEntered.await(5, TimeUnit.SECONDS),
        )

        // Start lock while the epoch-held Codex transaction is still blocked in
        // its failed slot-B write. The lock must wait until A/B/D rollback has
        // completed, otherwise it can advance the epoch and strand incoming A.
        val lockStarted = CountDownLatch(1)
        val lockDone = CountDownLatch(1)
        val lockThread = Thread {
            lockStarted.countDown()
            repo.lock()
            lockDone.countDown()
        }
        lockThread.start()
        assertTrue("lock thread did not start", lockStarted.await(5, TimeUnit.SECONDS))

        releaseFailedSlotB.countDown()
        assertTrue("Codex rollback did not finish", rollbackComplete.await(5, TimeUnit.SECONDS))
        assertEquals("rollback must restore prior slot A before lock completes", priorA, storage.slotA)
        assertEquals("rollback must restore prior noise slot B before lock completes", priorB, storage.slotB)
        assertEquals("rollback must restore prior duress slot before lock completes", priorD, storage.decoy)
        assertTrue(
            "lock must remain blocked until the epoch-held rollback returns",
            !lockDone.await(0, TimeUnit.NANOSECONDS),
        )
        releaseAfterRollback.countDown()
        assertTrue("Codex restore did not finish", restoreDone.await(5, TimeUnit.SECONDS))
        assertTrue("lock did not finish", lockDone.await(5, TimeUnit.SECONDS))
        restoreThread.join(5_000)
        lockThread.join(5_000)

        assertEquals(UnlockResult.Error("storage_write_failed"), restoreResult.get())
        assertEquals("rollback must restore prior slot A", priorA, storage.slotA)
        assertEquals("rollback must restore prior noise slot B", priorB, storage.slotB)
        assertEquals("rollback must restore prior duress slot", priorD, storage.decoy)
    }
}