#include "LuaSaveLoadApi.hpp"
#include <OgreLogManager.h>
#include <string>
#include <unordered_map>

namespace editScene
{

/* ===================================================================== */
/*  Internal state                                                       */
/* ===================================================================== */

static std::unordered_map<std::string, int> s_saveCallbacks;
static std::unordered_map<std::string, int> s_loadCallbacks;

/* ===================================================================== */
/*  Lua table <-> JSON conversion                                        */
/* ===================================================================== */

nlohmann::json luaTableToJson(lua_State *L, int index)
{
	nlohmann::json result;
	int absIdx = lua_absindex(L, index);

	if (lua_istable(L, absIdx)) {
		/* First, try to detect if it's an array */
		bool isArray = true;
		lua_len(L, absIdx);
		int len = (int)lua_tointeger(L, -1);
		lua_pop(L, 1);

		if (len > 0) {
			for (int i = 1; i <= len; i++) {
				lua_rawgeti(L, absIdx, i);
				if (lua_isnil(L, -1)) {
					isArray = false;
					lua_pop(L, 1);
					break;
				}
				lua_pop(L, 1);
			}
		}

		if (isArray && len > 0) {
			result = nlohmann::json::array();
			for (int i = 1; i <= len; i++) {
				lua_rawgeti(L, absIdx, i);
				result.push_back(luaTableToJson(L, -1));
				lua_pop(L, 1);
			}
		} else {
			result = nlohmann::json::object();
			lua_pushnil(L);
			while (lua_next(L, absIdx) != 0) {
				std::string key;
				if (lua_type(L, -2) == LUA_TNUMBER) {
					key = std::to_string(
						(lua_Integer)lua_tonumber(L, -2));
				} else if (lua_type(L, -2) == LUA_TSTRING) {
					key = lua_tostring(L, -2);
				} else {
					lua_pop(L, 1);
					continue;
				}
				result[key] = luaTableToJson(L, -1);
				lua_pop(L, 1);
			}
		}
	} else if (lua_isboolean(L, absIdx)) {
		result = lua_toboolean(L, absIdx) != 0;
	} else if (lua_isinteger(L, absIdx)) {
		result = (int64_t)lua_tointeger(L, absIdx);
	} else if (lua_isnumber(L, absIdx)) {
		result = lua_tonumber(L, absIdx);
	} else if (lua_isstring(L, absIdx)) {
		result = lua_tostring(L, absIdx);
	} else {
		result = nullptr;
	}
	return result;
}

void jsonToLuaValue(lua_State *L, const nlohmann::json &j)
{
	if (j.is_null()) {
		lua_pushnil(L);
	} else if (j.is_boolean()) {
		lua_pushboolean(L, j.get<bool>());
	} else if (j.is_number_integer()) {
		lua_pushinteger(L, (lua_Integer)j.get<int64_t>());
	} else if (j.is_number_float()) {
		lua_pushnumber(L, j.get<double>());
	} else if (j.is_string()) {
		lua_pushstring(L, j.get<std::string>().c_str());
	} else if (j.is_array()) {
		lua_newtable(L);
		int i = 1;
		for (const auto &elem : j) {
			jsonToLuaValue(L, elem);
			lua_rawseti(L, -2, i++);
		}
	} else if (j.is_object()) {
		lua_newtable(L);
		for (auto &[key, val] : j.items()) {
			lua_pushstring(L, key.c_str());
			jsonToLuaValue(L, val);
			lua_rawset(L, -3);
		}
	} else {
		lua_pushnil(L);
	}
}

/* ===================================================================== */
/*  Save / Load data collection                                          */
/* ===================================================================== */

nlohmann::json collectLuaSaveData(lua_State *L)
{
	nlohmann::json data;
	for (const auto &pair : s_saveCallbacks) {
		const std::string &name = pair.first;
		int ref = pair.second;

		lua_rawgeti(L, LUA_REGISTRYINDEX, ref);
		if (lua_pcall(L, 0, 1, 0) == 0) {
			if (lua_istable(L, -1)) {
				data[name] = luaTableToJson(L, -1);
			}
			lua_pop(L, 1);
		} else {
			Ogre::LogManager::getSingleton().logMessage(
				"Lua save callback '" + name +
				"' error: " + lua_tostring(L, -1));
			lua_pop(L, 1);
		}
	}
	return data;
}

void applyLuaLoadData(lua_State *L, const nlohmann::json &data)
{
	for (const auto &pair : s_loadCallbacks) {
		const std::string &name = pair.first;
		int ref = pair.second;

		if (!data.contains(name))
			continue;

		lua_rawgeti(L, LUA_REGISTRYINDEX, ref);
		jsonToLuaValue(L, data[name]);
		if (lua_pcall(L, 1, 0, 0) != 0) {
			Ogre::LogManager::getSingleton().logMessage(
				"Lua load callback '" + name +
				"' error: " + lua_tostring(L, -1));
			lua_pop(L, 1);
		}
	}
}

/* ===================================================================== */
/*  Lua C API functions                                                  */
/* ===================================================================== */

static int luaRegisterSaveCallback(lua_State *L)
{
	const char *name = luaL_checkstring(L, 1);
	luaL_checktype(L, 2, LUA_TFUNCTION);

	/* Unregister existing callback with same name */
	auto it = s_saveCallbacks.find(name);
	if (it != s_saveCallbacks.end()) {
		luaL_unref(L, LUA_REGISTRYINDEX, it->second);
	}

	/* Store new callback in registry */
	lua_pushvalue(L, 2);
	int ref = luaL_ref(L, LUA_REGISTRYINDEX);
	s_saveCallbacks[name] = ref;
	return 0;
}

static int luaRegisterLoadCallback(lua_State *L)
{
	const char *name = luaL_checkstring(L, 1);
	luaL_checktype(L, 2, LUA_TFUNCTION);

	auto it = s_loadCallbacks.find(name);
	if (it != s_loadCallbacks.end()) {
		luaL_unref(L, LUA_REGISTRYINDEX, it->second);
	}

	lua_pushvalue(L, 2);
	int ref = luaL_ref(L, LUA_REGISTRYINDEX);
	s_loadCallbacks[name] = ref;
	return 0;
}

/* ===================================================================== */
/*  Registration                                                         */
/* ===================================================================== */

void registerLuaSaveLoadApi(lua_State *L)
{
	lua_getglobal(L, "ecs");
	if (!lua_istable(L, -1)) {
		lua_pop(L, 1);
		return;
	}

	lua_pushcfunction(L, luaRegisterSaveCallback);
	lua_setfield(L, -2, "register_save_callback");

	lua_pushcfunction(L, luaRegisterLoadCallback);
	lua_setfield(L, -2, "register_load_callback");

	lua_pop(L, 1);
}

} // namespace editScene
