package com.nisemoe.nise.replays

import java.io.ByteArrayOutputStream
import java.nio.ByteBuffer
import java.nio.ByteOrder
import kotlin.math.abs
import kotlin.math.round

// JVM implementation of https://github.com/circleguard/wtc-lzma-compressor/tree/master



class WTC {
    companion object {
        private const val CURRENT_VERSION_HEADER: Short = 1

        private val VERSION_HEADER_BYTE_ARRAY = byteArrayOf((CURRENT_VERSION_HEADER.toInt() and 0xFF).toByte(), ((CURRENT_VERSION_HEADER.toInt() shr 8) and 0xFF).toByte())

        fun compress(stream: String): ByteArray {
            val lists = seperate(stream)

            val xs = unsortedDiffPackShortsToBytes(lists.x)
            val ys = unsortedDiffPackShortsToBytes(lists.y)

            val ws = packIntsToBytes(lists.w)
            val zs = lists.z

            fun packBytes(arr: ByteArray): ByteArray {
                val length = arr.size
                val buffer = ByteBuffer.allocate(4 + length).order(ByteOrder.LITTLE_ENDIAN)

                buffer.putInt(length)
                for (byte in arr) {
                    buffer.put(byte)
                }

                return buffer.array()
            }

            val byteStream = ByteArrayOutputStream()

            byteStream.writeBytes(VERSION_HEADER_BYTE_ARRAY)
            byteStream.writeBytes(packBytes(xs))
            byteStream.writeBytes(packBytes(ys))
            byteStream.writeBytes(packBytes(zs))
            byteStream.writeBytes(packBytes(ws))

            return byteStream.toByteArray()
        }

        fun decompress(data: ByteArray, hasVersionHeader: Boolean = true): String {
            val buffer = ByteBuffer.wrap(data).order(ByteOrder.LITTLE_ENDIAN)

            fun unpackBytes(): ByteArray {
                val size = buffer.getInt()

                val bytes = ByteArray(size)
                buffer.get(bytes, 0, size)

                return bytes
            }

            if (hasVersionHeader) {
                buffer.getShort() // Version - may be used in the future
            }

            val xs = unpackBytes()
            val ys = unpackBytes()
            val zs = unpackBytes()
            val ws = unpackBytes()

            val xxs = unsortedDiffUnpackBytesToShorts(xs)
            val yys = unsortedDiffUnpackBytesToShorts(ys)

            val wws = unpackBytesToInts(ws)

            return combine(FrameLists(
                x = xxs,
                y = yys,
                z = zs,
                w = wws,
            ))
        }

        data class FrameLists(
            val x: ShortArray,
            val y: ShortArray,
            val z: ByteArray,
            val w: IntArray,
        )

        private fun unsortedDiffPackShortsToBytes(shorts: ShortArray): ByteArray {
            val start = shorts.first()
            val diff = arrayDiff(shorts)
            val packed = mutableListOf<Byte>()

            fun pack(word: Short) {
                if (abs(word.toInt()) <= Byte.MAX_VALUE) {
                    packed.add(word.toByte())
                }
                else {
                    packed.add(Byte.MIN_VALUE)
                    packed.add((word.toInt() and 0xFF).toByte())
                    packed.add((word.toInt() shr 8).toByte())
                }
            }

            pack(start)
            for (word in diff) {
                pack(word)
            }

            return packed.toByteArray()
        }

        private fun unsortedDiffUnpackBytesToShorts(int8s: ByteArray): ShortArray {
            val decoded = mutableListOf<Short>()

            var i = 0
            while (i < int8s.size) {
                val byte = int8s[i]

                if (byte == Byte.MIN_VALUE) {
                    i++
                    var word = int8s[i].toInt() and 0xFF
                    i++
                    word += int8s[i].toInt() shl 8
                    decoded.add(word.toShort())
                }
                else {
                    decoded.add(byte.toShort())
                }

                i++
            }

            return cumSum(decoded.toShortArray())
        }

        private fun packIntsToBytes(int32s: IntArray): ByteArray {
            val packed = mutableListOf<Byte>()

            for (dw in int32s) {
                var dword = dw
                if (abs(dword) <= Byte.MAX_VALUE) {
                    packed.add(dword.toByte())
                }
                else {
                    packed.add(Byte.MIN_VALUE)
                    packed.add((dword and 0xFF).toByte())
                    dword = dword shr 8
                    packed.add((dword and 0xFF).toByte())
                    dword = dword shr 8
                    packed.add((dword and 0xFF).toByte())
                    dword = dword shr 8
                    packed.add(dword.toByte())
                }
            }

            return packed.toByteArray()
        }

        private fun unpackBytesToInts(int8s: ByteArray): IntArray {
            val unpacked = mutableListOf<Int>()

            var i = 0
            while (i < int8s.size) {
                val byte = int8s[i]

                if (byte == Byte.MIN_VALUE) {
                    i++
                    var dword = int8s[i].toInt() and 0xFF
                    i++
                    dword += (int8s[i].toInt() shl 8) and 0xFF00
                    i++
                    dword += (int8s[i].toInt() shl 16) and 0xFF0000
                    i++
                    dword += int8s[i].toInt() shl 24
                    unpacked.add(dword)
                }
                else {
                    unpacked.add(byte.toInt())
                }

                i++
            }

            return unpacked.toIntArray()
        }

        private fun seperate(stream: String): FrameLists {
            val frames = stream.split(',')
            val frameCount = frames.size

            val ws = IntArray(frameCount)
            val xs = ShortArray(frameCount)
            val ys = ShortArray(frameCount)
            val zs = ByteArray(frameCount)

            for ((i, frame) in frames.withIndex()) {
                if (frame.isEmpty()) {
                    continue
                }

                val splitFrame = frame.split('|')
                val w = splitFrame[0].toInt()
                val x = splitFrame[1].toFloat()
                val y = splitFrame[2].toFloat()
                val z = splitFrame[3].toInt()

                val zz = z and 0xFF

                var xx = round(x * 16).toInt()
                var yy = round(y * 16).toInt()

                if (xx <= -0x8000) xx = -0x8000
                else if (xx >= 0x7FFF) xx = 0x7FFF

                if (yy <= -0x8000) yy = -0x8000
                else if (yy >= 0x7FFF) yy = 0x7FFF

                ws[i] = w
                xs[i] = xx.toShort()
                ys[i] = yy.toShort()
                zs[i] = zz.toByte()
            }

            return FrameLists(
                x = xs,
                y = ys,
                z = zs,
                w = ws,
            )
        }

        private fun combine(lists: FrameLists): String {
            val xArr = lists.x.map { it.toFloat() / 16 }

            val yArr = lists.y.map { it.toFloat() / 16 }

            val frames = arrayOfNulls<String>(xArr.size)

            for (i in xArr.indices) {
                val x = xArr[i]
                val y = yArr[i]
                val z = lists.z[i]
                val w = lists.w[i]
                frames[i] = "$w|$x|$y|$z"
            }

            return frames.joinToString(",")
        }

        private fun arrayDiff(arr: ShortArray): ShortArray {
            if (arr.isEmpty()) {
                return emptyArray<Short>().toShortArray()
            }

            val diffed = ShortArray(arr.size - 1)

            for (index in 1..<arr.size) {
                diffed[index - 1] = (arr[index] - arr[index - 1]).toShort()
            }

            return diffed
        }

        private fun cumSum(arr: ShortArray): ShortArray {
            if (arr.isEmpty()) {
                return emptyArray<Short>().toShortArray()
            }

            val cumArr = ShortArray(arr.size)

            cumArr[0] = arr.first()
            for (index in 1..<arr.size) {
                cumArr[index] = (arr[index] + cumArr[index - 1]).toShort()
            }

            return cumArr
        }
    }
}
