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-01: a wrapping-key rotation (changePassword / resetPasswordWithRecoveryWords)
 * must either fully commit slot A + size-matched noise slot B + version stamp or
 * leave the at-rest state untouched. A partial commit leaves the live session on
 * a key that no longer matches disk, so the next save encrypts under a key that
 * cannot be unlocked next time.
 *
 * The shared fault storage surfaces the three places [commitVaultRotation] can
 * fail: writeSlotA returning false, writeSlotB returning false, and
 * stampLocalMutation throwing. setRecoveryWords is reachable here so the
 * recovery-reset path has a real prior slot A with both escrows.
 */
class PasswordRotationAtomicityTest {

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

    private fun repoWithPassword(storage: FaultStorage, password: String): VaultRepository = runBlocking {
        val repo = VaultRepository(storage)
        repo.createVault(password)
        // Add an entry so we can prove the rotation preserved it (or prove
        // a failed rotation did NOT touch it).
        repo.addEntry(VaultEntry(id = "", title = "Keep me", type = "login", password = "p",
            createdAt = 0, modifiedAt = 0))
        // Confirm recovery words so resetPasswordWithRecoveryWords has a real
        // escrow to rebuild against.
        val words = rootKeyToWords(repo.syncMasterKey()!!)
        assertTrue(repo.setRecoveryWords(words))
        repo
    }

    private fun recoveryWordsOf(repo: VaultRepository): List<String> {
        val key = repo.syncMasterKey()!!
        return try { rootKeyToWords(key) } finally { VaultCrypto.zero(key) }
    }

    // ---- changePassword ---------------------------------------------------

    @Test fun change_password_with_slot_b_failure_keeps_old_key() = runBlocking {
        val storage = FaultStorage()
        val repo = repoWithPassword(storage, "old-pass")
        val slotABefore = storage.slotA
        val slotBBefore = storage.slotB
        storage.failWriteSlotB = true

        // changePassword must return false AND leave the at-rest state intact:
        // the session is still on the old key, slot A on disk still opens under
        // the old password.
        assertFalse(repo.changePassword("old-pass", "new-pass"))

        // The half-written commit was rolled back: both slots match what was
        // there before the call, not the would-be new values.
        assertEquals(slotABefore, storage.slotA)
        assertEquals(slotBBefore, storage.slotB)
        // Round-trip under the OLD password still works — session key untouched.
        repo.lock()
        assertEquals(UnlockResult.Success, repo.unlock("old-pass"))
        assertEquals(1, repo.getEntries().size)
        assertTrue(repo.getEntries().any { it.title == "Keep me" })
        assertEquals(UnlockResult.WrongPassword, repo.unlock("new-pass"))
    }

    @Test fun change_password_happy_path_swaps_session_key_and_disk() = runBlocking {
        val storage = FaultStorage()
        val repo = repoWithPassword(storage, "old-pass")
        val slotABefore = storage.slotA!!

        assertTrue(repo.changePassword("old-pass", "new-pass"))

        // Slot A on disk is now sealed under the new salt/key, not the old one.
        assertNotEquals(slotABefore.salt, storage.slotA?.salt)
        assertNotEquals(slotABefore.blob, storage.slotA?.blob)
        repo.lock()
        assertEquals(UnlockResult.WrongPassword, repo.unlock("old-pass"))
        assertEquals(UnlockResult.Success, repo.unlock("new-pass"))
        assertTrue(repo.getEntries().any { it.title == "Keep me" })
    }

    // ---- resetPasswordWithRecoveryWords ----------------------------------

    @Test fun reset_password_with_slot_b_failure_restores_prior_slot_a() = runBlocking {
        val storage = FaultStorage()
        val repo = repoWithPassword(storage, "first-pass")
        val words = recoveryWordsOf(repo)
        repo.lock()

        assertTrue(repo.unlockWithRecoveryWords(words))
        val slotABefore = storage.slotA
        val slotBBefore = storage.slotB
        storage.failWriteSlotB = true

        assertFalse(repo.resetPasswordWithRecoveryWords(words, "second-pass"))

        // The half-written commit was rolled back; session has no derivedKey
        // yet (unlockWithRecoveryWords leaves it null until the reset succeeds).
        assertEquals(slotABefore, storage.slotA)
        assertEquals(slotBBefore, storage.slotB)
        // A subsequent reset under a healthy store must still work, proving
        // the failed rotation didn't strand the recovery session.
        storage.failWriteSlotB = false
        assertTrue(repo.resetPasswordWithRecoveryWords(words, "second-pass"))
        repo.lock()
        assertEquals(UnlockResult.Success, repo.unlock("second-pass"))
        assertTrue(repo.getEntries().any { it.title == "Keep me" })
    }

    @Test fun reset_password_with_stamp_failure_restores_both_slots() = runBlocking {
        val storage = FaultStorage()
        val repo = repoWithPassword(storage, "first-pass")
        val words = recoveryWordsOf(repo)
        repo.lock()

        assertTrue(repo.unlockWithRecoveryWords(words))
        val slotABefore = storage.slotA
        val slotBBefore = storage.slotB
        storage.failStamp = true

        assertFalse(repo.resetPasswordWithRecoveryWords(words, "second-pass"))

        // H-01: stamp failure happens AFTER slot B is written, so rollback
        // must restore BOTH slots — otherwise the at-rest pair is size-mismatched.
        assertEquals(slotABefore, storage.slotA)
        assertEquals(slotBBefore, storage.slotB)
        // Session state from unlockWithRecoveryWords is intact (entries loaded,
        // root key present, no derivedKey yet) so a clean retry can complete.
        assertTrue(repo.getEntries().any { it.title == "Keep me" })
        // A clean retry succeeds.
        storage.failStamp = false
        assertTrue(repo.resetPasswordWithRecoveryWords(words, "second-pass"))
        repo.lock()
        assertEquals(UnlockResult.Success, repo.unlock("second-pass"))
    }
}