package com.peachpasswords.vault

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

/**
 * H-02: a snapshot restore must either fully commit slot A + slot B + version
 * stamp or leave every byte of at-rest state (slot A, slot B, history)
 * untouched. A partial commit strands the session on a vault that no longer
 * matches disk, and a "restore" history entry written before the commit makes
 * a failed restore look like a successful one in the history list.
 */
class SnapshotRestoreAtomicityTest {

    private class FaultStorage : VaultStorage {
        var slotA: StoredVault? = null
        var slotB: String? = null
        var failWriteSlotA = false
        var failWriteSlotB = false
        var failStamp = false
        private var ver: Long = 0
        private var settings = VaultSettings()
        var history = VaultHistory()
        override fun loadSlotA() = slotA
        override fun hasVault() = slotA != null
        override fun writeSlotA(slot: StoredVault): Boolean {
            if (failWriteSlotA) return false
            slotA = slot; return true
        }
        override fun loadSlotB() = slotB
        override fun writeSlotB(blobBase64: String): Boolean {
            if (failWriteSlotB) 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 {
            if (failStamp) throw RuntimeException("stamp failed")
            ver = maxOf(System.currentTimeMillis(), ver + 1); return ver
        }
        override fun loadSettings() = settings
        override fun writeSettings(s: VaultSettings): Boolean { settings = s; return true }
        override fun loadVaultHistory() = history
        override fun writeVaultHistory(history: VaultHistory): Boolean { this.history = history; return true }
        override fun clearVaultHistory() { history = VaultHistory() }
        override fun clearVault(): Boolean { slotA = null; slotB = null; ver = 0; settings = VaultSettings(); history = VaultHistory(); return true }
    }

    /**
     * Build a repo with two entries plus a restorable snapshot. After setup,
     * the on-disk vault holds both entries; the captured snapshot holds only
     * the first entry (takeVaultSnapshot reads slot A from disk BEFORE the new
     * blob is written, so the snapshot from addEntry("Second") captures the
     * First-only state that was on disk just before Second was persisted).
     */
    private data class Fixture(
        val repo: VaultRepository,
        val storage: FaultStorage,
        val snapshotId: String,
    )

    private fun setup(): Fixture = runBlocking {
        val storage = FaultStorage()
        val repo = VaultRepository(storage)
        repo.createVault("pw")
        repo.addEntry(VaultEntry(id = "", title = "First", type = "login", password = "p", createdAt = 0, modifiedAt = 0))
        repo.addEntry(VaultEntry(id = "", title = "Second", type = "login", password = "p", createdAt = 0, modifiedAt = 0))
        val snapshotId = repo.getVaultHistory().first().id
        Fixture(repo, storage, snapshotId)
    }

    private fun loginEntry(title: String) =
        VaultEntry(id = "", title = title, type = "login", password = "p", createdAt = 0, modifiedAt = 0)

    @Test fun happy_path_swaps_slots_locks_and_appends_restore_snapshot() = runBlocking {
        val f = setup()
        val slotABefore = f.storage.slotA
        val slotBBefore = f.storage.slotB
        val historyBefore = f.storage.history.snapshots.toList()

        assertTrue(f.repo.restoreVaultSnapshot(f.snapshotId))

        // Disk now holds the restored pair, not the prior pair.
        assertNotEquals(slotABefore, f.storage.slotA)
        assertNotEquals(slotBBefore, f.storage.slotB)
        // Session is locked; the original "restore" entry was taken BEFORE the
        // commit (it captures the pre-restore state), so it remains in history.
        assertFalse(f.repo.isUnlocked.value)
        assertEquals(historyBefore.size + 1, f.storage.history.snapshots.size)
        assertEquals("restore", f.storage.history.snapshots.first().reason)
        // Vault still opens under the same password and shows the restored entry set.
        assertEquals(UnlockResult.Success, f.repo.unlock("pw"))
        assertEquals(1, f.repo.getEntries().size)
        assertTrue(f.repo.getEntries().any { it.title == "First" })
        assertFalse(f.repo.getEntries().any { it.title == "Second" })
    }

    @Test fun write_slot_a_failure_keeps_slots_and_history_unchanged() = runBlocking {
        val f = setup()
        val slotABefore = f.storage.slotA
        val slotBBefore = f.storage.slotB
        val historyBefore = f.storage.history.snapshots.toList()
        f.storage.failWriteSlotA = true

        assertFalse(f.repo.restoreVaultSnapshot(f.snapshotId))

        assertEquals(slotABefore, f.storage.slotA)
        assertEquals(slotBBefore, f.storage.slotB)
        assertEquals(historyBefore, f.storage.history.snapshots)
        assertTrue(f.storage.history.snapshots.none { it.reason == "restore" })
        // Session is still on the pre-restore key.
        assertTrue(f.repo.isUnlocked.value)
        assertEquals(UnlockResult.Success, f.repo.unlock("pw"))
        assertTrue(f.repo.getEntries().any { it.title == "Second" })
    }

    @Test fun write_slot_b_failure_rolls_back_slot_a_and_history() = runBlocking {
        val f = setup()
        val slotABefore = f.storage.slotA
        val slotBBefore = f.storage.slotB
        val historyBefore = f.storage.history.snapshots.toList()
        f.storage.failWriteSlotB = true

        assertFalse(f.repo.restoreVaultSnapshot(f.snapshotId))

        // The pre-restore pair is back on disk — without the rollback, the at-rest
        // pair would be (restored A, prior B) and an at-rest observer could pick
        // the real slot by size (MN-1).
        assertEquals(slotABefore, f.storage.slotA)
        assertEquals(slotBBefore, f.storage.slotB)
        assertEquals(historyBefore, f.storage.history.snapshots)
        assertTrue(f.storage.history.snapshots.none { it.reason == "restore" })
        // Session still on prior key.
        assertTrue(f.repo.isUnlocked.value)
        assertEquals(UnlockResult.Success, f.repo.unlock("pw"))
        assertTrue(f.repo.getEntries().any { it.title == "Second" })
    }

    @Test fun stamp_failure_rolls_back_both_slots_and_history() = runBlocking {
        val f = setup()
        val slotABefore = f.storage.slotA
        val slotBBefore = f.storage.slotB
        val historyBefore = f.storage.history.snapshots.toList()
        f.storage.failStamp = true

        assertFalse(f.repo.restoreVaultSnapshot(f.snapshotId))

        // Stamp happens AFTER slot B is written; rollback must restore BOTH slots.
        assertEquals(slotABefore, f.storage.slotA)
        assertEquals(slotBBefore, f.storage.slotB)
        assertEquals(historyBefore, f.storage.history.snapshots)
        assertTrue(f.storage.history.snapshots.none { it.reason == "restore" })
        // Vault still on prior key — both entries still readable.
        assertEquals(UnlockResult.Success, f.repo.unlock("pw"))
        assertTrue(f.repo.getEntries().any { it.title == "First" })
        assertTrue(f.repo.getEntries().any { it.title == "Second" })
    }

    @Test fun unknown_snapshot_id_is_a_no_op() = runBlocking {
        val f = setup()
        val slotABefore = f.storage.slotA
        val slotBBefore = f.storage.slotB
        val historyBefore = f.storage.history.snapshots.toList()
        assertTrue(f.repo.isUnlocked.value)

        assertFalse(f.repo.restoreVaultSnapshot("does-not-exist"))

        assertEquals(slotABefore, f.storage.slotA)
        assertEquals(slotBBefore, f.storage.slotB)
        assertEquals(historyBefore, f.storage.history.snapshots)
        assertTrue(f.repo.isUnlocked.value)
    }
}