package ru.dbotthepony.kstarbound.server.world

import com.github.benmanes.caffeine.cache.AsyncCacheLoader
import com.github.benmanes.caffeine.cache.Caffeine
import com.google.gson.JsonObject
import it.unimi.dsi.fastutil.io.FastByteArrayInputStream
import it.unimi.dsi.fastutil.io.FastByteArrayOutputStream
import kotlinx.coroutines.CoroutineScope
import kotlinx.coroutines.SupervisorJob
import kotlinx.coroutines.future.await
import kotlinx.coroutines.launch
import org.apache.logging.log4j.LogManager
import ru.dbotthepony.kommons.arrays.Object2DArray
import ru.dbotthepony.kommons.io.ByteKey
import ru.dbotthepony.kommons.io.readBinaryString
import ru.dbotthepony.kommons.io.readCollection
import ru.dbotthepony.kommons.io.readVarInt
import ru.dbotthepony.kommons.io.writeBinaryString
import ru.dbotthepony.kommons.io.writeCollection
import ru.dbotthepony.kommons.io.writeStruct2f
import ru.dbotthepony.kommons.io.writeVarInt
import ru.dbotthepony.kommons.util.xxhash32
import ru.dbotthepony.kstarbound.Starbound
import ru.dbotthepony.kstarbound.VersionRegistry
import ru.dbotthepony.kstarbound.defs.EntityType
import ru.dbotthepony.kstarbound.io.BTreeDB5
import ru.dbotthepony.kstarbound.io.readInternedString
import ru.dbotthepony.kstarbound.io.readVector2f
import ru.dbotthepony.kstarbound.json.VersionedJson
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.ScheduledCoroutineExecutor
import ru.dbotthepony.kstarbound.world.CHUNK_SIZE
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.ImmutableCell
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.ByteArrayInputStream
import java.io.DataInputStream
import java.io.DataOutputStream
import java.io.File
import java.lang.ref.Cleaner
import java.sql.DriverManager
import java.time.Duration
import java.util.concurrent.CompletableFuture
import java.util.concurrent.Executor
import java.util.concurrent.TimeUnit
import java.util.function.Function
import java.util.function.Supplier
import java.util.zip.Deflater
import java.util.zip.DeflaterOutputStream
import java.util.zip.Inflater
import java.util.zip.InflaterInputStream

sealed class LegacyWorldStorage() : WorldStorage() {
	protected abstract fun load(at: ByteKey): CompletableFuture<ByteArray?>
	protected abstract fun write(at: ByteKey, value: ByteArray)
	protected abstract val executor: Executor

	protected val scope by lazy {
		CoroutineScope(ScheduledCoroutineExecutor(executor) + SupervisorJob())
	}

	private val uniqueIndexBlockCache = Caffeine.newBuilder()
		.maximumSize(1024L)
		.scheduler(Starbound)
		.executor(Starbound.EXECUTOR)
		.evictionListener<ByteKey, HashMap<String, Pair<ChunkPos, Vector2d>>> { key, value, cause -> writeUniqueEntityIndex(key!!, value!!) }
		.buildAsync(AsyncCacheLoader<ByteKey, HashMap<String, Pair<ChunkPos, Vector2d>>> { key, executor ->
			load(key).thenApply {
				parseUniqueEntityIndex(it ?: return@thenApply HashMap())
			}
		})

	private fun uniqueEntityIndex(identifier: String): ByteKey {
		val hash = xxhash32(identifier)
		return ByteKey(3, hash[3], hash[2], hash[1], hash[0])
	}

	private fun parseUniqueEntityIndex(data: ByteArray): HashMap<String, Pair<ChunkPos, Vector2d>> {
		val inflater = Inflater()

		try {
			// why is index zlib compressed?
			// I know you guys don't have much of brain, but holy shit, man, why in the world would you
			// compress
			// AN INDEX
			val stream = DataInputStream(BufferedInputStream(InflaterInputStream(FastByteArrayInputStream(data), inflater, 0x10000), 0x40000))

			val keys = stream.readVarInt()

			val map = HashMap<String, Pair<ChunkPos, Vector2d>>(keys)

			for (i in 0 until keys) {
				val key = stream.readInternedString()
				val x = stream.readUnsignedShort()
				val y = stream.readUnsignedShort()
				val pos = stream.readVector2f().toDoubleVector()
				map[key] = ChunkPos(x, y) to pos
			}

			return map
		} finally {
			inflater.end()
		}
	}

	private fun writeUniqueEntityIndex(key: ByteKey, data: HashMap<String, Pair<ChunkPos, Vector2d>>) {
		executor.execute {
			val deflater = Deflater()
			val buffer = FastByteArrayOutputStream()
			// aaaaaaaaaaaaaaaaagrh
			val stream = DataOutputStream(BufferedOutputStream(DeflaterOutputStream(buffer, deflater, 0x10000), 0x40000))

			try {
				stream.writeVarInt(data.size)

				for ((id, value) in data) {
					val (pos, epos) = value

					stream.writeBinaryString(id)
					stream.writeShort(pos.x)
					stream.writeShort(pos.y)
					stream.writeStruct2f(epos.toFloatVector())
				}

				stream.close()
				write(key, buffer.array.copyOf(buffer.length))
			} finally {
				deflater.end()
			}
		}
	}

	override fun findUniqueEntity(identifier: String): CompletableFuture<UniqueEntitySearchResult?> {
		return uniqueIndexBlockCache[uniqueEntityIndex(identifier)].thenApply {
			it[identifier]?.let { UniqueEntitySearchResult(it.first, it.second) }
		}
	}

	override fun loadCells(pos: ChunkPos): CompletableFuture<ChunkCells?> {
		val chunkX = pos.x
		val chunkY = pos.y
		val key = ByteKey(1, (chunkX shr 8).toByte(), chunkX.toByte(), (chunkY shr 8).toByte(), chunkY.toByte())

		return load(key).thenApplyAsync(Function {
			it ?: return@Function null
			val reader = DataInputStream(BufferedInputStream(InflaterInputStream(ByteArrayInputStream(it))))
			val generationLevel = reader.readVarInt()
			val tileSerializationVersion = reader.readVarInt()

			val state = when (generationLevel) {
				0 -> ChunkState.EMPTY
				1 -> ChunkState.TERRAIN
				2 -> ChunkState.MICRO_DUNGEONS
				3 -> ChunkState.CAVE_LIQUID
				4 -> ChunkState.FULL
				else -> ChunkState.FULL
			}

			val result = Object2DArray.nulls<ImmutableCell>(CHUNK_SIZE, CHUNK_SIZE)

			for (y in 0 until CHUNK_SIZE) {
				for (x in 0 until CHUNK_SIZE) {
					val read = MutableCell().readLegacy(reader, tileSerializationVersion)
					result[x, y] = read.immutable()
				}
			}

			reader.close()
			return@Function ChunkCells(result as Object2DArray<out AbstractCell>, state)
		}, Starbound.EXECUTOR)
	}

	override fun loadEntities(pos: ChunkPos): CompletableFuture<Collection<AbstractEntity>> {
		val chunkX = pos.x
		val chunkY = pos.y
		val key = ByteKey(2, (chunkX shr 8).toByte(), chunkX.toByte(), (chunkY shr 8).toByte(), chunkY.toByte())

		return load(key).thenApplyAsync(Function {
			it ?: return@Function listOf()
			val reader = DataInputStream(BufferedInputStream(InflaterInputStream(ByteArrayInputStream(it))))
			val i = reader.readVarInt()
			val objects = ArrayList<AbstractEntity>()

			for (i2 in 0 until i) {
				val obj = VersionedJson(reader)

				val type = EntityType.entries.firstOrNull { it.storeName == obj.id }

				if (type != null) {
					try {
						objects.add(type.fromStorage(obj.content as JsonObject))
					} catch (err: Throwable) {
						LOGGER.error("Unable to deserialize entity in chunk $pos", err)
					}
				} else {
					LOGGER.error("Unknown entity type in chunk $pos: ${obj.id}")
				}
			}

			reader.close()
			return@Function objects
		}, Starbound.EXECUTOR)
	}

	override fun loadMetadata(): CompletableFuture<Metadata?> {
		return load(metadataKey).thenApplyAsync(Function {
			it ?: return@Function null
			val stream = DataInputStream(BufferedInputStream(InflaterInputStream(FastByteArrayInputStream(it))))

			val width = stream.readInt()
			val height = stream.readInt()
			val json = VersionedJson(stream)

			Metadata(WorldGeometry(Vector2i(width, height), true, false), json)
		}, Starbound.EXECUTOR)
	}

	override fun saveEntities(pos: ChunkPos, entities: Collection<AbstractEntity>) {
		scope.launch {
			val chunkX = pos.x
			val chunkY = pos.y
			val key = ByteKey(2, (chunkX shr 8).toByte(), chunkX.toByte(), (chunkY shr 8).toByte(), chunkY.toByte())
			val uniquesKey = ByteKey(4, (chunkX shr 8).toByte(), chunkX.toByte(), (chunkY shr 8).toByte(), chunkY.toByte())

			val buffEntities = FastByteArrayOutputStream()
			val buffUniques = FastByteArrayOutputStream()
			val streamEntities = DataOutputStream(BufferedOutputStream(DeflaterOutputStream(buffEntities)))
			val streamUniques = DataOutputStream(BufferedOutputStream(DeflaterOutputStream(buffUniques)))

			val uniques = HashMap<String, Vector2d>()

			try {
				streamEntities.writeVarInt(entities.size)

				for (entity in entities) {
					Starbound.legacyStoreJson {
						val data = JsonObject()
						entity.serialize(data)
						VersionRegistry.make(entity.type.storeName, data).write(streamEntities)
					}

					val uniqueID = entity.uniqueID.get()

					if (uniqueID != null) {
						uniques[uniqueID] = entity.position
					}
				}

				streamEntities.close()

				val readUniques = load(uniquesKey).await()
				val existingUniques: Set<String>

				if (readUniques == null) {
					existingUniques = setOf()
				} else {
					DataInputStream(BufferedInputStream(InflaterInputStream(FastByteArrayInputStream(readUniques)))).use {
						existingUniques = it.readCollection({ readBinaryString() }, ::HashSet)
					}
				}

				// FIXME: eviction can happen in between these calls... not good, since that might leave
				//  unique index and actual uniques in inconsistent state
				//  Lets try our best to avoid that!

				// TODO: we don't account for unique ID collisions in legacy world storage,
				//  which might corrupt legacy client shipworlds if there is malicious actor / broken mod
				val touchedUniqueKeys = HashMap<ByteKey, HashMap<String, Pair<ChunkPos, Vector2d>>>()

				for (existing in existingUniques) {
					if (existing !in uniques) {
						val indexKey = uniqueEntityIndex(existing)
						val load = touchedUniqueKeys[indexKey] ?: uniqueIndexBlockCache[indexKey].await()
						touchedUniqueKeys[indexKey] = load
						load.remove(existing)
					}
				}

				for ((newKey, newValue) in uniques) {
					if (newKey !in existingUniques) {
						val indexKey = uniqueEntityIndex(newKey)
						val load = touchedUniqueKeys[indexKey] ?: uniqueIndexBlockCache[indexKey].await()
						touchedUniqueKeys[indexKey] = load
						load[newKey] = pos to newValue
					}
				}

				streamUniques.writeCollection(uniques.keys) { writeBinaryString(it) }
				streamUniques.close()

				// non atomic updates, god damn it
				write(key, buffEntities.array.copyOf(buffEntities.length))
				write(uniquesKey, buffUniques.array.copyOf(buffUniques.length))

				for ((index, data) in touchedUniqueKeys) {
					writeUniqueEntityIndex(index, data)
				}
			} catch (err: Throwable) {
				streamEntities.close()
				streamUniques.close()
				throw err
			}
		}
	}

	override fun saveCells(pos: ChunkPos, data: ChunkCells) {
		executor.execute {
			val buff = FastByteArrayOutputStream()
			val stream = DataOutputStream(BufferedOutputStream(DeflaterOutputStream(buff)))

			stream.writeVarInt(
				when (data.state) {
					ChunkState.FRESH -> 0
					ChunkState.EMPTY -> 0
					ChunkState.TERRAIN -> 1
					ChunkState.MICRO_DUNGEONS -> 2
					ChunkState.CAVE_LIQUID -> 3
					ChunkState.FULL -> 4
				}
			)

			stream.writeVarInt(418)

			for (y in 0 until CHUNK_SIZE) {
				for (x in 0 until CHUNK_SIZE) {
					val cell = data.cells.getOrNull(x, y) ?: AbstractCell.NULL
					cell.writeLegacy(stream)
				}
			}

			val chunkX = pos.x
			val chunkY = pos.y
			val key = ByteKey(1, (chunkX shr 8).toByte(), chunkX.toByte(), (chunkY shr 8).toByte(), chunkY.toByte())

			stream.close()
			write(key, buff.array.copyOf(buff.length))
		}
	}

	override fun saveMetadata(data: Metadata) {
		val buff = FastByteArrayOutputStream()
		val stream = DataOutputStream(BufferedOutputStream(DeflaterOutputStream(buff)))

		stream.writeInt(data.geometry.size.x)
		stream.writeInt(data.geometry.size.y)
		data.data.write(stream)

		stream.close()
		write(metadataKey, buff.array.copyOf(buff.length))
	}

	class Memory(private val get: (ByteKey) -> ByteArray?, private val set: (ByteKey, ByteArray) -> Unit) : LegacyWorldStorage() {
		private val pending = HashMap<ByteKey, ByteArray>()
		override val executor: Executor = Executor { it.run() }

		override fun load(at: ByteKey): CompletableFuture<ByteArray?> {
			return CompletableFuture.completedFuture(pending[at] ?: get(at))
		}

		override fun write(at: ByteKey, value: ByteArray) {
			// set(at, value)
			pending[at] = value
		}

		override fun close() {}

		override fun commit() {
			pending.entries.forEach { (k, v) -> set(k, v) }
			pending.clear()
		}
	}

	class DB5(private val database: BTreeDB5) : LegacyWorldStorage() {
		override val executor = CarriedExecutor(Starbound.IO_EXECUTOR)

		override fun load(at: ByteKey): CompletableFuture<ByteArray?> {
			return CompletableFuture.supplyAsync(Supplier { database.read(at).orNull() }, executor)
		}

		override fun write(at: ByteKey, value: ByteArray) {
			throw UnsupportedOperationException()
		}

		override fun close() {
			executor.execute { database.close() }
			executor.wait(300L, TimeUnit.SECONDS)
		}

		override fun commit() {
			// do nothing
		}
	}

	class SQL(path: File) : LegacyWorldStorage() {
		override val executor = CarriedExecutor(Starbound.IO_EXECUTOR)
		private val connection = DriverManager.getConnection("jdbc:sqlite:${path.canonicalPath.replace('\\', '/')}")
		private val cleaner: Cleaner.Cleanable

		init {
			val connection = connection

			cleaner = Starbound.CLEANER.register(this) {
				connection.close()
			}

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

				it.execute("""CREATE TABLE IF NOT EXISTS "data" (
					|"key" BLOB NOT NULL PRIMARY KEY,
					|"value" BLOB NOT NULL
					|)""".trimMargin())
			}

			connection.autoCommit = false
		}

		// concurrent safety - load(), write() and close() are always called on one thread
		// and load()/write() will never be called after close()
		private val loader = connection.prepareStatement("SELECT `value` FROM `data` WHERE `key` = ? LIMIT 1")
		private val writer = connection.prepareStatement("REPLACE INTO `data` (`key`, `value`) VALUES (?, ?)")

		override fun load(at: ByteKey): CompletableFuture<ByteArray?> {
			return CompletableFuture.supplyAsync(Supplier {
				loader.setBytes(1, at.toByteArray())

				loader.executeQuery().use {
					if (it.next()) {
						it.getBytes(1)
					} else {
						null
					}
				}
			}, executor)
		}

		override fun write(at: ByteKey, value: ByteArray) {
			executor.execute {
				writer.setBytes(1, at.toByteArray())
				writer.setBytes(2, value)
				writer.execute()
			}
		}

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

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

	companion object {
		private val LOGGER = LogManager.getLogger()
		private val metadataKey = ByteKey(0, 0, 0, 0, 0)

		fun memory(): Memory {
			val map = HashMap<ByteKey, ByteArray>()
			return Memory(map::get, map::set)
		}
	}
}
