package com.infendro.hash.blake2

import com.infendro.bytearray.ByteOrder.LITTLE_ENDIAN
import com.infendro.bytearray.toUByteArray
import com.infendro.bytearray.toUIntArray
import com.infendro.hash.HashFunction

class Blake2s(
    override val bytes: Int,
) : HashFunction() {
    init {
        if (bytes !in 1..32) throw IllegalArgumentException()
    }

    override val block: Int
        get() = 64

    private val initial = uintArrayOf(
        0x6a09e667U, 0xbb67ae85U,
        0x3c6ef372U, 0xa54ff53aU,
        0x510e527fU, 0x9b05688cU,
        0x1f83d9abU, 0x5be0cd19U,
    )

    private val sigma = arrayOf(
        intArrayOf(0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15),
        intArrayOf(14, 10, 4, 8, 9, 15, 13, 6, 1, 12, 0, 2, 11, 7, 5, 3),
        intArrayOf(11, 8, 12, 0, 5, 2, 15, 13, 10, 14, 3, 6, 7, 1, 9, 4),
        intArrayOf(7, 9, 3, 1, 13, 12, 11, 14, 2, 6, 5, 10, 4, 0, 15, 8),
        intArrayOf(9, 0, 5, 7, 2, 4, 10, 15, 14, 1, 11, 12, 6, 8, 3, 13),
        intArrayOf(2, 12, 6, 10, 0, 11, 8, 3, 4, 13, 7, 5, 15, 14, 1, 9),
        intArrayOf(12, 5, 1, 15, 14, 13, 4, 10, 0, 7, 6, 3, 9, 2, 8, 11),
        intArrayOf(13, 11, 7, 14, 12, 1, 3, 9, 5, 0, 15, 4, 8, 6, 2, 10),
        intArrayOf(6, 15, 14, 9, 11, 3, 0, 8, 12, 2, 13, 7, 1, 4, 10, 5),
        intArrayOf(10, 2, 8, 4, 7, 6, 1, 5, 15, 11, 9, 14, 3, 12, 13, 0),
    )

    override fun hash(value: UByteArray): UByteArray {
        val l = value.size.toULong()
        val h = initial.copyOf()
        h[0] = h[0] xor (0x01010000U or bytes.toUInt())
        var t = 0U

        val blocks = pad(value)
            .chunked(64) { it.toUByteArray() }
        for ((i, block) in blocks.withIndex()) {
            t += minOf(l - t, 64UL).toUInt()
            compress(h, block, t, i == blocks.lastIndex)
        }

        return h.toUByteArray(LITTLE_ENDIAN).sliceArray(0..<bytes)
    }

    private fun pad(
        value: UByteArray,
    ): UByteArray {
        return buildList {
            addAll(value)
            if (isEmpty()) add(0x00U)
            while (size % 64 != 0) {
                add(0x00U)
            }
        }.toUByteArray()
    }

    private fun compress(
        h: UIntArray,
        block: UByteArray,
        t: UInt,
        last: Boolean,
    ) {
        val m = block.toUIntArray()
        val v = UIntArray(16)
        for (i in 0..<8) v[i] = h[i]
        for (i in 0..<8) v[i + 8] = initial[i]

        v[12] = v[12] xor t
        if (last) v[14] = v[14].inv()

        repeat(10) { r ->
            val s = sigma[r]

            mix(v, 0, 4, 8, 12, m[s[0]], m[s[1]])
            mix(v, 1, 5, 9, 13, m[s[2]], m[s[3]])
            mix(v, 2, 6, 10, 14, m[s[4]], m[s[5]])
            mix(v, 3, 7, 11, 15, m[s[6]], m[s[7]])

            mix(v, 0, 5, 10, 15, m[s[8]], m[s[9]])
            mix(v, 1, 6, 11, 12, m[s[10]], m[s[11]])
            mix(v, 2, 7, 8, 13, m[s[12]], m[s[13]])
            mix(v, 3, 4, 9, 14, m[s[14]], m[s[15]])
        }

        for (i in 0..7) {
            h[i] = h[i] xor v[i] xor v[i + 8]
        }
    }

    private fun mix(
        v: UIntArray,
        a: Int,
        b: Int,
        c: Int,
        d: Int,
        x: UInt,
        y: UInt,
    ) {
        v[a] += v[b] + x
        v[d] = (v[d] xor v[a]).rotateRight(16)
        v[c] += v[d]
        v[b] = (v[b] xor v[c]).rotateRight(12)
        v[a] += v[b] + y
        v[d] = (v[d] xor v[a]).rotateRight(8)
        v[c] += v[d]
        v[b] = (v[b] xor v[c]).rotateRight(7)
    }
}