package com.infendro.encoding

import com.infendro.encoding.util.lcm
import com.infendro.encoding.util.log2
import kotlin.math.pow

sealed class Encoding {
    abstract fun encode(
        bytes: ByteArray,
    ): ByteArray

    abstract fun decode(
        bytes: ByteArray,
    ): ByteArray

    protected fun encode(
        bytes: ByteArray,
        characters: Array<Byte>,
        padding: Byte? = null,
    ): ByteArray {
        val length = log2(characters.size).toInt()
        val mask = (2.0.pow(length) - 1).toUInt()

        val encoded = mutableListOf<Byte>()

        var buffer = 0U
        var bits = 0

        for (byte in bytes) {
            buffer = (buffer shl 8) + byte.toUByte()
            bits += 8

            while (bits >= length) {
                val offset = bits - length
                val index = (buffer shr offset).toInt()
                buffer = buffer and (mask shl offset).inv()
                bits -= length

                encoded.add(characters[index])
            }
        }
        if (bits > 0) {
            buffer = buffer shl (length - bits)
            val index = buffer.toInt()
            encoded.add(characters[index])
        }

        if (padding != null) {
            val block = lcm(length, 8) / length
            while (encoded.size % block != 0) {
                encoded.add(padding)
            }
        }

        return encoded.toByteArray()
    }

    protected fun decode(
        bytes: ByteArray,
        characters: Array<Byte>,
        padding: Byte? = null,
    ): ByteArray {
        val length = log2(characters.size).toInt()
        val mask = 0xFFU

        val decoded = mutableListOf<Byte>()

        var buffer = 0U
        var bits = 0

        for (byte in bytes) {
            if (byte == padding) break

            val index = characters.indexOf(byte)
            if (index == -1) throw Exception()

            buffer = (buffer shl length) + index.toUByte()
            bits += length

            if (bits >= 8) {
                val offset = bits - 8
                val character = (buffer shr offset).toByte()
                buffer = buffer and (mask shl offset).inv()
                bits -= 8

                decoded.add(character)
            }
        }

        return decoded.toByteArray()
    }
}