package com.infendro.random.csprng

import com.infendro.hash.HashFunction
import kotlin.math.min
import kotlin.properties.Delegates.notNull

class HMAC internal constructor(
    private val function: HashFunction,
) : CSPRNG() {
    private val hmac = com.infendro.mac.HMAC(function)

    private var k: UByteArray by notNull()
    private var v: UByteArray by notNull()

    override fun seed(
        seed: UByteArray,
    ) {
        k = UByteArray(function.bytes) { 0x00U }
        v = UByteArray(function.bytes) { 0x01U }
        reseed(seed)
    }

    override fun reseed(
        entropy: UByteArray,
    ) {
        k = hmac.hash(k, v + ubyteArrayOf(0x00U) + entropy)
        v = hmac.hash(k, v)
        k = hmac.hash(k, v + ubyteArrayOf(0x01U) + entropy)
        v = hmac.hash(k, v)
    }

    override fun generate(
        count: Int,
    ): UByteArray {
        val bytes = buildList {
            while (size < count) {
                v = hmac.hash(k, v)
                addAll(v.take(min(count - size, function.bytes)))
            }
        }.toUByteArray()

        k = hmac.hash(k, v + ubyteArrayOf(0x00U))
        v = hmac.hash(k, v)

        return bytes
    }
}