package com.infendro.encoding

import com.infendro.encoding.exception.IllegalByteException
import com.infendro.encoding.util.align
import com.infendro.encoding.util.ceilDiv
import com.infendro.encoding.util.lcm
import com.infendro.encoding.util.log2

class StreamEncoding internal constructor(
    private val alphabet: ByteArray,
    private val padding: Byte? = null,
) : Encoding {
    private val base: Int = alphabet.size
    private val bits: Int = log2(base).toInt()
    private val block: Int = lcm(bits, 8) / bits
    private val mask: Int = (0x01 shl bits) - 1

    override fun encode(bytes: ByteArray): ByteArray {
        val size = when (padding) {
            null -> (bytes.size * 8).ceilDiv(bits)
            else -> (bytes.size * 8).ceilDiv(bits).align(block)
        }
        val encoded = ByteArray(size)
        var i = 0

        var buffer = 0
        var bufferBits = 0
        for (byte in bytes) {
            buffer = (buffer shl 8) or (byte.toInt() and 0xFF)
            bufferBits += 8

            while (bufferBits >= bits) {
                val index = (buffer shr (bufferBits - bits)) and mask
                bufferBits -= bits
                encoded[i++] = alphabet[index]
            }
        }
        if (bufferBits > 0) {
            val index = (buffer shl (bits - bufferBits)) and mask
            encoded[i++] = alphabet[index]
        }

        while (i < encoded.size) {
            encoded[i++] = padding!!
        }

        return encoded
    }

    override fun decode(bytes: ByteArray): ByteArray {
        val end = when (padding) {
            null -> bytes.size
            else -> when (val index = bytes.indexOf(padding)) {
                -1 -> bytes.size
                else -> index
            }
        }
        val size = end * bits / 8
        val decoded = ByteArray(size)
        var i = 0

        var buffer = 0
        var bufferBits = 0
        for (b in 0..<end) {
            val byte = bytes[b]
            val index = alphabet.indexOf(byte)
            if (index == -1) throw IllegalByteException(byte)
            buffer = (buffer shl bits) or index
            bufferBits += bits

            if (bufferBits >= 8) {
                val byte = (buffer shr (bufferBits - 8)) and 0xFF
                bufferBits -= 8
                decoded[i++] = byte.toByte()
            }
        }

        return decoded
    }
}