package ru.dbotthepony.kstarbound.server.world

import com.google.gson.JsonObject
import it.unimi.dsi.fastutil.ints.Int2ObjectAVLTreeMap
import it.unimi.dsi.fastutil.io.FastByteArrayInputStream
import it.unimi.dsi.fastutil.io.FastByteArrayOutputStream
import it.unimi.dsi.fastutil.objects.Object2IntFunction
import it.unimi.dsi.fastutil.objects.Object2IntOpenHashMap
import org.apache.logging.log4j.LogManager
import ru.dbotthepony.kommons.arrays.Object2DArray
import ru.dbotthepony.kommons.io.readBinaryString
import ru.dbotthepony.kommons.io.readCollection
import ru.dbotthepony.kommons.io.writeBinaryString
import ru.dbotthepony.kommons.io.writeCollection
import ru.dbotthepony.kstarbound.Starbound
import ru.dbotthepony.kstarbound.VersionRegistry
import ru.dbotthepony.kstarbound.defs.EntityType
import ru.dbotthepony.kstarbound.io.createSavepoint
import ru.dbotthepony.kstarbound.json.VersionedJson
import ru.dbotthepony.kstarbound.json.readJsonElement
import ru.dbotthepony.kstarbound.json.readJsonObjectInflated
import ru.dbotthepony.kstarbound.json.writeJsonElement
import ru.dbotthepony.kstarbound.json.writeJsonObjectDeflated
import ru.dbotthepony.kstarbound.math.vector.Vector2d
import ru.dbotthepony.kstarbound.math.vector.Vector2i
import ru.dbotthepony.kstarbound.util.CarriedExecutor
import ru.dbotthepony.kstarbound.util.supplyAsync
import ru.dbotthepony.kstarbound.world.ChunkPos
import ru.dbotthepony.kstarbound.world.ChunkState
import ru.dbotthepony.kstarbound.world.WorldGeometry
import ru.dbotthepony.kstarbound.world.api.AbstractCell
import ru.dbotthepony.kstarbound.world.api.MutableCell
import ru.dbotthepony.kstarbound.world.entities.AbstractEntity
import java.io.BufferedInputStream
import java.io.BufferedOutputStream
import java.io.DataInputStream
import java.io.DataOutputStream
import java.io.File
import java.lang.ref.Cleaner
import java.sql.Connection
import java.sql.DriverManager
import java.sql.PreparedStatement
import java.util.concurrent.CompletableFuture
import java.util.concurrent.TimeUnit
import java.util.zip.Deflater
import java.util.zip.DeflaterOutputStream
import java.util.zip.Inflater
import java.util.zip.InflaterInputStream

class NativeLocalWorldStorage(file: File?) : WorldStorage() {
	private val connection: Connection
	private val executor = CarriedExecutor(Starbound.IO_EXECUTOR)

	init {
		if (file == null) {
			connection = DriverManager.getConnection("jdbc:sqlite:")
		} else {
			connection = DriverManager.getConnection("jdbc:sqlite:${file.absolutePath.replace('\\', '/')}")
		}
	}

	private val cleaner: Cleaner.Cleanable = Starbound.CLEANER.register(this, connection::close)

	override fun commit() {
		executor.execute { connection.commit() }
	}

	override fun close() {
		executor.execute { connection.commit(); cleaner.clean() }
		executor.wait(300L, TimeUnit.SECONDS)
	}

	init {
		connection.createStatement().use {
			it.execute("PRAGMA journal_mode=WAL")
			it.execute("PRAGMA synchronous=NORMAL")

			it.execute("""
				CREATE TABLE IF NOT EXISTS "metadata" (
					"key" VARCHAR NOT NULL PRIMARY KEY,
					"version" INTEGER NOT NULL,
					"data" BLOB NOT NULL
				)
			""".trimIndent())

			it.execute("""
				CREATE TABLE IF NOT EXISTS "cells" (
					"x" INTEGER NOT NULL,
					"y" INTEGER NOT NULL,
					"generation_stage" INTEGER NOT NULL,
					"version" INTEGER NOT NULL,
					"data" BLOB NOT NULL,
					PRIMARY KEY ("x", "y")
				)
			""".trimIndent())

			it.execute("""
				CREATE TABLE IF NOT EXISTS "entities" (
					-- store chunks because rules for entities belonging to specific chunk might get different over time
					"chunkX" INTEGER NOT NULL,
					"chunkY" INTEGER NOT NULL,
					"x" REAL NOT NULL,
					"y" REAL NOT NULL,
					"unique_id" VARCHAR,
					"type" VARCHAR NOT NULL,
					"version" INTEGER NOT NULL,
					"data" BLOB NOT NULL
				)
			""".trimIndent())

			// Enforce unique-ness of unique entity IDs
			// If another entity pops-up with same unique id, then it will overwrite entity in other chunk
			// (unless colliding entities are both loaded into world's memory, which will cause runtime exception to be thrown;
			// and no overwrite will happen)
			it.execute("""
				CREATE UNIQUE INDEX IF NOT EXISTS "entities_unique_id" ON "entities" ("unique_id")
			""".trimIndent())

			it.execute("""
				CREATE INDEX IF NOT EXISTS "entities_chunk_pos" ON "entities" ("chunkX", "chunkY")
			""".trimIndent())
		}

		connection.autoCommit = false
	}

	private val readCells = connection.prepareStatement("""
		SELECT "generation_stage", "version", "data" FROM "cells" WHERE "x" = ? AND "y" = ?
	""".trimIndent())

	override fun loadCells(pos: ChunkPos): CompletableFuture<ChunkCells?> {
		return executor.supplyAsync {
			readCells.setInt(1, pos.x)
			readCells.setInt(2, pos.y)

			readCells.executeQuery().use {
				if (!it.next())
					return@supplyAsync null

				val stage = ChunkState.entries[it.getInt(1)]
				val version = it.getInt(2)
				val data = it.getBytes(3)

				val inflater = Inflater()
				val stream = DataInputStream(BufferedInputStream(InflaterInputStream(FastByteArrayInputStream(data), inflater, 0x10000), 0x40000))

				try {
					val palette = PaletteSet()
					palette.read(stream)

					val width = stream.readInt()
					val height = stream.readInt()
					val array = Object2DArray.nulls<AbstractCell>(width, height)

					for (x in 0 until width) {
						for (y in 0 until height) {
							array[x, y] = MutableCell().readNative(stream, palette, version).immutable()
						}
					}

					return@supplyAsync ChunkCells(array as Object2DArray<out AbstractCell>, stage)
				} finally {
					inflater.end()
				}
			}
		}
	}

	private val readEntities = connection.prepareStatement("""
		SELECT "type", "version", "data" FROM "entities" WHERE "chunkX" = ? AND "chunkY" = ?
	""".trimIndent())

	override fun loadEntities(pos: ChunkPos): CompletableFuture<Collection<AbstractEntity>> {
		return executor.supplyAsync {
			readEntities.setInt(1, pos.x)
			readEntities.setInt(2, pos.y)

			val entities = ArrayList<AbstractEntity>()

			readEntities.executeQuery().use {
				while (it.next()) {
					val rtype = it.getString(1)
					val version = it.getInt(2)
					val type = EntityType.entries.firstOrNull { it.storeName == rtype }

					if (type != null) {
						try {
							val data = it.getBytes(3).readJsonObjectInflated()
							entities.add(type.fromStorage(VersionRegistry.migrate(VersionedJson(rtype, version, data)) as JsonObject))
						} catch (err: Throwable) {
							LOGGER.error("Unable to deserialize entity in chunk $pos", err)
						}
					} else {
						LOGGER.error("Unknown entity type $rtype in chunk $pos")
					}
				}
			}

			return@supplyAsync entities
		}
	}

	private val readMetadata = connection.prepareStatement("""
		SELECT "version", "data" FROM "metadata" WHERE "key" = ?
	""".trimIndent())

	private val writeMetadata = connection.prepareStatement("""
		REPLACE INTO "metadata" ("key", "version", "data") VALUES (?, ?, ?)
	""".trimIndent())

	private fun readMetadata(key: String): Pair<Int, ByteArray>? {
		readMetadata.setString(1, key)

		return readMetadata.executeQuery().use {
			if (it.next()) it.getInt(1) to it.getBytes(2) else null
		}
	}

	private fun writeMetadata(key: String, version: Int, value: ByteArray) {
		writeMetadata.setString(1, key)
		writeMetadata.setInt(2, version)
		writeMetadata.setBytes(3, value)
		writeMetadata.execute()
	}

	override fun loadMetadata(): CompletableFuture<Metadata?> {
		return executor.supplyAsync {
			val (version, metadata) = readMetadata("metadata") ?: return@supplyAsync null

			val stream = DataInputStream(BufferedInputStream(InflaterInputStream(FastByteArrayInputStream(metadata))))

			val width = stream.readInt()
			val height = stream.readInt()
			val loopX = stream.readBoolean()
			val loopY = stream.readBoolean()
			val json = VersionedJson("WorldMetadata", version, stream.readJsonElement())

			stream.close()
			Metadata(WorldGeometry(Vector2i(width, height), loopX, loopY), json)
		}
	}

	private val clearEntities = connection.prepareStatement("""
		DELETE FROM "entities" WHERE "chunkX" = ? AND "chunkY" = ?
	""".trimIndent())

	private val beginSaveEntities: PreparedStatement
	private val finishSaveEntities: PreparedStatement
	private val rollbackSaveEntities: PreparedStatement

	init {
		val (begin, commit, rollback) = connection.createSavepoint("save_entities")
		beginSaveEntities = begin
		finishSaveEntities = commit
		rollbackSaveEntities = rollback
	}

	private val writeEntity = connection.prepareStatement("""
		REPLACE INTO "entities" ("chunkX", "chunkY", "x", "y", "unique_id", "type", "version", "data")
			VALUES (?, ?, ?, ?, ?, ?, ?, ?)
	""".trimIndent())

	override fun saveEntities(pos: ChunkPos, entities: Collection<AbstractEntity>) {
		executor.execute {
			beginSaveEntities.execute()

			try {
				clearEntities.setInt(1, pos.x)
				clearEntities.setInt(2, pos.y)
				clearEntities.execute()

				for (entity in entities) {
					Starbound.storeJson {
						val data = JsonObject()
						entity.serialize(data)

						writeEntity.setInt(1, pos.x)
						writeEntity.setInt(2, pos.y)
						writeEntity.setDouble(3, entity.position.x)
						writeEntity.setDouble(4, entity.position.y)
						writeEntity.setString(5, entity.uniqueID.get())
						writeEntity.setString(6, entity.type.storeName)
						writeEntity.setInt(7, VersionRegistry.currentVersion(entity.type.storeName))
						writeEntity.setBytes(8, data.writeJsonObjectDeflated())

						writeEntity.execute()
					}
				}

				finishSaveEntities.execute()
			} catch (err: Throwable) {
				rollbackSaveEntities.execute()
				throw err
			}
		}
	}

	class Palette {
		private val key2index = Object2IntOpenHashMap<String>()
		private val index2key = Int2ObjectAVLTreeMap<String>()

		init {
			key2index.defaultReturnValue(-1)
		}

		fun add(value: String) {
			index2key[key2index.computeIfAbsent(value, Object2IntFunction { key2index.size })] = value
		}

		fun get(value: String): Int {
			val get = key2index.getInt(value)
			if (get == -1) throw NoSuchElementException("No such palette element with key '$value'")
			return get
		}

		fun get(value: Int): String {
			return index2key[value] ?: throw NoSuchElementException("No such palette element with index $value")
		}

		fun writeKey(stream: DataOutputStream, key: String) {
			if (key2index.size <= 256) {
				stream.writeByte(get(key))
			} else {
				stream.writeShort(get(key))
			}
		}

		fun readKey(stream: DataInputStream): String {
			if (key2index.size <= 256) {
				return get(stream.readUnsignedByte())
			} else {
				return get(stream.readUnsignedShort())
			}
		}

		fun write(stream: DataOutputStream) {
			stream.writeCollection(index2key.values) { writeBinaryString(it) }
		}

		fun read(stream: DataInputStream) {
			key2index.clear()
			index2key.clear()

			for (key in stream.readCollection { readBinaryString() }) {
				if (key in key2index)
					throw IllegalStateException("Duplicate palette key $key")

				index2key[key2index.computeIfAbsent(key, Object2IntFunction { key2index.size })] = key
			}
		}
	}

	class PaletteSet {
		val tiles = Palette()
		val mods = Palette()
		val liquids = Palette()

		fun write(stream: DataOutputStream) {
			tiles.write(stream)
			mods.write(stream)
			liquids.write(stream)
		}

		fun read(stream: DataInputStream) {
			tiles.read(stream)
			mods.read(stream)
			liquids.read(stream)
		}
	}

	private val writeCells = connection.prepareStatement("""
		REPLACE INTO "cells" ("x", "y", "generation_stage", "version", "data")
			VALUES (?, ?, ?, ?, ?)
	""".trimIndent())

	override fun saveCells(pos: ChunkPos, data: ChunkCells) {
		executor.execute {
			val palette = PaletteSet()

			for (x in 0 until data.cells.columns) {
				for (y in 0 until data.cells.rows) {
					data.cells[x, y].prepare(palette)
				}
			}

			val deflater = Deflater()
			val buff = FastByteArrayOutputStream()
			val stream = DataOutputStream(BufferedOutputStream(DeflaterOutputStream(buff, deflater, 0x10000), 0x40000))

			try {
				palette.write(stream)

				stream.writeInt(data.cells.columns)
				stream.writeInt(data.cells.rows)

				for (x in 0 until data.cells.columns) {
					for (y in 0 until data.cells.rows) {
						data.cells[x, y].writeNative(stream, palette)
					}
				}

				stream.close()

				writeCells.setInt(1, pos.x)
				writeCells.setInt(2, pos.y)
				writeCells.setInt(3, data.state.ordinal)
				writeCells.setInt(4, 0)
				writeCells.setBytes(5, buff.array.copyOf(buff.length))
				writeCells.execute()
			} finally {
				deflater.end()
			}
		}
	}

	override fun saveMetadata(data: Metadata) {
		executor.execute {
			val buff = FastByteArrayOutputStream()
			val stream = DataOutputStream(BufferedOutputStream(DeflaterOutputStream(buff, Deflater(), 0x10000), 0x40000))

			stream.writeInt(data.geometry.size.x)
			stream.writeInt(data.geometry.size.y)
			stream.writeBoolean(data.geometry.loopX)
			stream.writeBoolean(data.geometry.loopY)
			stream.writeJsonElement(data.data.content)

			stream.close()
			writeMetadata("metadata", data.data.version ?: 0, buff.array.copyOf(buff.length))
		}
	}

	private val findUniqueEntity = connection.prepareStatement("""
		SELECT "chunkX", "chunkY", "x", "y" FROM "entities" WHERE "unique_id" = ?
	""".trimIndent())

	override fun findUniqueEntity(identifier: String): CompletableFuture<UniqueEntitySearchResult?> {
		return executor.supplyAsync {
			findUniqueEntity.setString(1, identifier)

			findUniqueEntity.executeQuery().use {
				if (it.next()) {
					UniqueEntitySearchResult(ChunkPos(it.getInt(1), it.getInt(2)), Vector2d(it.getDouble(3), it.getDouble(4)))
				} else {
					null
				}
			}
		}
	}

	companion object {
		private val LOGGER = LogManager.getLogger()
	}
}
