package com.peachpasswords.vault

import kotlinx.coroutines.runBlocking
import org.junit.Assert.assertArrayEquals
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

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

class RecoveryEscrowTest {

    /** Recovery words this vault would actually show — derived from its sync root key. */
    private fun recoveryWordsOf(repo: VaultRepository): List<String> {
        val key = repo.syncMasterKey()!!
        return try { rootKeyToWords(key) } finally { VaultCrypto.zero(key) }
    }

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

    // 1 — setRecoveryWords writes BOTH the offline-restore escrow and the
    // browser-compatible root-key escrow (MF-1).
    @Test fun set_recovery_words_writes_both_escrows() = runBlocking {
        val storage = RecoveryFakeStorage()
        val repo = VaultRepository(storage)
        repo.createVault("pw")
        assertNull(storage.slotA?.recoveryEscrow)
        assertNull(storage.slotA?.rootKeyEscrow)
        assertTrue(repo.setRecoveryWords(recoveryWordsOf(repo)))
        assertNotNull(storage.slotA?.recoveryEscrow)
        assertNotNull(storage.slotA?.rootKeyEscrow)
        assertTrue(repo.hasRecoveryEscrow())
    }

    // 2 — unlockWithRecoveryWords with correct words restores entries (offline escrow).
    @Test fun unlock_with_correct_words_succeeds() = runBlocking {
        val storage = RecoveryFakeStorage()
        val repo = VaultRepository(storage)
        repo.createVault("pw")
        repo.addEntry(login("GitHub"))
        val words = recoveryWordsOf(repo)
        repo.setRecoveryWords(words)
        repo.lock()
        assertTrue(repo.unlockWithRecoveryWords(words))
        assertTrue(repo.getEntries().any { it.title == "GitHub" })
    }

    // 3 — unlockWithRecoveryWords with wrong (still checksum-valid) words fails.
    @Test fun unlock_with_wrong_words_fails() = runBlocking {
        val storage = RecoveryFakeStorage()
        val repo = VaultRepository(storage)
        repo.createVault("pw")
        repo.setRecoveryWords(recoveryWordsOf(repo))
        repo.lock()
        val wrong = rootKeyToWords(ByteArray(32) { (it + 7).toByte() })
        assertFalse(repo.unlockWithRecoveryWords(wrong))
    }

    // 4 — unlockWithRecoveryWords with no escrow set returns false.
    @Test fun unlock_with_no_escrow_fails() = runBlocking {
        val storage = RecoveryFakeStorage()
        val repo = VaultRepository(storage)
        repo.createVault("pw")
        val words = recoveryWordsOf(repo) // never call setRecoveryWords
        repo.lock()
        assertFalse(repo.hasRecoveryEscrow())
        assertFalse(repo.unlockWithRecoveryWords(words))
    }

    // 5 — Full round-trip: create → set words → reset password via recovery → unlock with new password.
    @Test fun full_recovery_resets_password() = runBlocking {
        val storage = RecoveryFakeStorage()
        val repo = VaultRepository(storage)
        repo.createVault("old-pass")
        repo.addEntry(login("Bank"))
        val words = recoveryWordsOf(repo)
        repo.setRecoveryWords(words)
        repo.lock()

        assertTrue(repo.unlockWithRecoveryWords(words))
        assertTrue(repo.resetPasswordWithRecoveryWords(words, "brand-new-pass"))
        assertTrue(repo.isUnlocked.value)
        repo.lock()

        assertEquals(UnlockResult.WrongPassword, repo.unlock("old-pass"))
        assertEquals(UnlockResult.Success, repo.unlock("brand-new-pass"))
        assertTrue(repo.getEntries().any { it.title == "Bank" })
        assertTrue(repo.hasRecoveryEscrow())
    }

    // 6 — A checksum-INVALID phrase (e.g. all-"abandon") is rejected, matching the
    // browser. Under browser-parity IKM (`wordsToRootKeyMaterial`) such a phrase
    // fails the BIP39 checksum, so no escrow is sealed and recovery cannot proceed.
    // (Previously this asserted the old BLAKE3 behavior, which round-tripped any
    // phrase — that divergence is exactly what MF-1 fixes.)
    @Test fun checksum_invalid_phrase_is_rejected() = runBlocking {
        val storage = RecoveryFakeStorage()
        val repo = VaultRepository(storage)
        repo.createVault("pw")
        repo.addEntry(login("Mail"))
        val allAbandon = List(12) { "abandon" } // valid words, invalid checksum
        // No escrow gets sealed for an invalid phrase.
        repo.setRecoveryWords(allAbandon)
        assertNull(storage.slotA?.recoveryEscrow)
        assertNull(storage.slotA?.rootKeyEscrow)
        repo.lock()
        assertFalse(repo.unlockWithRecoveryWords(allAbandon))
    }

    // 7 — CROSS-PLATFORM KAT. Vector generated from the browser's exact algorithm
    // (recovery-words.ts) against the real browser word list (/tmp/mf1_kat_gen.mjs):
    // proves the same 12 words derive the same recovery key and unwrap the same
    // root key on both platforms.
    @Test fun kat_matches_browser_algorithm() {
        val rootKeyHex = "bfb2d5731a470e1b1a0ae9521f501ac5edad8b4ce8e064b76da4cb3cdf00b937"
        // Corrected in v2.4.13. These were previously this platform's own
        // output under a word list that duplicated "abandon" at index 1, so the
        // vector was one position off the browser at every index >= 1 — this
        // test asserted browser parity while pinning the divergence. Recomputed
        // against the browser's canonical list; the material/HKDF/ciphertext
        // vectors below are unchanged because they key off the root key bytes,
        // not the word ordering.
        val expectedWords = listOf(
            "sausage", "note", "fresh", "cruel", "ill", "assault",
            "habit", "frog", "faith", "whisper", "allow", "merit",
        )
        val wordMaterialHex = "bfb2d5731a470e1b1a0ae9521f501ac500000000000000000000000000000000"
        val recoveryKeyHex = "06880fb677e5b1cd1a65de6340a605681a55c86268885d644117401a9effd990"
        val saltHex = "000102030405060708090a0b0c0d0e0f"
        val nonceHex = "0b0a090807060504030201ff"
        val ciphertextHex = "bba39ac8dfaaff82b916ff434df340ee458c5e6b9b6a18b3a81f4748191cbf03" +
            "ac9f84707cb10c17430b9d63cd83960d2d2cd9da89aa907d2e3bd0cb78bd3628"
        val tagHex = "6f82150ebdbddef6312e5e630975089c"

        val rootKey = hexToRootKey(rootKeyHex)

        // (a) words derive from the root key exactly as the browser.
        assertEquals(expectedWords, rootKeyToWords(rootKey))

        // (b) IKM = 16 entropy bytes followed by 16 zero bytes.
        assertEquals(wordMaterialHex, rootKeyToHex(wordsToRootKeyMaterial(expectedWords)))

        // (c) HKDF-derived recovery key matches the browser's.
        val recoveryKey = deriveRecoveryKey(wordsToRootKeyMaterial(expectedWords), hexToRootKey(saltHex)).encoded
        assertEquals(recoveryKeyHex, rootKeyToHex(recoveryKey))

        // (d) the browser-wrapped escrow unwraps to the original root key under AAD 0x01.
        val combined = hexToRootKey(nonceHex) + hexToRootKey(ciphertextHex) + hexToRootKey(tagHex)
        val plaintext = VaultCrypto.decrypt(combined, recoveryKey, VaultCrypto.AAD_RECOVERY_ROOTKEY)
        assertEquals(rootKeyHex, String(plaintext, Charsets.US_ASCII))
        assertArrayEquals(rootKey, hexToRootKey(String(plaintext, Charsets.US_ASCII)))
    }

    // 8 — A browser-format vault (only the root-key escrow present) recovers the
    // root key and resets the password; entries are empty (they return via sync).
    @Test fun browser_format_vault_recovers_root_key_only() = runBlocking {
        val storage = RecoveryFakeStorage()
        val repo = VaultRepository(storage)
        repo.createVault("pw")
        repo.addEntry(login("Local"))
        val words = recoveryWordsOf(repo)
        val rootKeyBefore = repo.sessionRootKey()!!
        repo.setRecoveryWords(words)
        // Simulate a browser-made vault: drop the Android-only whole-payload escrow.
        storage.slotA = storage.slotA!!.copy(recoveryEscrow = null)
        repo.lock()

        assertTrue(repo.unlockWithRecoveryWords(words))
        // sessionRootKey, not syncMasterKey: a recovery unlock has no password key
        // yet, so the sync accessors fail closed until the reset completes (M-04).
        assertArrayEquals(rootKeyBefore, repo.sessionRootKey())  // same root key recovered
        assertTrue(repo.getEntries().isEmpty())                  // entries return via sync
    }

    // 9 — changePassword keeps the recovery escrow valid for the SAME words.
    @Test fun change_password_keeps_recovery_valid() = runBlocking {
        val storage = RecoveryFakeStorage()
        val repo = VaultRepository(storage)
        repo.createVault("old")
        repo.addEntry(login("Email"))
        val words = recoveryWordsOf(repo)
        repo.setRecoveryWords(words)
        assertTrue(repo.changePassword("old", "new"))
        // Root-key escrow is carried unchanged (browser parity); offline escrow too.
        assertNotNull(storage.slotA?.rootKeyEscrow)
        repo.lock()
        assertTrue(repo.unlockWithRecoveryWords(words))
        assertTrue(repo.getEntries().any { it.title == "Email" })
    }

    // 10 — Backward compat: a vault with neither escrow still opens with its password.
    @Test fun vault_without_escrow_still_opens() = runBlocking {
        val storage = RecoveryFakeStorage()
        val repo = VaultRepository(storage)
        repo.createVault("pw")
        assertNull(storage.slotA?.recoveryEscrow)
        assertNull(storage.slotA?.rootKeyEscrow)
        repo.lock()
        assertEquals(UnlockResult.Success, repo.unlock("pw"))
    }
}