package com.peachpasswords.sync

import com.peachpasswords.vault.StoredVault
import com.peachpasswords.vault.CreateVaultResult
import com.peachpasswords.vault.VaultRepository
import com.peachpasswords.vault.VaultSettings
import com.peachpasswords.vault.VaultStorage
import kotlinx.coroutines.runBlocking
import org.junit.Assert.assertEquals
import org.junit.Assert.assertFalse
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

/** Regression for a stale signed manifest mutating only a prefix before relock. */
class VaultSyncManifestEpochRaceTest {
    private class Storage : VaultStorage {
        private var slotA: StoredVault? = null
        private var slotB: String? = null
        private var version = 0L
        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() = Unit
        override fun version() = version
        override fun stampLocalMutation(): Long { version += 1; return version }
        override fun loadSettings() = settings
        override fun writeSettings(settings: VaultSettings): Boolean { this.settings = settings; return true }
        override fun clearVault(): Boolean { slotA = null; slotB = null; version = 0; settings = VaultSettings(); return true }
    }

    private class Registry(initial: List<PairedDevice>) : PairedDeviceRegistry {
        private val devices = initial.associateBy { it.deviceId }.toMutableMap()
        private val revocations = mutableListOf<DeviceRevocation>()
        var batchCalls = 0

        override fun getAll(): List<PairedDevice> = devices.values.toList()
        override fun get(id: String): PairedDevice? = devices[id]
        override fun upsert(device: PairedDevice) { devices[device.deviceId] = device }
        override fun remove(deviceId: String) { devices.remove(deviceId) }
        override fun listRevocations(): List<DeviceRevocation> = revocations.toList()
        override fun applyManifestBatch(batch: DeviceManifestRegistryBatch): Boolean {
            batchCalls += 1
            batch.removals.forEach(devices::remove)
            batch.upserts.forEach { devices[it.deviceId] = it }
            revocations += batch.incomingRevocations
            return true
        }
    }

    private class Pending : PendingVouchedRegistry {
        override fun getAll(): List<PendingVouchedDevice> = emptyList()
        override fun upsert(device: PendingVouchedDevice) = Unit
        override fun remove(id: String) = Unit
    }

    @Test
    fun stale_manifest_is_rejected_as_a_whole_batch_after_lock_and_reunlock() = runBlocking {
        val repo = VaultRepository(Storage())
        assertTrue(repo.createVault("correct horse battery staple") is CreateVaultResult.Success)
        repo.lock()

        val master = repo.syncMasterKey() // locked, so derive the issuer from a fresh unlock below
        assertEquals(null, master)
        assertTrue(repo.unlock("correct horse battery staple") is com.peachpasswords.vault.UnlockResult.Success)
        val syncMaster = repo.syncMasterKey()!!
        val issuerSigning = deriveSyncSigningKeypair(syncMaster, "issuer")
        val registry = Registry(
            listOf(
                PairedDevice(
                    deviceId = "issuer",
                    deviceName = "Issuer",
                    deviceType = "android",
                    longTermPublicKeyHex = "11".repeat(32),
                    signingPublicKeyHex = pubkeyToHex(issuerSigning.publicKeyRaw),
                    trustSource = "direct",
                    pairedAt = 1L,
                ),
                PairedDevice(
                    deviceId = "old-device",
                    deviceName = "Old",
                    deviceType = "extension",
                    longTermPublicKeyHex = "dd".repeat(32),
                    signingPublicKeyHex = "ee".repeat(32),
                    trustSource = "direct",
                    pairedAt = 1L,
                ),
            ),
        )
        val adapter = VaultSyncAdapter(
            repo,
            "android-self",
            registry,
            Pending(),
            localCredentials = randomLocalIdentityCredentials("android-self"),
        )
        val manifest = buildSignedDeviceManifest(
            issuerDeviceId = "issuer",
            devices = listOf(
                ManifestDeviceEntry(
                    deviceId = "new-device",
                    deviceName = "New",
                    deviceType = "extension",
                    longTermPublicKeyHex = "aa".repeat(32),
                    signingPublicKeyHex = "bb".repeat(32),
                    relayUrl = "wss://relay.peachpasswords.com",
                    pairedAt = 2L,
                ),
            ),
            signingPrivateKey = issuerSigning.privateKey,
            revoked = listOf(
                DeviceRevocation(
                    deviceId = "old-device",
                    longTermPublicKeyHex = "dd".repeat(32),
                    revokedAt = System.currentTimeMillis(),
                    revokedByDeviceId = "issuer",
                ),
            ),
        )

        val enteredFence = CountDownLatch(1)
        val releaseFence = CountDownLatch(1)
        try {
            repo.beforeSessionEpochCommit = {
                enteredFence.countDown()
                check(releaseFence.await(5, TimeUnit.SECONDS)) { "manifest did not release" }
            }
            val applyFailure = AtomicReference<Throwable?>()
            val applyDone = CountDownLatch(1)
            Thread {
                try {
                    runBlocking {
                        adapter.applyIncomingDeviceManifest(manifest, ByteArray(32) { 0x11 })
                    }
                } catch (t: Throwable) {
                    applyFailure.set(t)
                } finally {
                    applyDone.countDown()
                }
            }.start()
            assertTrue(enteredFence.await(5, TimeUnit.SECONDS))

            val relockDone = CountDownLatch(1)
            Thread {
                repo.lock()
                runBlocking {
                    assertTrue(repo.unlock("correct horse battery staple") is com.peachpasswords.vault.UnlockResult.Success)
                }
                relockDone.countDown()
            }.start()
            assertTrue(relockDone.await(5, TimeUnit.SECONDS))
            releaseFence.countDown()
            assertTrue(applyDone.await(5, TimeUnit.SECONDS))
            assertEquals(null, applyFailure.get())
            assertEquals("the stale epoch must never reach the registry", 0, registry.batchCalls)
            assertTrue(registry.get("old-device") != null)
            assertTrue(registry.get("new-device") == null)
            assertTrue(registry.listRevocations().isEmpty())
        } finally {
            releaseFence.countDown()
            repo.beforeSessionEpochCommit = {}
        }
    }
}