package ru.dbotthepony.kstarbound.io

import it.unimi.dsi.fastutil.ints.IntArraySet
import java.io.*
import java.util.*
import java.util.concurrent.locks.ReentrantLock
import kotlin.concurrent.withLock

private fun readHeader(reader: RandomAccessFile, required: Char) {
	val read = reader.read()
	require(read.toChar() == required) { "Bad Starbound Pak header, expected ${required.code}, got $read" }
}

enum class TreeBlockType(val identity: String) {
	INDEX("II"),
	LEAF("LL"),
	FREE("FF");

	companion object {
		operator fun get(index: String): TreeBlockType {
			return when (index) {
				INDEX.identity -> INDEX
				LEAF.identity -> LEAF
				FREE.identity -> FREE
				else -> throw NoSuchElementException("Unknown block type $index")
			}
		}
	}
}

private operator fun ByteArray.compareTo(b: ByteArray): Int {
	require(size == b.size) { "Keys are not of same size (${size} vs ${b.size})" }

	for (i in indices) {
		if (this[i].toInt() and 0xFF > b[i].toInt() and 0xFF) {
			return 1
		} else if (this[i].toInt() and 0xFF < b[i].toInt() and 0xFF) {
			return -1
		}
	}

	return 0
}

/**
 * Класс, который позволяет читать и записывать в файлы Starbound BTReeDB 5
 *
 * Big credit for https://github.com/blixt/py-starbound/blob/master/FORMATS.md#btreedb5 !
 */
class BTreeDB(val path: File) {
	val reader = RandomAccessFile(path, "r")
	private val lock = ReentrantLock()

	init {
		readHeader(reader, 'B')
		readHeader(reader, 'T')
		readHeader(reader, 'r')
		readHeader(reader, 'e')
		readHeader(reader, 'e')
		readHeader(reader, 'D')
		readHeader(reader, 'B')
		readHeader(reader, '5')
	}

	val blockSize = reader.readInt()
	val dbNameRaw = ByteArray(16).also { reader.read(it) }
	val indexKeySize = reader.readInt()
	val useNodeTwo = reader.readBoolean()
	val freeNodeIndex1 = reader.readInt().toLong()

	init { reader.skipBytes(4) }
	val freeBlockOffset1 = reader.readInt().toLong()
	val rootNode1Index = reader.readInt().toLong()
	val rootNode1IsLeaf = reader.readBoolean()
	val freeNodeIndex2 = reader.readInt().toLong()
	init { reader.skipBytes(4) }
	val freeBlockOffset2 = reader.readInt().toLong()
	val rootNode2Index = reader.readInt().toLong()
	val rootNode2IsLeaf = reader.readBoolean()
	init { reader.skipBytes(445) }

	val blocksOffsetStart = reader.filePointer

	init {
		// check(reader.length() - blocksOffsetStart == 512L) { "Unexpected header size of ${reader.length() - blocksOffsetStart} bytes" }
		check((reader.length() - 512L) % blockSize == 0L) { "Junk data somewhere in file (${(reader.length() - 512L) % blockSize} lingering bytes)" }
	}

	val rootNodeIndex get() = if (useNodeTwo) rootNode2Index else rootNode1Index
	val rootNodeIsLeaf get() = if (useNodeTwo) rootNode2IsLeaf else rootNode1IsLeaf

	fun readBlockType() = lock.withLock { TreeBlockType[reader.readString(2)] }

	fun findAllKeys(index: Long = rootNodeIndex): List<ByteArray> {
		lock.withLock {
			seekBlock(index)

			val list = ArrayList<ByteArray>()
			val type = readBlockType()

			if (type == TreeBlockType.LEAF) {
				val keyAmount = reader.readInt()
				// offset внутри лепестка в байтах
				var offset = 6

				for (i in 0 until keyAmount) {
					// читаем ключ
					list.add(ByteArray(indexKeySize).also { reader.read(it) })
					offset += indexKeySize

					// читаем размер данных внутри ключа
					var (dataLength, readBytes) = reader.readVarIntInfo()
					offset += readBytes

					while (true) {
						// если конец данных внутри текущего блока, останавливаемся
						if (offset + dataLength <= blockSize - 4) {
							reader.skipBytes(dataLength)
							offset += dataLength
							break
						}

						// иначе, ищем следующий блок

						// пропускаем оставшиеся данные, переходим на границу текущего блока-лепестка
						val delta = (blockSize - 4 - offset)
						reader.skipBytes(delta)

						// ищем следующий блок с нашими данными
						val nextBlockIndex = reader.readInt()
						seekBlock(nextBlockIndex.toLong())

						// удостоверяемся что мы попали в лепесток
						check(readBlockType() == TreeBlockType.LEAF) { "Did not hit leaf block" }
						offset = 2
						dataLength -= delta
					}
				}
			} else if (type == TreeBlockType.INDEX) {
				reader.skipBytes(1)
				val keyAmount = reader.readInt()

				val blockList = IntArraySet()
				blockList.add(reader.readInt())

				for (i in 0 until keyAmount) {
					// ключ
					reader.skipBytes(indexKeySize)

					// указатель на блок
					blockList.add(reader.readInt())
				}

				// читаем все дочерние блоки на ключи
				for (block in blockList.intIterator()) {
					for (key in findAllKeys(block.toLong())) {
						list.add(key)
					}
				}
			}

			return list
		}
	}

	fun read(key: ByteArray): ByteArray? {
		require(key.size == indexKeySize) { "Key provided is ${key.size} in size, while $indexKeySize is required" }

		lock.withLock {
			seekBlock(rootNodeIndex)
			var type = readBlockType()
			var iterations = 1000

			val keyLoader = ByteArray(indexKeySize)

			// сканирование индекса
			while (iterations-- > 0 && type != TreeBlockType.LEAF) {
				if (type == TreeBlockType.FREE) {
					throw IllegalStateException("Hit free block while scanning index for ${key.joinToString(", ")}")
				}

				reader.skipBytes(1)

				val keyCount = reader.readInt()
				// if keyAmount == 4 then
				// B a B b B c B d B
				val readKeys = ByteArray((keyCount + 1) * 4 + keyCount * indexKeySize)
				reader.readFully(readKeys)

				val stream = DataInputStream(ByteArrayInputStream(readKeys))

				var read = false

				// B a
				// B b
				// B c
				// B d
				for (keyIndex in 0 until keyCount) {
					// указатель на левый блок
					val pointer = stream.readInt()

					// левый ключ, всё что меньше него находится в левом блоке
					stream.readFully(keyLoader)

					// нужный ключ меньше самого первого ключа, поэтому он находится где то в левом блоке
					if (key < keyLoader) {
						seekBlock(pointer.toLong())
						type = readBlockType()
						read = true
						break
					}
				}

				if (!read) {
					// ... B
					seekBlock(stream.readInt().toLong())
					type = readBlockType()
				}
			}

			// мы пришли в лепесток, теперь прямолинейно ищем в linked list
			val leafStream = DataInputStream(BufferedInputStream(LeafInputStream(2)))
			val keyCount = leafStream.readInt()

			for (keyIndex in 0 until keyCount) {
				// читаем ключ
				leafStream.read(keyLoader)

				// читаем размер данных
				val dataLength = leafStream.readVarInt()

				// это наш блок
				if (keyLoader.contentEquals(key)) {
					val binary = ByteArray(dataLength)

					if (dataLength == 0) {
						// нет данных (?)
						return binary
					}

					leafStream.readFully(binary)

					return binary
				} else {
					leafStream.skipBytes(dataLength)
				}
			}

			return null
		}
	}

	fun seekBlock(id: Long) {
		require(id >= 0) { "Negative id $id" }
		require(id * blockSize + blocksOffsetStart < reader.length()) { "Tried to seek block with $id, but it is outside of file's bounds (file size ${reader.length()} bytes, seeking ${id * blockSize + blocksOffsetStart})! (does not exist)" }

		lock.withLock {
			reader.seek(id * blockSize + blocksOffsetStart)
		}
	}

	private inner class LeafInputStream(private var offset: Int) : InputStream() {
		private var canRead = true

		override fun read(): Int {
			if (offset + 4 >= blockSize) {
				if (!seekNextBlock()) {
					return -1
				}
			}

			offset++
			return reader.read()
		}

		override fun read(b: ByteArray, off: Int, len: Int): Int {
			Objects.checkFromIndexSize(off, len, b.size)
			var totalRead = 0
			var index = off

			while (canRead && totalRead < len) {
				if (offset + 4 >= blockSize) {
					if (!seekNextBlock()) {
						return totalRead
					}
				}

				val readAtMost = (blockSize - 4 - offset).coerceAtMost(len - totalRead)
				val readBytes = reader.read(b, index, readAtMost)

				if (readBytes <= 0) {
					canRead = false
					break
				}

				totalRead += readBytes
				index += readBytes
				offset += readBytes
			}

			return totalRead
		}

		private fun seekNextBlock(): Boolean {
			if (!canRead) {
				return false
			}

			val nextBlockIndex = reader.readInt()

			if (nextBlockIndex < 0L) {
				canRead = false
				return false
			}

			seekBlock(nextBlockIndex.toLong())

			check(readBlockType() == TreeBlockType.LEAF) { "Did not hit leaf block" }
			offset = 2

			return true
		}
	}
}