package com.infendro.random

import com.infendro.bytearray.addAll
import com.infendro.hash.HashFunction

class HMAC internal constructor(
    private val function: HashFunction,
) : PRNG() {
    private val random = Default
    private val hmac = com.infendro.mac.HMAC(function)
    private lateinit var K: ByteArray
    private lateinit var V: ByteArray

    init {
        val bytes = random.nextBytes(function.bytes)
        seed(bytes)
    }

    override fun seed(
        seed: ByteArray,
    ) {
        K = ByteArray(function.bytes) { 0x00 }
        V = ByteArray(function.bytes) { 0x01 }
        reseed(seed)
    }

    override fun reseed(
        entropy: ByteArray,
    ) {
        K = hmac.hash(K, V + byteArrayOf(0x00) + entropy)
        V = hmac.hash(K, V)
        K = hmac.hash(K, V + byteArrayOf(0x01) + entropy)
        V = hmac.hash(K, V)
    }

    override fun generate(
        count: Int,
    ): ByteArray {
        val bytes = buildList {
            while (size < count) {
                V = hmac.hash(K, V)
                addAll(V)
            }
        }

        K = hmac.hash(K, V + byteArrayOf(0x00))
        V = hmac.hash(K, V)

        return bytes
            .take(count)
            .toByteArray()
    }
}