package com.infendro.hash.blake2

import com.infendro.bytes.ByteOrder.LITTLE_ENDIAN
import com.infendro.bytes.bytearray.copyInto
import com.infendro.bytes.longarray.toByteArray
import com.infendro.hash.HashFunction

class Blake2b(bytes: Int) : HashFunction(bytes, 128) {
    init {
        require(bytes in 1..64)
    }

    private val h = initial.copyOf()
        .also { h ->
            h[0] = h[0] xor (0x01010000L or bytes.toLong())
        }

    private val m = LongArray(16)
    private val v = LongArray(16)

    override fun reset() {
        super.reset()
        initial.copyInto(h)
        h[0] = h[0] xor (0x01010000L or bytes.toLong())
    }

    override fun update(byte: Byte) {
        if (buffer.isFull()) {
            process()
            buffer.reset()
        }
        buffer.append(byte)
        length++
    }

    override fun update(bytes: ByteArray) {
        if (buffer.isFull() && bytes.isNotEmpty()) {
            process()
            buffer.reset()
        }

        var i = 0

        var free = buffer.free
        while (i + free < bytes.size) {
            buffer.append(bytes, startIndex = i, endIndex = i + free)
            length += free
            process()
            buffer.reset()
            free = buffer.free
            i += block
        }

        buffer.append(bytes, startIndex = i)
        length += (bytes.size - i)
    }

    override fun process() = compress(buffer.value, length, false)

    private fun finalize() = compress(buffer.value, length, true)

    override fun digestInto(destination: ByteArray, destinationOffset: Int) {
        finalize()
        h.toByteArray(order = LITTLE_ENDIAN).copyInto(destination, destinationOffset, endIndex = bytes)
        reset()
    }

    private fun compress(block: ByteArray, t: Long, last: Boolean) {
        block.copyInto(m)
        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(12) { r ->
            val s = sigma[r % 10]

            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..<8) {
            h[i] = h[i] xor v[i] xor v[i + 8]
        }
    }

    private fun mix(v: LongArray, a: Int, b: Int, c: Int, d: Int, x: Long, y: Long) {
        v[a] += v[b] + x
        v[d] = (v[d] xor v[a]).rotateRight(32)
        v[c] += v[d]
        v[b] = (v[b] xor v[c]).rotateRight(24)
        v[a] += v[b] + y
        v[d] = (v[d] xor v[a]).rotateRight(16)
        v[c] += v[d]
        v[b] = (v[b] xor v[c]).rotateRight(63)
    }

    companion object {
        private val initial = longArrayOf(
            +0x6a09e667f3bcc908L, -0x4498517a7b3558c5L,
            +0x3c6ef372fe94f82bL, -0x5ab00ac5a0e2c90fL,
            +0x510e527fade682d1L, -0x64fa9773d4c193e1L,
            +0x1f83d9abfb41bd6bL, +0x5be0cd19137e2179L,
        )

        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),
        )
    }
}