package ch.threema.localcrypto

import ch.threema.common.mapState
import ch.threema.localcrypto.models.MasterKeyState
import ch.threema.localcrypto.models.RemoteSecretParameters
import ch.threema.localcrypto.models.RemoteSecretProtectionState
import kotlinx.coroutines.flow.Flow
import kotlinx.coroutines.flow.MutableStateFlow
import kotlinx.coroutines.flow.StateFlow
import kotlinx.coroutines.flow.filterIsInstance
import kotlinx.coroutines.flow.first
import kotlinx.coroutines.flow.map
import kotlinx.coroutines.flow.update

/**
 * The [MasterKeyLockStateHolder] keeps track of if and how the master key is currently locked.
 */
class MasterKeyLockStateHolder {
    private val lockStateFlow = MutableStateFlow<LockState>(LockState.Unknown)
    private var lockState: LockState
        get() = lockStateFlow.value
        set(newState) {
            lockStateFlow.update { previousState ->
                if (previousState is LockState.PermanentlyLocked) {
                    error("Cannot unlock, permanently locked")
                }

                // wipe the master key from memory by filling it with zeroes
                (previousState as? LockState.Unlocked)?.masterKey?.let { previousMasterKey ->
                    if ((newState as? LockState.Unlocked)?.masterKey != previousMasterKey) {
                        previousMasterKey.invalidate()
                    }
                }

                newState
            }
        }

    val masterKeyFlow: StateFlow<MasterKey?> =
        lockStateFlow.mapState { lockState ->
            (lockState as? LockState.Unlocked)?.masterKey
        }

    val remoteSecretProtectionFlow: Flow<RemoteSecretProtectionState>
        get() = lockStateFlow
            .filterIsInstance<RemoteSecretLockState>()
            .map { lockState ->
                when (lockState.remoteSecretLockData) {
                    null -> RemoteSecretProtectionState.INACTIVE
                    else -> RemoteSecretProtectionState.ACTIVE
                }
            }

    val remoteSecretParametersFlow: StateFlow<RemoteSecretParameters?> =
        lockStateFlow.mapState { lockState ->
            (lockState as? RemoteSecretLockState)?.remoteSecretLockData?.parameters
        }

    val passphraseLockedFlow: StateFlow<Boolean> =
        lockStateFlow.mapState { lockState ->
            lockState is LockState.LockedWithPassphrase
        }

    fun setLockedWithPassphrase(lockData: MasterKeyState.WithPassphrase) {
        lockState = LockState.LockedWithPassphrase(lockData)
    }

    fun setLockedWithRemoteSecret(lockData: MasterKeyState.WithRemoteSecret) {
        lockState = LockState.LockedWithRemoteSecret(lockData)
    }

    fun setUnlocked(masterKey: InvalidateableMasterKey, remoteSecretLockData: MasterKeyState.WithRemoteSecret?) {
        lockState = LockState.Unlocked(masterKey, remoteSecretLockData)
    }

    fun setPermanentlyLocked() {
        lockState = LockState.PermanentlyLocked
    }

    fun isLockedWithPassphrase() =
        when (lockState) {
            LockState.Unknown -> error("Lock state not yet known")
            is LockState.LockedWithPassphrase -> true
            else -> false
        }

    suspend fun isLockedWithRemoteSecret() =
        awaitRemoteSecretLockState() is LockState.LockedWithRemoteSecret

    suspend fun awaitRemoteSecretLockState(): RemoteSecretLockState =
        lockStateFlow
            .filterIsInstance<RemoteSecretLockState>()
            .first()

    fun getRemoteSecretLockState(): RemoteSecretLockState? =
        lockState as? RemoteSecretLockState

    fun getMasterKey(): MasterKey? =
        (lockState as? LockState.Unlocked)?.masterKey

    fun getPassphraseLock(): MasterKeyState.WithPassphrase? =
        (lockState as? LockState.LockedWithPassphrase)?.lockData

    private sealed class LockState {
        data object Unknown : LockState()

        data class LockedWithPassphrase(
            val lockData: MasterKeyState.WithPassphrase,
        ) : LockState()

        data class LockedWithRemoteSecret(
            val lockData: MasterKeyState.WithRemoteSecret,
        ) : LockState(), RemoteSecretLockState {
            override val remoteSecretLockData: MasterKeyState.WithRemoteSecret
                get() = lockData
        }

        data object PermanentlyLocked : LockState()

        class Unlocked(
            val masterKey: InvalidateableMasterKey,
            override val remoteSecretLockData: MasterKeyState.WithRemoteSecret?,
        ) : LockState(), RemoteSecretLockState
    }

    sealed interface RemoteSecretLockState {
        /**
         * If this is non-null, it indicates that the master key is protected by a remote secret, but not necessarily locked with it.
         * I.e., it means that this lock state is a remote secret lock itself or was reached by unlocking a remote secret lock.
         * When a lock state with remote secret lock data becomes active, its parameters must be used to monitor the remote secret.
         */
        val remoteSecretLockData: MasterKeyState.WithRemoteSecret?
    }
}
