package ru.dbotthepony.mc.otm.capability.drive

import it.unimi.dsi.fastutil.objects.Object2ObjectOpenHashMap
import net.minecraft.CrashReport
import net.minecraft.ReportedException
import net.minecraft.Util
import net.minecraft.core.HolderLookup
import net.minecraft.core.HolderLookup.Provider
import net.minecraft.nbt.CompoundTag
import net.minecraft.nbt.NbtAccounter
import net.minecraft.nbt.NbtIo
import net.minecraft.world.level.storage.LevelResource
import net.neoforged.neoforge.event.level.LevelEvent
import org.apache.logging.log4j.LogManager
import ru.dbotthepony.mc.otm.MINECRAFT_SERVER
import ru.dbotthepony.mc.otm.Registries
import ru.dbotthepony.mc.otm.SERVER_IS_LIVE
import ru.dbotthepony.mc.otm.isServerThread
import java.lang.ref.WeakReference
import java.nio.file.Files
import java.nio.file.Path
import java.util.*
import java.util.concurrent.ConcurrentLinkedQueue
import kotlin.io.path.exists

private class WeakDriveReference(drive: IMatteryDrive<*>) {
	private var drive: IMatteryDrive<*>? = drive
	private val weak = WeakReference(drive)

	val isValid: Boolean get() {
		return drive != null || weak.get() != null
	}

	fun drive(): IMatteryDrive<*>? {
		return drive ?: weak.get()
	}

	fun sync(registry: HolderLookup.Provider): CompoundTag? {
		val drive = drive() ?: return null

		if (!drive.isDirty) {
			this.drive = null
			return null
		}

		val tag = drive.serializeNBT(registry)
		drive.isDirty = false
		this.drive = null

		return tag
	}

	fun access(): IMatteryDrive<*>? {
		this.drive = weak.get()
		return drive
	}
}

private data class BacklogLine(val uuid: UUID, val tag: CompoundTag)

/**
 * Separate thread and separate files are way faster for enormous volumes of data (if mod is being run on public dedicated server)
 */
object DrivePool {
	private val resource = LevelResource("otm_drives")
	private val LOGGER = LogManager.getLogger()
	private val pool = Object2ObjectOpenHashMap<UUID, WeakDriveReference>()

	private val backlog = ConcurrentLinkedQueue<BacklogLine>()
	private var knownBaseDirectory: Path? = null

	private val baseDirectory: Path get() {
		val baseDirectory = MINECRAFT_SERVER.storageSource.getLevelPath(resource)

		if (knownBaseDirectory != baseDirectory) {
			Files.createDirectories(baseDirectory)
			knownBaseDirectory = baseDirectory
		}

		return baseDirectory
	}

	fun <T : IMatteryDrive<*>> get(id: UUID, deserializer: (CompoundTag, HolderLookup.Provider) -> T, factory: () -> T): T {
		if (!isServerThread())
			throw IllegalAccessException("Can not access drive pool from outside of server thread.")

		if (!SERVER_IS_LIVE)
			throw IllegalStateException("Server is not running")

		val get = pool[id]?.access()

		if (get != null) {
			return get as T
		}

		val file = baseDirectory.resolve("$id.dat")

		if (file.exists()) {
			try {
				val factorize = deserializer(NbtIo.readCompressed(file, NbtAccounter.unlimitedHeap()), MINECRAFT_SERVER.registryAccess())
				pool[id] = WeakDriveReference(factorize)
				return factorize
			} catch (error: Throwable) {
				val report = CrashReport("Reading OTM Drive from disk", error)
				val category = report.addCategory("Corrupt drive")
				category.setDetail("UUID", id)
				category.setDetail("File location", file)
				throw ReportedException(report)
			}
		}

		return factory().also { pool[id] = WeakDriveReference(it) }
	}

	fun markDirty(id: UUID) {
		if (isServerThread()) {
			pool[id]?.access()
		}
	}

	fun onWorldSave(event: LevelEvent.Save) {
		writeBacklog(event.level.server!!.registryAccess())
	}

	private fun writeBacklog(registry: Provider) {
		var needsSync = false
		val removeKeys = ArrayList<UUID>()

		for ((key, value) in pool) {
			try {
				val tag = value.sync(registry)

				if (tag != null) {
					LOGGER.debug("Serializing OTM Drive {}", key)
					backlog.add(BacklogLine(key, tag))
					needsSync = true
				} else if (!value.isValid) {
					removeKeys.add(key)
				}
			} catch (err: Throwable) {
				LOGGER.error("Failed to serialize OTM Drive for persistent storage with ID $key", err)
			}
		}

		for (key in removeKeys) {
			pool.remove(key)
		}

		if (needsSync) {
			Util.backgroundExecutor().execute { sync(registry) }
		}
	}

	private fun sync(registry: HolderLookup.Provider) {
		while (true) {
			val next = backlog.poll() ?: return
			val (uuid, tag) = next

			try {
				LOGGER.debug("Syncing OTM Drive {}", uuid)

				val oldFile = baseDirectory.resolve("$uuid.dat_old")
				val newFile = baseDirectory.resolve("$uuid.dat")

				if (oldFile.exists()) Files.delete(oldFile)
				if (newFile.exists()) Files.move(newFile, oldFile)

				NbtIo.writeCompressed(tag, newFile)
			} catch (error: Throwable) {
				LOGGER.error("Failed to sync OTM Drive to persistent storage with ID $uuid", error)
			}
		}
	}
}
