package ch.threema.app.voip.groupcall.sfu.webrtc

import ch.threema.app.voip.groupcall.CryptoCallUtils
import ch.threema.common.generateRandomBytes
import ch.threema.common.secureRandom
import ch.threema.domain.protocol.csp.ProtocolDefines.GC_PCMK_LENGTH

/**
 * State snapshot of a participant's media key (PCMK).
 *
 * All properties are to be considered immutable.
 */
data class ParticipantCallMediaKeyState(
    val epoch: UInt,
    val ratchetCounter: UInt,
    val pcmk: ByteArray,
) {
    override fun toString(): String = "ParticipantCallMediaKeyState(epoch=$epoch, ratchetCounter=$ratchetCounter)"

    override fun equals(other: Any?): Boolean {
        if (this === other) return true
        if (javaClass != other?.javaClass) return false

        other as ParticipantCallMediaKeyState

        if (epoch != other.epoch) return false
        if (ratchetCounter != other.ratchetCounter) return false
        if (!pcmk.contentEquals(other.pcmk)) return false

        return true
    }

    override fun hashCode(): Int {
        var result = epoch.hashCode()
        result = 31 * result + ratchetCounter.hashCode()
        result = 31 * result + pcmk.contentHashCode()
        return result
    }
}

/**
 * A participant's media key (PCMK) that will be applied soon.
 */
class PendingParticipantCallMediaKeyState(
    private val creator: LocalParticipantCallMediaKey,
    private var wasApplied: Boolean = false,
    internal var stale: Boolean = false,
    val state: ParticipantCallMediaKeyState,
) {
    fun applied() {
        assert(!wasApplied)
        wasApplied = true
        creator.applied(this)
    }
}

class LocalParticipantCallMediaKey {
    private var _current: ParticipantCallMediaKeyState = ParticipantCallMediaKeyState(
        epoch = 0u,
        ratchetCounter = 0u,
        pcmk = secureRandom().generateRandomBytes(GC_PCMK_LENGTH),
    )
    private var _pending: PendingParticipantCallMediaKeyState? = null

    val current: ParticipantCallMediaKeyState
        get() = _current
    val pending: PendingParticipantCallMediaKeyState?
        get() = _pending

    fun all(): List<ParticipantCallMediaKeyState> {
        return pending.let { if (it === null) listOf(current) else listOf(current, it.state) }
    }

    fun nextEpoch(): PendingParticipantCallMediaKeyState {
        // If another successor state is pending to be applied, mark it as _stale_ and abort.
        _pending?.let {
            it.stale = true
            return it
        }

        // Increase the epoch. Note that it is allowed to wrap.
        return _current.let { currentKeyState ->
            val state = ParticipantCallMediaKeyState(
                epoch = if (currentKeyState.epoch == 255u) 0u else currentKeyState.epoch + 1u,
                ratchetCounter = 0u,
                pcmk = secureRandom().generateRandomBytes(GC_PCMK_LENGTH),
            )
            PendingParticipantCallMediaKeyState(
                state = state,
                creator = this,
            ).also {
                _pending = it
            }
        }
    }

    fun nextRatchetCounter(): ParticipantCallMediaKeyState {
        // Note: The ratchet will be applied to the current PCMK, even when a
        // successor state is pending. This is perfectly fine though.
        _current = _current.let {
            // Ensure the ratchet counter does not overflow
            // Note: This is considered unreachable
            if (it.ratchetCounter == 255u) {
                throw Error("Ratchet counter would overflow")
            }

            // Ratchet once
            //
            // PCMK' = BLAKE2b(key=PCMK, salt="m'", personal='3ma-call')
            val pcmk = CryptoCallUtils.gcBlake2b256(
                key = it.pcmk,
                salt = CryptoCallUtils.SALT_CURRENT_PCMK,
            )
            ParticipantCallMediaKeyState(
                epoch = it.epoch,
                ratchetCounter = it.ratchetCounter + 1u,
                pcmk,
            )
        }
        return _current
    }

    internal fun applied(pending: PendingParticipantCallMediaKeyState) {
        // When the sucessor PCMK has been applied, replace it
        assert(pending === this._pending)
        this._pending = null
        _current = pending.state
    }
}
