package com.infendro.hash.keccak

import com.infendro.bytearray.ByteOrder.LITTLE_ENDIAN
import com.infendro.bytearray.addAll
import com.infendro.bytearray.toByteArray
import com.infendro.bytearray.toULongArray
import com.infendro.hash.HashFunction
import kotlin.math.min

@Suppress("UNUSED")
open class Keccak internal constructor(
    override val bytes: Int,
    private val rate: Int,
    private val domain: Byte,
) : HashFunction() {
    init {
        if (rate % 8 != 0) {
            throw IllegalArgumentException("rate must be divisible by 8")
        }
    }

    private val rc = ulongArrayOf(
        0x0000000000000001UL, 0x0000000000008082UL, 0x800000000000808aUL, 0x8000000080008000UL,
        0x000000000000808bUL, 0x0000000080000001UL, 0x8000000080008081UL, 0x8000000000008009UL,
        0x000000000000008aUL, 0x0000000000000088UL, 0x0000000080008009UL, 0x000000008000000aUL,
        0x000000008000808bUL, 0x800000000000008bUL, 0x8000000000008089UL, 0x8000000000008003UL,
        0x8000000000008002UL, 0x8000000000000080UL, 0x000000000000800aUL, 0x800000008000000aUL,
        0x8000000080008081UL, 0x8000000000008080UL, 0x0000000080000001UL, 0x8000000080008008UL
    )

    private val rotation = arrayOf(
        intArrayOf(0, 36, 3, 41, 18),
        intArrayOf(1, 44, 10, 45, 2),
        intArrayOf(62, 6, 43, 15, 61),
        intArrayOf(28, 55, 25, 21, 56),
        intArrayOf(27, 20, 39, 8, 14),
    )

    override fun hash(
        value: ByteArray,
    ): ByteArray {
        val paddedValue = buildList {
            addAll(value)
            add(domain)
            while ((size % rate) != (rate - 1)) {
                add(0x00.toByte())
            }
            add(0x80.toByte())
        }

        val state = ULongArray(25)

        // absorb
        val chunks = paddedValue
            .chunked(rate)
        for (chunk in chunks) {
            val numbers = chunk.toULongArray(LITTLE_ENDIAN)
            for ((i, number) in numbers.withIndex()) {
                state[i] = state[i] xor number
            }
            permute(state)
        }

        // squeeze
        return buildList(bytes) {
            var i = 0
            while (size < bytes) {
                if (i == rate / 8) {
                    permute(state)
                    i = 0
                }
                addAll(
                    state[i++].toByteArray(LITTLE_ENDIAN)
                        .take(min(bytes - size, 8))
                )
            }
        }.toByteArray()
    }

    private fun permute(
        a: ULongArray,
    ) {
        val b = ULongArray(25)
        val c = ULongArray(5)
        val d = ULongArray(5)

        repeat(24) { round ->
            // θ
            for (x in 0..4) {
                c[x] = a[x] xor a[x + 5] xor a[x + 10] xor a[x + 15] xor a[x + 20]
            }
            for (x in 0..4) {
                d[x] = c[(x + 4) % 5] xor c[(x + 1) % 5].rotateLeft(1)
            }
            for (x in 0..4) for (y in 0..4) {
                a[x + 5 * y] = a[x + 5 * y] xor d[x]
            }

            // ρ and π
            for (x in 0..4) for (y in 0..4) {
                b[y + 5 * ((2 * x + 3 * y) % 5)] = a[x + 5 * y].rotateLeft(rotation[x][y])
            }

            // χ
            for (x in 0..4) for (y in 0..4) {
                a[x + 5 * y] = b[x + 5 * y] xor ((b[((x + 1) % 5) + 5 * y].inv()) and b[((x + 2) % 5) + 5 * y])
            }

            // ι
            a[0] = a[0] xor rc[round]
        }
    }
}