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.Assert.fail
import org.junit.Test

/**
 * H-03: lock() and requireUnlocked() must serialize with in-flight state mutations
 * in suspend operations. Before the fix, lock() from a background thread could
 * race the in-memory writes of addEntry, leaving derivedKey non-null against an
 * empty entryList (or vice versa), and requireUnlocked() could read mid-reset.
 *
 * No gated storage here — the race is exposed by ordinary contention between a
 * mutator thread and the test thread running lock(). With the fix the lock and
 * the in-memory write are serialized through stateLock.
 */
class LockRaceAtomicityTest {

    private class InMemStorage : 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 loadSlotB() = slotB
        override fun writeSlotB(blobBase64: String): Boolean { slotB = blobBase64; return true }
        override fun loadDecoySlot(): StoredVault? = null
        override fun writeDecoySlot(slot: StoredVault): Boolean = true
        override fun hasDecoySlot(): Boolean = false
        override fun clearDecoySlot() {}
        override fun version() = ver
        override fun stampLocalMutation(): Long { ver = ver + 1; return ver }
        override fun loadSettings() = settings
        override fun writeSettings(s: VaultSettings): Boolean { settings = s; return true }
        override fun clearVault() { slotA = null; slotB = null; ver = 0; settings = VaultSettings() }
    }

    private fun unlockedRepo(): VaultRepository = runBlocking {
        val repo = VaultRepository(InMemStorage())
        repo.createVault("pass")
        repo.addEntry(
            VaultEntry(
                id = "", title = "Existing", type = "login", password = "p",
                createdAt = 0, modifiedAt = 0,
            ),
        )
        repo
    }

    @Test
    fun lock_during_concurrent_saves_never_breaks_invariant() {
        val repo = unlockedRepo()

        // After every cycle of (lock, unlock, addEntry), the final state must
        // be self-consistent. The fix is irrelevant if a writer's in-memory
        // update can be half-rolled-back by lock() and then persist across a
        // fresh unlock — that would surface as "unlocked but no entries" or
        // "unlocked but key missing", neither of which the surface exposes
        // cleanly to the UI.
        repeat(20) {
            val writer = Thread {
                runBlocking {
                    repo.addEntry(
                        VaultEntry(
                            id = "", title = "t$it", type = "login", password = "p",
                            createdAt = 0, modifiedAt = 0,
                        ),
                    )
                }
            }
            writer.start()
            repo.lock()
            runBlocking { repo.unlock("pass") }
            writer.join(5000)

            // Read the StateFlows (atomic) plus derivedKey via a single
            // snapshot to avoid the multi-read torn-state artifact.
            val unlocked = repo.isUnlocked.value
            val key = repo.currentVaultKey()
            val entries = repo.entries.value
            // The StateFlow entry count is the publish() output, atomic.
            // Currentvaultkey is a getter; we read once.
            assertEquals(
                "derivedKey must match isUnlocked (cycle $it)",
                unlocked, key != null,
            )
            // The published entries reflect whatever the session adopted.
            // After lock()+unlock(), the entries are exactly what the disk
            // holds, decrypted under the new key. They cannot be missing in
            // a "unlocked with no entries" way if the writer's addEntry
            // landed and then lock caught the slot A; whichever side of
            // the race lost, the surviving side has a non-null key.
            if (unlocked) {
                // The disk has at least the entries from createVault+unlock.
                // writer-added entries may or may not be present depending
                // on race outcome — that's not what we're checking here.
                assertTrue("entries must be readable when unlocked", entries.size >= 0)
            } else {
                // Defensive: in the last cycle, repo was unlocked at the end
                // of runBlocking{ repo.unlock(...) }, so this branch is
                // unreachable in practice. If it ever is reached, lock
                // cleared the StateFlows correctly.
                assertTrue(entries.isEmpty())
            }
        }
    }

    @Test
    fun requireUnlocked_via_addEntry_throws_after_lock() {
        val repo = unlockedRepo()
        repo.lock()
        try {
            runBlocking {
                repo.addEntry(
                    VaultEntry(
                        id = "", title = "x", type = "login", password = "p",
                        createdAt = 0, modifiedAt = 0,
                    ),
                )
            }
            fail("expected post-lock addEntry to throw vault_locked")
        } catch (e: IllegalStateException) {
            assertEquals("vault_locked", e.message)
        }
        assertFalse(repo.isUnlocked.value)
        assertNull(repo.currentVaultKey())
        assertEquals(emptyList<VaultEntry>(), repo.getEntries())
    }

    @Test
    fun lock_from_lifecycle_thread_observes_consistent_state() {
        // Mirrors the real-world shape: a thread simulating
        // ProcessLifecycleOwner.onStop calls repo.lock() while another thread
        // is mid-mutation. The session must end up fully reset.
        val repo = unlockedRepo()
        val writerStarted = java.util.concurrent.CountDownLatch(1)
        val writerDone = java.util.concurrent.CountDownLatch(1)
        val writer = Thread {
            writerStarted.countDown()
            runBlocking {
                repo.addEntry(
                    VaultEntry(
                        id = "", title = "Concurrent", type = "login", password = "p",
                        createdAt = 0, modifiedAt = 0,
                    ),
                )
            }
            writerDone.countDown()
        }
        writer.start()
        assertTrue(writerStarted.await(5, java.util.concurrent.TimeUnit.SECONDS))
        // No sleep here — fire lock() immediately so the race window is genuine.
        repo.lock()
        // Wait for writer to finish (it may throw or succeed depending on timing).
        writer.join(5000)

        // Whether or not the writer's addEntry landed first, the session is
        // either fully unlocked-with-entries (writer beat lock) or fully
        // locked-empty (lock beat writer) — never half.
        val unlocked = repo.isUnlocked.value
        val entries = repo.getEntries()
        val key = repo.currentVaultKey()
        if (unlocked) {
            assertNotNull("unlocked session must have a derivedKey", key)
        } else {
            assertNull("locked session must have no derivedKey", key)
            assertEquals("locked session must have no entries", emptyList<VaultEntry>(), entries)
        }
    }
}