package dev.privacyllc.period.feature.lock

import androidx.datastore.preferences.core.PreferenceDataStoreFactory
import dev.privacyllc.period.core.datastore.UserPreferencesRepository
import dev.privacyllc.period.core.security.AppLockRepository
import dev.privacyllc.period.core.security.MacProvider
import dev.privacyllc.period.core.security.UnlockResult
import dev.privacyllc.period.lock.AppLockController
import dev.privacyllc.period.lock.LockState
import dev.privacyllc.period.lock.lockStateOf
import kotlinx.coroutines.CoroutineScope
import kotlinx.coroutines.Dispatchers
import kotlinx.coroutines.delay
import kotlinx.coroutines.flow.combine
import kotlinx.coroutines.flow.first
import kotlinx.coroutines.Job
import kotlinx.coroutines.launch
import kotlinx.coroutines.runBlocking
import kotlinx.coroutines.test.UnconfinedTestDispatcher
import kotlinx.coroutines.test.resetMain
import kotlinx.coroutines.test.setMain
import kotlinx.coroutines.withTimeout
import org.junit.After
import org.junit.Assert.assertEquals
import org.junit.Assert.assertFalse
import org.junit.Assert.assertTrue
import org.junit.Before
import org.junit.Rule
import org.junit.Test
import org.junit.rules.TemporaryFolder
import org.junit.runner.RunWith
import org.robolectric.RobolectricTestRunner
import org.robolectric.annotation.Config
import javax.crypto.Mac
import javax.crypto.spec.SecretKeySpec

/**
 * The settings side of the app lock, which had no test and two defects.
 *
 * One was a security hole: "Change PIN" fired the check of the current PIN and
 * moved to "choose a PIN" in the same breath, without waiting for the answer, so
 * any four digits reached the replacement screen and overwrote the PIN. Under
 * this project's no-recovery policy that costs the owner her entire history.
 *
 * The other was a race: the PIN was written first and the session unlocked
 * second, and the gate could shut in between — disposing the subtree this very
 * ViewModel lives in, and cancelling the unlock with it. The user was thrown to
 * the lock screen to type the PIN she had chosen a second earlier.
 *
 * Both are asserted here against the real repository over a host-JVM signing
 * key, not a fake of the thing under test.
 */
@RunWith(RobolectricTestRunner::class)
@Config(sdk = [34])
class LockSettingsViewModelTest {

    @get:Rule val temp = TemporaryFolder()

    private val dispatcher = UnconfinedTestDispatcher()
    private lateinit var lock: AppLockRepository
    private lateinit var controller: AppLockController
    private lateinit var vm: LockSettingsViewModel
    private lateinit var subscription: Job

    /**
     * A signing key that exists off-device.
     *
     * `AndroidKeyStore` cannot be reached on the host JVM and Robolectric ships
     * no crypto shadows, so this is the seam `AppLockRepository`'s two-argument
     * constructor exists for. Deliberately a duplicate of core/security's own
     * fake rather than a shared fixture: fifteen lines, against a Gradle
     * test-fixtures setup whose `internal` visibility rules are their own risk.
     */
    private class InMemoryMacProvider : MacProvider {
        private var key: ByteArray? = null
        override fun hasKey() = key != null
        override fun ensureKey() { if (key == null) key = ByteArray(32) { it.toByte() } }
        override fun mac(data: ByteArray): ByteArray {
            val k = key ?: error("the verifier key is absent")
            return Mac.getInstance("HmacSHA256").apply { init(SecretKeySpec(k, "HmacSHA256")) }.doFinal(data)
        }
        override fun deleteKey() { key = null }
    }

    @Before fun setUp() {
        Dispatchers.setMain(dispatcher)
        // The stores get a real background dispatcher, not the test one: the
        // assertions below block the test thread while they wait, and a
        // DataStore whose own scope is that thread can never finish the write
        // they are waiting for. Main stays the test dispatcher so viewModelScope
        // still runs eagerly.
        val scope = CoroutineScope(Dispatchers.IO)
        lock = AppLockRepository(
            PreferenceDataStoreFactory.create(scope = scope) { temp.newFile("app_lock.preferences_pb") },
            InMemoryMacProvider(),
        )
        val prefs = UserPreferencesRepository(
            PreferenceDataStoreFactory.create(scope = scope) { temp.newFile("prefs.preferences_pb") },
        )
        controller = AppLockController()
        vm = LockSettingsViewModel(lock, prefs, controller)

        // `state` is stateIn(WhileSubscribed): with nobody collecting, the
        // upstream never runs and `.value` sits on its initial value forever.
        // The screen is that collector in the app; here it is this.
        subscription = CoroutineScope(Dispatchers.IO).launch { vm.state.collect {} }
        await { vm.state.value == LockSettings() }
    }

    @After fun tearDown() {
        subscription.cancel()
        Dispatchers.resetMain()
    }

    /**
     * Generous on purpose. Every PIN here costs a real PBKDF2 derivation at
     * 210,000 iterations, and a test that asks for three of them can pass alone
     * and time out when the whole suite runs beside it. This budget is for
     * catching a genuine hang, not for measuring the crypto.
     */
    private fun await(predicate: suspend () -> Boolean) = runBlocking {
        withTimeout(30_000) { while (!predicate()) delay(10) }
    }

    private fun pin(value: String) = value.toCharArray()

    // -----------------------------------------------------------------------
    // The race
    // -----------------------------------------------------------------------

    @Test
    fun `setting the first PIN never closes the gate`() {
        val seen = mutableListOf<LockState>()
        val watching = CoroutineScope(dispatcher).launch {
            combine(lock.hasPin, controller.unlocked, ::lockStateOf).collect { seen += it }
        }

        vm.setPin(pin("2468"))
        await { vm.state.value.message == LockSettings.Message.PIN_SET }
        watching.cancel()

        // Not "ends unlocked" — never locked, at any instant. The gate disposes
        // the whole app subtree when it closes, so a single frame of Locked is
        // the user back at the lock screen typing the PIN she just chose.
        assertFalse("the gate closed while the PIN was being set: $seen", LockState.Locked in seen)
        assertTrue(runBlocking { lock.hasPin.first() })
        assertTrue(controller.unlocked.value)
    }

    // -----------------------------------------------------------------------
    // The bypass
    // -----------------------------------------------------------------------

    @Test
    fun `changing the PIN needs the current one and cannot be skipped`() {
        vm.setPin(pin("2468"))
        await { lock.hasPin.first() }

        var authorised = false
        vm.authoriseChange(pin("1111")) { authorised = true }
        await { vm.state.value.message == LockSettings.Message.WRONG_PIN }
        assertFalse("a wrong current PIN authorised a replacement", authorised)

        // The screen is not the only guard: even a caller that ignored the
        // callback and asked for the write anyway must be refused.
        vm.setPin(pin("1357"))
        await { !vm.state.value.busy }

        // The old PIN still opens the app; the attempted replacement does not.
        assertTrue(runBlocking { lock.check(pin("2468")) } is UnlockResult.Unlocked)
        assertTrue(runBlocking { lock.check(pin("1357")) } is UnlockResult.Wrong)
    }

    @Test
    fun `the right current PIN authorises exactly one replacement`() {
        vm.setPin(pin("2468"))
        await { lock.hasPin.first() }

        var reached = false
        vm.authoriseChange(pin("2468")) { reached = true }
        await { reached }

        vm.setPin(pin("1357"))
        await { vm.state.value.message == LockSettings.Message.PIN_SET }
        assertTrue(runBlocking { lock.check(pin("1357")) } is UnlockResult.Unlocked)

        // Permission is spent. A second write without asking again is refused.
        vm.setPin(pin("9999"))
        await { vm.state.value.message == LockSettings.Message.WRONG_PIN }
        assertTrue(runBlocking { lock.check(pin("1357")) } is UnlockResult.Unlocked)
    }

    @Test
    fun `abandoning a change withdraws the permission it was granted`() {
        vm.setPin(pin("2468"))
        await { lock.hasPin.first() }

        vm.authoriseChange(pin("2468")) { }
        await { !vm.state.value.busy }
        vm.cancelChange()

        vm.setPin(pin("1357"))
        await { vm.state.value.message == LockSettings.Message.WRONG_PIN }
        assertTrue(runBlocking { lock.check(pin("2468")) } is UnlockResult.Unlocked)
    }

    // -----------------------------------------------------------------------
    // Turning it off
    // -----------------------------------------------------------------------

    @Test
    fun `the lock comes off only for the right PIN`() {
        vm.setPin(pin("2468"))
        await { lock.hasPin.first() }

        var removed = false
        vm.removePin(pin("1111")) { removed = true }
        await { vm.state.value.message == LockSettings.Message.WRONG_PIN }
        assertFalse(removed)
        assertTrue(runBlocking { lock.hasPin.first() })

        vm.removePin(pin("2468")) { removed = true }
        await { removed }
        assertFalse(runBlocking { lock.hasPin.first() })
        assertEquals(LockSettings.Message.PIN_REMOVED, vm.state.value.message)
    }
}
