package ru.dbotthepony.kstarbound.lua

import com.google.gson.JsonArray
import com.google.gson.JsonElement
import com.google.gson.JsonNull
import com.google.gson.JsonObject
import com.google.gson.JsonPrimitive
import com.kenai.jffi.CallContext
import com.kenai.jffi.CallingConvention
import com.kenai.jffi.Closure
import com.kenai.jffi.ClosureManager
import com.kenai.jffi.MemoryIO
import com.kenai.jffi.Type
import jnr.ffi.Memory
import jnr.ffi.NativeType
import org.apache.logging.log4j.LogManager
import org.lwjgl.system.MemoryStack
import org.lwjgl.system.MemoryUtil
import ru.dbotthepony.kstarbound.io.json.InternedJsonElementAdapter
import java.io.Closeable
import java.io.PrintWriter
import java.io.StringWriter
import java.lang.ref.Cleaner
import java.nio.ByteBuffer
import java.nio.ByteOrder
import kotlin.system.exitProcess

private fun stringToBuffer(str: String): ByteBuffer {
	val bytes = str.toByteArray(charset = Charsets.UTF_8)
	val buf = ByteBuffer.allocateDirect(bytes.size)
	buf.order(ByteOrder.nativeOrder())
	bytes.forEach(buf::put)
	buf.position(0)

	return buf
}

@Suppress("unused")
class LuaState : Closeable {
	private val pointer = LuaJNR.INSTANCE.luaL_newstate() ?: throw OutOfMemoryError("Unable to allocate new LuaState")
	private val cleanable: Cleaner.Cleanable
	private val sharedStringBufferPtr = MemoryIO.getInstance().allocateMemory(2L shl 16, false)

	init {
		if (sharedStringBufferPtr == 0L) {
			LuaJNR.INSTANCE.lua_close(pointer)
			throw OutOfMemoryError("Unable to allocate new string shared buffer")
		}

		val pointer = pointer
		val sharedStringBufferPtr = sharedStringBufferPtr

		cleanable = CLEANER.register(this) {
			LuaJNR.INSTANCE.lua_close(pointer)
			MemoryIO.getInstance().freeMemory(sharedStringBufferPtr)
		}
	}

	private val panicHandler = ClosureManager.getInstance().newClosure(
		{
			LOGGER.fatal("${this@LuaState} at $pointer has panicked! This should be impossible!")
			exitProcess(1)
		},

		CallContext.getCallContext(Type.SINT, arrayOf(Type.POINTER), CallingConvention.DEFAULT, false)
	)

	override fun close() {
		cleanable.clean()
	}

	init {
		LuaJNR.INSTANCE.lua_atpanic(pointer, panicHandler.address)

		LuaJNR.INSTANCE.luaopen_base(pointer)
		storeGlobal("_G")
		LuaJNR.INSTANCE.luaopen_package(pointer)
		storeGlobal("package")
		LuaJNR.INSTANCE.luaopen_table(pointer)
		storeGlobal("table")
		LuaJNR.INSTANCE.luaopen_coroutine(pointer)
		storeGlobal("coroutine")
		LuaJNR.INSTANCE.luaopen_string(pointer)
		storeGlobal("string")
		LuaJNR.INSTANCE.luaopen_math(pointer)
		storeGlobal("math")
		LuaJNR.INSTANCE.luaopen_utf8(pointer)
		storeGlobal("utf8")
		LuaJNR.INSTANCE.luaopen_debug(pointer)
		storeGlobal("debug")
	}

	val stackTop: Int get() {
		val value = LuaJNR.INSTANCE.lua_gettop(pointer)
		check(value >= 0) { "Invalid stack top $value" }
		return value
	}

	/**
	 * Converts the acceptable index idx into an equivalent absolute index (that is, one that does not depend on the stack size).
	 */
	fun absStackIndex(index: Int): Int {
		if (index >= 0)
			return index

		return LuaJNR.INSTANCE.lua_absindex(pointer, index)
	}

	fun load(code: String, chunkName: String = "main chunk") {
		val buf = stringToBuffer(code)

		val closure = ClosureManager.getInstance().newClosure(
			object : Closure {
				override fun invoke(buffer: Closure.Buffer) {
					val remainingSize = LuaJNR.RUNTIME.memoryManager.newPointer(buffer.getAddress(2))

					if (buf.remaining() == 0) {
						remainingSize.putLong(0L, 0L)
						buffer.setAddressReturn(0L)
						return
					}

					remainingSize.putLongLong(0L, buf.remaining().toLong())
					val p = MemoryUtil.memAddress(buf)
					buf.position(buf.remaining())
					buffer.setAddressReturn(p)
				}
			},

			CallContext.getCallContext(Type.POINTER, arrayOf(Type.POINTER, Type.ULONG_LONG, Type.POINTER), CallingConvention.DEFAULT, false)
		)

		throwLoadError(LuaJNR.INSTANCE.lua_load(pointer, closure.address, 0L, chunkName, "t"))
		closure.dispose()
	}

	fun call(numArgs: Int = 0, numResults: Int = 0): Int {
		val status = LuaJNR.INSTANCE.lua_pcallk(pointer, numArgs, numResults, 0, 0L, 0L)

		if (status == LUA_ERRRUN) {
			throw LuaRuntimeException(getString())
		}

		return status
	}

	fun getString(stackIndex: Int = -1, limit: Long = 2 shl 16): String? {
		val len = Memory.allocateDirect(LuaJNR.RUNTIME, NativeType.SLONGLONG)
		val p = LuaJNR.INSTANCE.lua_tolstring(pointer, absStackIndex(stackIndex), len) ?: return null

		if (len.getLong(0L) == 0L) {
			return ""
		}

		if (len.getLong(0L) >= limit) {
			throw IllegalStateException("Unreasonably long Lua string: ${len.getLong(0L)}")
		}

		val readBytes = ByteArray(len.getLong(0L).toInt())
		p.get(0L, readBytes, 0, readBytes.size)
		return readBytes.toString(charset = Charsets.UTF_8)
	}

	fun isCFunction(stackIndex: Int = -1): Boolean = LuaJNR.INSTANCE.lua_iscfunction(pointer, absStackIndex(stackIndex)) > 0
	fun isFunction(stackIndex: Int = -1): Boolean = typeAt(stackIndex) == LuaType.FUNCTION
	fun isInteger(stackIndex: Int = -1): Boolean = LuaJNR.INSTANCE.lua_isinteger(pointer, absStackIndex(stackIndex)) > 0
	fun isLightUserdata(stackIndex: Int = -1): Boolean = typeAt(stackIndex) == LuaType.LIGHTUSERDATA
	fun isNil(stackIndex: Int = -1): Boolean = typeAt(stackIndex) == LuaType.NIL
	fun isNone(stackIndex: Int = -1): Boolean = typeAt(stackIndex) == LuaType.NONE
	fun isNoneOrNil(stackIndex: Int = -1): Boolean = typeAt(stackIndex).let { it == LuaType.NIL || it == LuaType.NONE }
	fun isNumber(stackIndex: Int = -1): Boolean = LuaJNR.INSTANCE.lua_isnumber(pointer, absStackIndex(stackIndex)) > 0
	fun isString(stackIndex: Int = -1): Boolean = LuaJNR.INSTANCE.lua_isstring(pointer, absStackIndex(stackIndex)) > 0
	fun isTable(stackIndex: Int = -1): Boolean = typeAt(stackIndex) == LuaType.TABLE
	fun isThread(stackIndex: Int = -1): Boolean = typeAt(stackIndex) == LuaType.THREAD
	fun isUserdata(stackIndex: Int = -1): Boolean = LuaJNR.INSTANCE.lua_isuserdata(pointer, absStackIndex(stackIndex)) > 0
	fun isBoolean(stackIndex: Int = -1): Boolean = typeAt(stackIndex) == LuaType.BOOLEAN

	fun getBoolean(stackIndex: Int = -1): Boolean? {
		if (!isBoolean(stackIndex))
			return null

		return LuaJNR.INSTANCE.lua_toboolean(pointer, stackIndex) > 0
	}

	fun getLong(stackIndex: Int = -1): Long? {
		if (!isInteger(stackIndex))
			return null

		val stack = MemoryStack.stackPush()
		val status = stack.mallocInt(1)
		val value = LuaJNR.INSTANCE.lua_tointegerx(pointer, stackIndex, MemoryUtil.memAddress(status))
		val b = status[0] > 0
		stack.close()

		if (!b)
			return null

		return value
	}

	fun getDouble(stackIndex: Int = -1): Double? {
		if (!isNumber(stackIndex))
			return null

		val stack = MemoryStack.stackPush()
		val status = stack.mallocInt(1)
		val value = LuaJNR.INSTANCE.lua_tonumberx(pointer, stackIndex, MemoryUtil.memAddress(status))
		val b = status[0] > 0
		stack.close()

		if (!b)
			return null

		return value
	}

	fun typeAt(stackIndex: Int = -1): LuaType {
		return when (val value = LuaJNR.INSTANCE.lua_type(pointer, stackIndex)) {
			LUA_TNONE -> LuaType.NONE
			LUA_TNIL -> LuaType.NIL
			LUA_TBOOLEAN -> LuaType.BOOLEAN
			LUA_TLIGHTUSERDATA -> LuaType.LIGHTUSERDATA
			LUA_TNUMBER -> LuaType.NUMBER
			LUA_TSTRING -> LuaType.STRING
			LUA_TTABLE -> LuaType.TABLE
			LUA_TFUNCTION -> LuaType.FUNCTION
			LUA_TUSERDATA -> LuaType.USERDATA
			LUA_TTHREAD -> LuaType.THREAD
			LUA_NUMTYPES -> LuaType.UMTYPES
			else -> throw RuntimeException("Invalid Lua type: $value")
		}
	}

	fun getValue(stackIndex: Int = -1): JsonElement? {
		val abs = absStackIndex(stackIndex)

		if (abs == 0)
			return null

		return when (typeAt(abs)) {
			LuaType.NONE -> null
			LuaType.NIL -> JsonNull.INSTANCE
			LuaType.BOOLEAN -> InternedJsonElementAdapter.of(getBoolean(abs)!!)
			LuaType.LIGHTUSERDATA -> null
			LuaType.NUMBER -> JsonPrimitive(if (isInteger(abs)) getLong(abs)!! else getDouble(abs)!!)
			LuaType.STRING -> JsonPrimitive(getString(abs))
			LuaType.TABLE -> getTable(abs)!!
			LuaType.FUNCTION -> null
			LuaType.USERDATA -> null
			LuaType.THREAD -> null
			LuaType.UMTYPES -> null
		}
	}

	fun getTable(stackIndex: Int = -1): JsonObject? {
		val abs = absStackIndex(stackIndex)

		if (!isTable(abs))
			return null

		val pairs = JsonObject()
		push()

		while (LuaJNR.INSTANCE.lua_next(pointer, abs) != 0) {
			val key = getValue(abs + 1)
			val value = getValue(abs + 2)

			if (key is JsonPrimitive && value != null) {
				pairs.add(key.asString, value)
			}

			pop()
		}

		return pairs
	}

	fun popBoolean(): Boolean? {
		try {
			return getBoolean()
		} finally {
			pop()
		}
	}

	fun popLong(): Long? {
		try {
			return getLong()
		} finally {
			pop()
		}
	}

	fun popDouble(): Double? {
		try {
			return getDouble()
		} finally {
			pop()
		}
	}

	fun popValue(): JsonElement? {
		try {
			return getValue()
		} finally {
			pop()
		}
	}

	fun popTable(): JsonObject? {
		try {
			return getTable()
		} finally {
			pop()
		}
	}

	fun pop(amount: Int = 1): Int {
		if (amount == 0) return 0
		check(amount > 0) { "Invalid amount to pop: $amount" }
		val old = stackTop
		val new = (old - amount).coerceAtLeast(0)
		LuaJNR.INSTANCE.lua_settop(pointer, new)
		return old - new
	}

	fun storeGlobal(name: String) {
		LuaJNR.INSTANCE.lua_setglobal(pointer, name)
	}

	fun loadGlobal(name: String) {
		LuaJNR.INSTANCE.lua_getglobal(pointer, name)
	}

	fun push(closure: (state: LuaState) -> Unit) {
		LuaJNI.lua_pushcclosure(pointer.address()) lazy@{
			try {
				closure.invoke(this@LuaState)
			} catch (err: Throwable) {
				val builder = StringWriter()
				val printWriter = PrintWriter(builder)
				err.printStackTrace(printWriter)
				push(builder.toString())
				return@lazy 1
			}

			return@lazy 0
		}
	}

	fun push() {
		LuaJNR.INSTANCE.lua_pushnil(pointer)
	}

	fun push(value: Int) {
		LuaJNR.INSTANCE.lua_pushinteger(pointer, value.toLong())
	}

	fun push(value: Long) {
		LuaJNR.INSTANCE.lua_pushinteger(pointer, value)
	}

	fun push(value: Double) {
		LuaJNR.INSTANCE.lua_pushnumber(pointer, value)
	}

	fun push(value: Float) {
		LuaJNR.INSTANCE.lua_pushnumber(pointer, value.toDouble())
	}

	fun push(value: Boolean) {
		LuaJNR.INSTANCE.lua_pushboolean(pointer, if (value) 1 else 0)
	}

	fun push(value: String) {
		val bytes = value.toByteArray(Charsets.UTF_8)

		if (bytes.size < 2 shl 16) {
			MemoryIO.getInstance().putByteArray(sharedStringBufferPtr, bytes, 0, bytes.size)
			LuaJNR.INSTANCE.lua_pushlstring(pointer, sharedStringBufferPtr, bytes.size.toLong())
		} else {
			val mem = MemoryIO.getInstance()
			val block = mem.allocateMemory(bytes.size.toLong(), false)

			if (block == 0L)
				throw OutOfMemoryError("Unable to allocate ${bytes.size} bytes on heap")

			try {
				mem.putByteArray(block, bytes, 0, bytes.size)
				LuaJNR.INSTANCE.lua_pushlstring(pointer, block, bytes.size.toLong())
			} finally {
				mem.freeMemory(block)
			}
		}
	}

	fun pushTable(arraySize: Int = 0, hashSize: Int = 0): Int {
		LuaJNR.INSTANCE.lua_createtable(pointer, arraySize, hashSize)
		return stackTop
	}

	fun setTableValue(stackIndex: Int) {
		LuaJNR.INSTANCE.lua_settable(pointer, stackIndex)
	}

	fun push(value: JsonElement) {
		when (value) {
			JsonNull.INSTANCE -> {
				push()
			}

			is JsonPrimitive -> {
				if (value.isNumber) {
					val num = value.asNumber

					when (num) {
						is Int, is Long -> push(num.toLong())
						else -> push(num.toDouble())
					}
				} else if (value.isString) {
					push(value.asString)
				} else if (value.isBoolean) {
					push(value.asBoolean)
				} else {
					throw IllegalArgumentException(value.toString())
				}
			}

			is JsonArray -> {
				val index = pushTable(arraySize = value.size())

				for ((i, v) in value.withIndex()) {
					push(i + 1L)
					push(v)

					setTableValue(index)
				}
			}

			is JsonObject -> {
				val index = pushTable(hashSize = value.size())

				for ((k, v) in value.entrySet()) {
					push(k)
					push(v)

					setTableValue(index)
				}
			}

			else -> {
				throw IllegalArgumentException(value.toString())
			}
		}
	}

	companion object {
		private val LOGGER = LogManager.getLogger()
		private val CLEANER = Cleaner.create {
			val thread = Thread(it, "LuaState cleaner")
			thread.priority = 1
			thread
		}


		const val LUA_TNONE =			-1

		const val LUA_TNIL =			0
		const val LUA_TBOOLEAN =		1
		const val LUA_TLIGHTUSERDATA =	2
		const val LUA_TNUMBER =			3
		const val LUA_TSTRING =			4
		const val LUA_TTABLE =			5
		const val LUA_TFUNCTION =		6
		const val LUA_TUSERDATA =		7
		const val LUA_TTHREAD =			8

		const val LUA_NUMTYPES =		9
	}
}