package com.peachpasswords.vault

import kotlinx.coroutines.runBlocking
import org.junit.Assert.assertEquals
import org.junit.Assert.assertFalse
import org.junit.Assert.assertNotNull
import org.junit.Assert.assertNull
import org.junit.Assert.assertTrue
import org.junit.Test
import java.util.concurrent.CountDownLatch
import java.util.concurrent.TimeUnit
import java.util.concurrent.atomic.AtomicReference

/**
 * M-04: no operation may leave in-memory session state that disagrees with what
 * is on disk, and the sync-key accessors must fail closed whenever the session
 * is not a fully-unlocked real vault.
 */
class HiddenSessionStateTest {

    private class FaultStorage : VaultStorage {
        var slotA: StoredVault? = null
        var slotB: String? = null
        var failWriteSlotA = false
        private var ver: Long = 0
        var settings = VaultSettings()
        var failLoadSettings = false
        var failWriteSettings = false
        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 { 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(): VaultSettings {
            if (failLoadSettings) throw IllegalStateException("settings_read_failed")
            return settings
        }
        override fun writeSettings(s: VaultSettings): Boolean {
            if (failWriteSettings) return false
            settings = s
            return true
        }
        override fun clearVault(): Boolean { slotA = null; slotB = null; ver = 0; settings = VaultSettings(); return true }
    }

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

    // A failed escrow write must not leave the session claiming words are set:
    // the recovery screen would stop offering the confirm step against a slot A
    // that holds no escrow at all.
    @Test fun failed_escrow_write_restores_prior_recovery_flags() = runBlocking {
        val storage = FaultStorage()
        val repo = VaultRepository(storage)
        repo.createVault("pw")
        val words = recoveryWordsOf(repo)
        assertFalse(repo.isRecoveryWordsSet())

        storage.failWriteSlotA = true
        assertFalse(repo.setRecoveryWords(words))
        assertFalse("flags must roll back when the escrow never landed", repo.isRecoveryWordsSet())
        assertFalse(repo.verifyRecoveryWords(words))
        assertFalse(repo.hasRecoveryEscrow())

        // And the retry still works once storage recovers.
        storage.failWriteSlotA = false
        assertTrue(repo.setRecoveryWords(words))
        assertTrue(repo.isRecoveryWordsSet())
        assertTrue(repo.verifyRecoveryWords(words))
        assertTrue(repo.hasRecoveryEscrow())
    }

    @Test fun sync_accessors_are_null_when_locked() = runBlocking {
        val repo = VaultRepository(FaultStorage())
        repo.createVault("pw")
        repo.addEntry(
            VaultEntry(
                id = "secret",
                title = "Secret",
                type = "login",
                username = "u",
                password = "p",
                createdAt = 0,
                modifiedAt = 0,
            ),
        )
        assertNotNull(repo.syncMasterKey())
        assertNotNull(repo.syncVaultId())
        repo.lock()
        assertNull(repo.syncMasterKey())
        assertNull(repo.syncVaultId())
        assertTrue(repo.getEntries().isEmpty())
        assertNull(repo.getEntry("secret"))
        assertTrue(repo.getDeletedEntries().isEmpty())
        assertTrue(repo.getActiveEntries().isEmpty())
        assertTrue(repo.getDistinctTags().isEmpty())
        assertTrue(repo.getDistinctFieldValues("username").isEmpty())
        assertTrue(repo.getTombstones().isEmpty())
        assertNull(repo.currentVaultKey())
    }

    // A recovery unlock installs a root key but no password key — the session is
    // not unlocked yet, so nothing may leave the device under that identity.
    @Test fun sync_accessors_are_null_during_recovery_half_unlock() = runBlocking {
        val storage = FaultStorage()
        val repo = VaultRepository(storage)
        repo.createVault("pw")
        val words = recoveryWordsOf(repo)
        assertTrue(repo.setRecoveryWords(words))
        repo.lock()

        assertTrue(repo.unlockWithRecoveryWords(words))
        assertFalse(repo.isUnlocked.value)
        assertNull(repo.syncMasterKey())
        assertNull(repo.syncVaultId())
        assertNotNull("the local-only accessor still sees the recovered key", repo.sessionRootKey())

        // Finishing the reset unlocks the session and the sync identity returns.
        assertTrue(repo.resetPasswordWithRecoveryWords(words, "new-pw"))
        assertTrue(repo.isUnlocked.value)
        assertNotNull(repo.syncMasterKey())
        assertNotNull(repo.syncVaultId())
    }

    @Test fun failed_create_retains_no_session_state() = runBlocking {
        val storage = FaultStorage()
        storage.failWriteSlotA = true
        val repo = VaultRepository(storage)

        assertTrue(repo.createVault("pw") is CreateVaultResult.Error)
        assertFalse(repo.isUnlocked.value)
        assertNull("a failed create must not retain a root key", repo.sessionRootKey())
        assertNull(repo.syncMasterKey())
        assertNull(repo.currentVaultKey())
    }

    @Test fun locked_settings_are_safe_default_and_cannot_be_mutated() = runBlocking {
        val storage = FaultStorage()
        val repo = VaultRepository(storage)
        assertTrue(repo.createVault("pw") is CreateVaultResult.Success)
        val requested = VaultSettings(accentColor = "blue", autoLockMinutes = 5)
        repo.lock()

        var failed = false
        try {
            repo.updateSettings(requested)
        } catch (_: IllegalStateException) {
            failed = true
        }
        assertTrue("locked settings writes must fail", failed)
        assertEquals(VaultSettings(), repo.settings.value)
        assertEquals("locked writes must not reach storage", VaultSettings(), storage.settings)
    }

    @Test fun stale_settings_write_cannot_publish_or_persist_after_lock() = runBlocking {
        val storage = FaultStorage()
        val repo = VaultRepository(storage)
        assertTrue(repo.createVault("pw") is CreateVaultResult.Success)
        val requested = VaultSettings(accentColor = "blue", autoLockMinutes = 5)
        val paused = CountDownLatch(1)
        val release = CountDownLatch(1)
        val failure = AtomicReference<Throwable?>()
        repo.beforeDurableCommit = {
            paused.countDown()
            check(release.await(5, TimeUnit.SECONDS)) { "settings writer did not release" }
        }
        val writer = Thread {
            try {
                runBlocking { repo.updateSettings(requested) }
            } catch (t: Throwable) {
                failure.set(t)
            }
        }
        writer.start()
        assertTrue("settings writer did not reach durable seam", paused.await(5, TimeUnit.SECONDS))
        repo.lock()
        assertEquals(VaultSettings(), repo.settings.value)
        release.countDown()
        writer.join(5_000)
        assertFalse(writer.isAlive)
        assertTrue("stale settings writer must fail closed", failure.get() is IllegalStateException)
        assertEquals("stale settings must not persist", VaultSettings(), storage.settings)
    }

    @Test fun failed_full_recovery_settings_load_leaves_no_candidate_session() = runBlocking {
        val storage = FaultStorage()
        val repo = VaultRepository(storage)
        assertTrue(repo.createVault("pw") is CreateVaultResult.Success)
        repo.addEntry(VaultEntry(id = "secret", title = "Secret", type = "login", password = "p", createdAt = 0, modifiedAt = 0))
        val key = repo.sessionRootKey()!!
        val words = try { rootKeyToWords(key) } finally { VaultCrypto.zero(key) }
        assertTrue(repo.setRecoveryWords(words))
        repo.lock()
        storage.failLoadSettings = true

        assertFalse(repo.unlockWithRecoveryWords(words))
        assertFalse(repo.isUnlocked.value)
        assertNull(repo.sessionRootKey())
        assertTrue(repo.getEntries().isEmpty())
        assertFalse(repo.isRecoveryWordsSet())
        assertEquals(VaultSettings(), repo.settings.value)
    }

    @Test fun invalid_recovery_root_key_leaves_no_candidate_session() = runBlocking {
        val storage = FaultStorage()
        val repo = VaultRepository(storage)
        assertTrue(repo.createVault("pw") is CreateVaultResult.Success)
        val root = repo.sessionRootKey()!!
        val words = try { rootKeyToWords(root) } finally { VaultCrypto.zero(root) }

        // Build an authenticated Android recovery escrow whose payload decodes but
        // carries a one-byte root key. The failure must not publish a half-adopted
        // candidate, and the implementation must scrub the rejected decoded bytes.
        val material = wordsToRootKeyMaterial(words)
        val salt = ByteArray(VaultCrypto.SALT_SIZE) { it.toByte() }
        val recoveryEscrow = try {
            val recoveryKey = deriveRecoveryKey(material, salt)
            val keyBytes = recoveryKey.encoded
            try {
                val escrowPlaintext = VaultCodec.encode(VaultPayload(rootKeyB64 = "AQ"))
                try {
                    val sealed = VaultCrypto.encrypt(escrowPlaintext, keyBytes, VaultCrypto.AAD_RECOVERY)
                    try {
                        RecoveryEscrow(
                            salt = java.util.Base64.getEncoder().withoutPadding().encodeToString(salt),
                            blob = java.util.Base64.getEncoder().withoutPadding().encodeToString(sealed),
                        )
                    } finally {
                        VaultCrypto.zero(sealed)
                    }
                } finally {
                    VaultCrypto.zero(escrowPlaintext)
                }
            } finally {
                VaultCrypto.zero(keyBytes)
            }
        } finally {
            VaultCrypto.zero(material)
        }
        storage.slotA = storage.slotA!!.copy(recoveryEscrow = recoveryEscrow, rootKeyEscrow = null)
        repo.lock()

        assertFalse(repo.unlockWithRecoveryWords(words))
        assertFalse(repo.isUnlocked.value)
        assertNull("invalid root key must not become a hidden session key", repo.sessionRootKey())
        assertTrue(repo.getEntries().isEmpty())
        assertFalse(repo.isRecoveryWordsSet())
        assertEquals(VaultSettings(), repo.settings.value)
    }

    @Test fun failed_root_only_recovery_settings_load_leaves_no_candidate_session() = runBlocking {
        val storage = FaultStorage()
        val repo = VaultRepository(storage)
        assertTrue(repo.createVault("pw") is CreateVaultResult.Success)
        val key = repo.sessionRootKey()!!
        val words = try { rootKeyToWords(key) } finally { VaultCrypto.zero(key) }
        assertTrue(repo.setRecoveryWords(words))
        storage.slotA = storage.slotA!!.copy(recoveryEscrow = null)
        repo.lock()
        storage.failLoadSettings = true

        assertFalse(repo.unlockWithRecoveryWords(words))
        assertFalse(repo.isUnlocked.value)
        assertNull(repo.sessionRootKey())
        assertTrue(repo.getEntries().isEmpty())
        assertFalse(repo.isRecoveryWordsSet())
        assertEquals(VaultSettings(), repo.settings.value)
    }

    @Test fun unlock_with_key_malformed_payload_is_error_not_wrong_password() = runBlocking {
        val storage = FaultStorage()
        val repo = VaultRepository(storage)
        assertTrue(repo.createVault("pw") is CreateVaultResult.Success)
        val slot = storage.slotA!!
        val key = VaultCrypto.deriveKey("pw", java.util.Base64.getDecoder().decode(slot.salt))
        try {
            val malformed = VaultCrypto.encrypt("not-json".toByteArray(), key, VaultCrypto.AAD_SLOT_A)
            storage.slotA = slot.copy(blob = java.util.Base64.getEncoder().withoutPadding().encodeToString(malformed))
        } finally {
            VaultCrypto.zero(key)
        }
        repo.lock()
        val supplied = VaultCrypto.deriveKey("pw", java.util.Base64.getDecoder().decode(storage.slotA!!.salt))
        val result = try { repo.unlockWithKey(supplied) } finally { VaultCrypto.zero(supplied) }
        assertTrue("malformed authenticated payload must be an Error", result is UnlockResult.Error)
    }
}