package dev.privacyllc.period.core.security

import android.os.SystemClock
import androidx.datastore.core.DataStore
import androidx.datastore.preferences.core.Preferences
import androidx.datastore.preferences.core.edit
import androidx.datastore.preferences.core.emptyPreferences
import androidx.datastore.preferences.core.stringPreferencesKey
import kotlinx.coroutines.flow.Flow
import kotlinx.coroutines.flow.catch
import kotlinx.coroutines.flow.first
import kotlinx.coroutines.flow.map
import java.io.IOException
import java.security.MessageDigest
import java.util.Base64

/** What happened when a PIN was offered. */
sealed interface UnlockResult {
    data object Unlocked : UnlockResult

    /** Wrong PIN. [waitMillis] is zero while the user still has free attempts. */
    data class Wrong(val waitMillis: Long) : UnlockResult

    /** Refused without checking, because a lockout is still running. */
    data class TooSoon(val waitMillis: Long) : UnlockResult

    /** No PIN has been set on this device. */
    data object NoPin : UnlockResult

    /**
     * The Keystore entry is gone, so no PIN can ever match again.
     *
     * Deliberately distinct from [Wrong]. Telling somebody their correct PIN is
     * wrong, forever, with an ever-growing delay, is the cruellest failure this
     * feature could produce — and under the no-recovery policy they would have
     * no way to find out otherwise.
     */
    data object KeyUnavailable : UnlockResult
}

/**
 * Everything the app lock persists, in its own DataStore file.
 *
 * ## Why not `UserPreferences`
 *
 * `UserPreferencesRepository.resetToDefaults()` is `edit { it.clear() }`. It has
 * no production caller today, and the day it gets one, a verifier record living
 * in that store would be silently removed — locking the user out of their own
 * history with no way back. The two stores also have deliberately different
 * deletion semantics: Delete My Data keeps your settings, and the erase hatch
 * destroys the lock. Keeping them apart is what lets both of those be true.
 *
 * ## There is no `pinSet` boolean
 *
 * It is derived from whether a verifier record exists. A separate flag would be
 * a second record of one fact, and the two would eventually disagree — with the
 * disagreement showing up as either a lock nobody can open or an app that never
 * locks.
 *
 * The store is injected rather than built from a `Context`, mirroring
 * `UserPreferencesRepository` for the same reason: it can then be exercised
 * against a temporary file without an emulator.
 */
class AppLockRepository internal constructor(
    private val store: DataStore<Preferences>,
    private val macs: MacProvider,
    private val clocks: Clocks,
) {

    /**
     * The constructor callers outside this module use.
     *
     * `internal` on the primary one is the same mechanism `CycleRepository`
     * uses to keep Room off everyone else's compile classpath: `MacProvider` and
     * `Clocks` are seams for this module's own tests, and nothing above needs to
     * be able to name them — or to substitute them.
     */
    constructor(store: DataStore<Preferences>) :
        this(store, AndroidKeyStoreMacProvider(), SystemClocks)

    /**
     * The real repository over a signing key that exists on the host JVM.
     *
     * For tests **above** this module. `AndroidKeyStore` cannot be reached
     * without a device, so an `app`-module test of the settings flow — the one
     * that proves a PIN cannot be changed without knowing it, and that setting
     * one does not lock you out of the session — would otherwise have to fake
     * the whole repository and prove nothing about it.
     *
     * [Clocks] stays internal: the backoff timing is this module's own concern
     * and is tested here, against a fake clock that can be moved and rebooted.
     */
    constructor(store: DataStore<Preferences>, macs: MacProvider) :
        this(store, macs, SystemClocks)

    private val verifier: PinVerifier = PinVerifier(macs)

    /**
     * True once a PIN exists. Derived, never stored.
     *
     * An unreadable file reads as **no lock**, and that is a decision rather
     * than an oversight.
     *
     * This flow had no `.catch` at all, unlike the one in
     * `UserPreferencesRepository`, and it is collected on the startup path in
     * two places that cannot handle a throw — the gate's `stateIn`, which
     * launches outside its own exception handler, and `MainActivity`'s
     * `FLAG_SECURE` collector. A corrupt `app_lock` file crashed the app before
     * any UI existed to say so.
     *
     * Opening is the safe direction, for the same reason `VerifierRecord.decode`
     * already treats a corrupt record as "no PIN set": whoever can corrupt this
     * file has the app's private storage, and so can read the database the lock
     * does not encrypt anyway. The alternative is an app that cannot be opened
     * *or* erased, because the way out is behind the lock that is broken.
     */
    val hasPin: Flow<Boolean> = store.data
        .catch { failure -> if (failure is IOException) emit(emptyPreferences()) else throw failure }
        .map { it[Keys.Verifier] != null }

    suspend fun hasPinNow(): Boolean = hasPin.first()

    /**
     * Set the first PIN, or replace one after the current PIN has been checked.
     *
     * Returns false only if the Keystore refuses to produce a key at all, which
     * is the one case where the honest thing is to leave the lock off rather
     * than enable a lock that cannot be opened.
     */
    suspend fun setPin(pin: CharArray): Boolean {
        val record = runCatching { verifier.enroll(pin) }.getOrNull() ?: return false
        store.edit {
            it[Keys.Verifier] = record.encode()
            it.remove(Keys.Lockout)
            it.remove(Keys.LockoutMac)
        }
        return true
    }

    /**
     * Offer a PIN.
     *
     * A running lockout is checked **before** the PIN is, so guessing during a
     * lockout costs an attempt and tells them nothing.
     */
    suspend fun check(pin: CharArray): UnlockResult {
        val prefs = store.data.first()
        val record = VerifierRecord.decode(prefs[Keys.Verifier]) ?: return UnlockResult.NoPin

        val state = readLockout(prefs)
        val wait = LockoutPolicy.remainingMillis(state, clocks)
        if (wait > 0L) return UnlockResult.TooSoon(wait)

        return when (val outcome = runCatching { verifier.verify(pin, record) }.getOrNull()) {
            null -> UnlockResult.KeyUnavailable
            is PinCheck.KeyUnavailable -> UnlockResult.KeyUnavailable

            is PinCheck.Correct -> {
                store.edit { edited ->
                    outcome.rehashed?.let { edited[Keys.Verifier] = it.encode() }
                    edited.remove(Keys.Lockout)
                    edited.remove(Keys.LockoutMac)
                }
                UnlockResult.Unlocked
            }

            is PinCheck.Wrong -> {
                val next = LockoutPolicy.onFailure(state ?: LockoutState(), clocks)
                writeLockout(next)
                UnlockResult.Wrong(LockoutPolicy.remainingMillis(next, clocks))
            }
        }
    }

    /** Milliseconds still to wait, for a screen that has to render a countdown. */
    suspend fun lockoutRemainingMillis(): Long =
        LockoutPolicy.remainingMillis(readLockout(store.data.first()), clocks)

    /**
     * Remove the lock entirely: the record, the counter, and the key itself.
     *
     * The key deletion is the step that must not be dropped. Without it a user
     * who erases everything keeps a working lock over an empty database —
     * erased *and* still locked out — which is the single worst outcome this
     * feature can produce, and is why a test mutates this line out and requires
     * a red.
     */
    suspend fun clearLock() {
        macs.deleteKey()
        store.edit {
            it.remove(Keys.Verifier)
            it.remove(Keys.Lockout)
            it.remove(Keys.LockoutMac)
        }
    }

    /**
     * Null means "there is a counter and it does not check out", which
     * [LockoutPolicy] treats as maximum backoff.
     *
     * It must **not** mean "there is no counter". A fresh install has neither
     * key, and so does a device where the last unlock succeeded and cleared
     * them — reporting either as tampering would open the app in a
     * fifteen-minute lockout that nobody had earned. Both keys absent is the
     * ordinary empty state; one of them absent is somebody having removed it.
     */
    private fun readLockout(prefs: Preferences): LockoutState? {
        val blob = prefs[Keys.Lockout]
        val stored = prefs[Keys.LockoutMac]
        if (blob == null && stored == null) return LockoutState().longerOf(unsavedLockout)
        if (blob == null || stored == null) return null
        val expected = runCatching { macOf(blob) }.getOrNull() ?: return null
        // A tampered or unverifiable counter reads as null, and null is maximum
        // backoff in LockoutPolicy — so editing the file is the worst available
        // move rather than the best one.
        if (!MessageDigest.isEqual(expected, stored.decodeB64())) return null
        return LockoutState.decode(blob).longerOf(unsavedLockout)
    }

    /** Whichever of the two has seen more failures — see [unsavedLockout]. */
    private fun LockoutState?.longerOf(other: LockoutState?): LockoutState? = when {
        this == null || other == null -> this ?: other
        other.failedAttempts > failedAttempts -> other
        else -> this
    }

    /**
     * The failed attempt this process could not sign.
     *
     * Signing needs the Keystore, and the Keystore can be unavailable for a
     * moment. When it is, the counter cannot be written — and the old code
     * simply returned, so the attempt cost nothing and wrong PINs became free
     * for as long as the condition lasted. That is the one direction this must
     * not fail in.
     *
     * Held in memory instead, and read back by [readLockout] whenever it is the
     * longer wait. Deliberately not persisted: it is the earned delay, not the
     * maximum a tampered counter earns, and it is forgotten when the process
     * dies — which is the same bound an attacker already has on a reboot.
     */
    @Volatile
    private var unsavedLockout: LockoutState? = null

    private suspend fun writeLockout(state: LockoutState) {
        val blob = state.encode()
        val mac = runCatching { macOf(blob) }.getOrNull() ?: run {
            unsavedLockout = state
            return
        }
        unsavedLockout = null
        store.edit {
            it[Keys.Lockout] = blob
            it[Keys.LockoutMac] = mac.encodeB64()
        }
    }

    /** Tag 0x02, so a counter value can never be replayed as a verifier value. */
    private fun macOf(blob: String): ByteArray =
        macs.mac(byteArrayOf(PinVerifier.TAG_BACKOFF) + blob.toByteArray(Charsets.UTF_8))

    private fun ByteArray.encodeB64(): String =
        Base64.getEncoder().withoutPadding().encodeToString(this)

    private fun String.decodeB64(): ByteArray =
        runCatching { Base64.getDecoder().decode(this) }.getOrDefault(ByteArray(0))

    private object Keys {
        val Verifier = stringPreferencesKey("pin_verifier_v1")
        val Lockout = stringPreferencesKey("lockout_state_v1")
        val LockoutMac = stringPreferencesKey("lockout_state_mac_v1")
    }
}

/** The real clocks. Split out so tests can move time without waiting. */
internal object SystemClocks : Clocks {
    override fun wallMillis(): Long = System.currentTimeMillis()
    override fun elapsedMillis(): Long = SystemClock.elapsedRealtime()
}
