#include "onebit_lua_runtime.h"

#include <godot_cpp/classes/file_access.hpp>
#include <godot_cpp/core/class_db.hpp>
#include <godot_cpp/core/object.hpp>
#include <godot_cpp/variant/utility_functions.hpp>
#include <godot_cpp/variant/vector2i.hpp>

extern "C" {
#include <lauxlib.h>
#include <lualib.h>
}

#ifndef LUA_OK
#define LUA_OK 0
#endif

using namespace godot;

void OneBitLuaRuntime::_bind_methods() {
	ClassDB::bind_method(D_METHOD("load_string", "source"), &OneBitLuaRuntime::load_string);
	ClassDB::bind_method(D_METHOD("load_script", "path"), &OneBitLuaRuntime::load_script);
	ClassDB::bind_method(D_METHOD("call_event", "event_name", "ctx", "state", "args"), &OneBitLuaRuntime::call_event, DEFVAL(Array()));
	ClassDB::bind_method(D_METHOD("set_entity_id", "entity_id"), &OneBitLuaRuntime::set_entity_id);
	ClassDB::bind_method(D_METHOD("set_query_host", "host"), &OneBitLuaRuntime::set_query_host);
	ClassDB::bind_method(D_METHOD("set_instruction_budget", "budget"), &OneBitLuaRuntime::set_instruction_budget);
}

OneBitLuaRuntime::OneBitLuaRuntime() {
	_reset_lua();
}

OneBitLuaRuntime::~OneBitLuaRuntime() {
	_reset_lua();
}

void OneBitLuaRuntime::_reset_lua() {
	if (lua) {
		_clear_handlers();
		lua_close(lua);
		lua = nullptr;
	}

	lua = luaL_newstate();
	instruction_budget_exceeded = false;
	_install_api();
}

void OneBitLuaRuntime::_clear_handlers() {
	if (!lua) {
		handlers.clear();
		return;
	}

	for (KeyValue<String, int> &entry : handlers) {
		luaL_unref(lua, LUA_REGISTRYINDEX, entry.value);
	}
	handlers.clear();
}

void OneBitLuaRuntime::_install_api() {
	lua_pushlightuserdata(lua, this);
	lua_setglobal(lua, "__onebit_runtime");

	lua_newtable(lua);
	lua_pushlightuserdata(lua, this);
	lua_pushcclosure(lua, &OneBitLuaRuntime::_lua_sense_is_walkable, 1);
	lua_setfield(lua, -2, "is_walkable");
	lua_pushlightuserdata(lua, this);
	lua_pushcclosure(lua, &OneBitLuaRuntime::_lua_sense_find_nearest, 1);
	lua_setfield(lua, -2, "find_nearest");
	lua_setglobal(lua, "sense");

	lua_newtable(lua);
	lua_pushlightuserdata(lua, this);
	lua_pushcclosure(lua, &OneBitLuaRuntime::_lua_move_step, 1);
	lua_setfield(lua, -2, "step");
	lua_setglobal(lua, "move");

	lua_newtable(lua);
	lua_pushlightuserdata(lua, this);
	lua_pushcclosure(lua, &OneBitLuaRuntime::_lua_log_say, 1);
	lua_setfield(lua, -2, "say");
	lua_setglobal(lua, "log");
}

Dictionary OneBitLuaRuntime::_base_result(bool ok, const String &error) const {
	Dictionary result;
	result["ok"] = ok;
	result["api_version"] = 2;
	result["commands"] = commands;
	result["state"] = Dictionary();
	result["logs"] = logs;
	result["error"] = error;
	return result;
}

Dictionary OneBitLuaRuntime::_error_result(const String &error) const {
	return _base_result(false, error);
}

Dictionary OneBitLuaRuntime::load_string(const String &source) {
	_reset_lua();
	commands.clear();
	logs.clear();

	CharString utf8 = source.utf8();
	if (luaL_loadstring(lua, utf8.get_data()) != LUA_OK) {
		String error = String(lua_tostring(lua, -1));
		lua_pop(lua, 1);
		return _error_result(error);
	}

	if (lua_pcall(lua, 0, 1, 0) != LUA_OK) {
		String error = String(lua_tostring(lua, -1));
		lua_pop(lua, 1);
		return _error_result(error);
	}

	if (!lua_istable(lua, -1)) {
		lua_pop(lua, 1);
		return _error_result("Lua script must return a handler table");
	}

	Dictionary result = _extract_handlers(lua_gettop(lua), "<string>");
	lua_pop(lua, 1);
	return result;
}

Dictionary OneBitLuaRuntime::load_script(const String &path) {
	Ref<FileAccess> file = FileAccess::open(path, FileAccess::READ);
	if (file.is_null()) {
		return _error_result("Could not open Lua script: " + path);
	}

	String source = file->get_as_text();
	return load_string(source);
}

Dictionary OneBitLuaRuntime::_extract_handlers(int table_index, const String &source_name) {
	lua_getfield(lua, table_index, "api_version");
	if (!lua_isnumber(lua, -1) || int(lua_tointeger(lua, -1)) != 2) {
		lua_pop(lua, 1);
		return _error_result("Lua script " + source_name + " must return api_version = 2");
	}
	lua_pop(lua, 1);

	_clear_handlers();
	lua_pushnil(lua);
	while (lua_next(lua, table_index) != 0) {
		if (lua_type(lua, -2) == LUA_TSTRING && lua_isfunction(lua, -1)) {
			String name = String(lua_tostring(lua, -2));
			lua_pushvalue(lua, -1);
			handlers[name] = luaL_ref(lua, LUA_REGISTRYINDEX);
		}
		lua_pop(lua, 1);
	}

	Dictionary result = _base_result(true);
	Array handler_names;
	for (const KeyValue<String, int> &entry : handlers) {
		handler_names.append(entry.key);
	}
	result["handlers"] = handler_names;
	return result;
}

Dictionary OneBitLuaRuntime::call_event(const String &event_name, const Dictionary &ctx, const Dictionary &state, const Array &args) {
	commands.clear();
	logs.clear();
	instruction_budget_exceeded = false;

	if (!handlers.has(event_name)) {
		return _error_result("No Lua handler registered for event: " + event_name);
	}

	lua_rawgeti(lua, LUA_REGISTRYINDEX, handlers[event_name]);
	if (!_push_variant(ctx)) {
		lua_pop(lua, 1);
		return _error_result("Failed to push ctx table");
	}
	if (!_push_variant(state)) {
		lua_pop(lua, 2);
		return _error_result("Failed to push state table");
	}

	lua_pushvalue(lua, -1);
	int state_ref = luaL_ref(lua, LUA_REGISTRYINDEX);

	for (int64_t i = 0; i < args.size(); i++) {
		_push_variant(args[i]);
	}

	lua_sethook(lua, &OneBitLuaRuntime::_instruction_hook, LUA_MASKCOUNT, int(instruction_budget));
	int status = lua_pcall(lua, 2 + int(args.size()), 0, 0);
	lua_sethook(lua, nullptr, 0, 0);

	if (status != LUA_OK) {
		String error = instruction_budget_exceeded ? String("Lua instruction budget exceeded") : String(lua_tostring(lua, -1));
		lua_pop(lua, 1);
		luaL_unref(lua, LUA_REGISTRYINDEX, state_ref);
		return _error_result(error);
	}

	lua_rawgeti(lua, LUA_REGISTRYINDEX, state_ref);
	Variant updated_state = _read_lua_value(-1);
	lua_pop(lua, 1);
	luaL_unref(lua, LUA_REGISTRYINDEX, state_ref);

	Dictionary result = _base_result(true);
	result["commands"] = commands;
	result["logs"] = logs;
	result["state"] = updated_state;
	return result;
}

void OneBitLuaRuntime::set_entity_id(int64_t p_entity_id) {
	entity_id = p_entity_id;
}

void OneBitLuaRuntime::set_query_host(Object *p_host) {
	query_host_id = p_host ? p_host->get_instance_id() : 0;
}

void OneBitLuaRuntime::set_instruction_budget(int64_t p_budget) {
	instruction_budget = p_budget > 0 ? p_budget : 1;
}

void OneBitLuaRuntime::_instruction_hook(lua_State *L, lua_Debug *debug) {
	(void)debug;
	OneBitLuaRuntime *runtime = nullptr;
	lua_getglobal(L, "__onebit_runtime");
	if (lua_islightuserdata(L, -1)) {
		runtime = static_cast<OneBitLuaRuntime *>(lua_touserdata(L, -1));
	}
	lua_pop(L, 1);

	if (runtime) {
		runtime->instruction_budget_exceeded = true;
	}
	luaL_error(L, "Lua instruction budget exceeded");
}

Variant OneBitLuaRuntime::_query(const String &query_name, const Dictionary &params) {
	if (query_host_id == 0) {
		return Variant();
	}

	Object *query_host = ObjectDB::get_instance(query_host_id);
	if (!query_host) {
		return Variant();
	}
	return query_host->call("lua_query", entity_id, query_name, params);
}

void OneBitLuaRuntime::_append_move_step(int dx, int dy) {
	Dictionary params;
	params["direction"] = Vector2i(dx, dy);

	Dictionary command;
	command["type"] = "move";
	command["entity_id"] = entity_id;
	command["params"] = params;
	commands.append(command);
}

void OneBitLuaRuntime::_append_log_say(const String &message, double duration) {
	Dictionary log_entry;
	log_entry["message"] = message;
	log_entry["duration"] = duration;
	logs.append(log_entry);
}

int OneBitLuaRuntime::_lua_move_step(lua_State *L) {
	OneBitLuaRuntime *runtime = static_cast<OneBitLuaRuntime *>(lua_touserdata(L, lua_upvalueindex(1)));
	int dx = int(luaL_checkinteger(L, 1));
	int dy = int(luaL_checkinteger(L, 2));
	runtime->_append_move_step(dx, dy);
	lua_pushboolean(L, 1);
	return 1;
}

int OneBitLuaRuntime::_lua_log_say(lua_State *L) {
	OneBitLuaRuntime *runtime = static_cast<OneBitLuaRuntime *>(lua_touserdata(L, lua_upvalueindex(1)));
	String message = String(luaL_checkstring(L, 1));
	double duration = luaL_optnumber(L, 2, 3.0);
	runtime->_append_log_say(message, duration);
	lua_pushboolean(L, 1);
	return 1;
}

int OneBitLuaRuntime::_lua_sense_is_walkable(lua_State *L) {
	OneBitLuaRuntime *runtime = static_cast<OneBitLuaRuntime *>(lua_touserdata(L, lua_upvalueindex(1)));
	Dictionary params;
	params["x"] = int(luaL_checkinteger(L, 1));
	params["y"] = int(luaL_checkinteger(L, 2));
	Variant result = runtime->_query("sense.is_walkable", params);
	lua_pushboolean(L, result.get_type() == Variant::BOOL && bool(result));
	return 1;
}

int OneBitLuaRuntime::_lua_sense_find_nearest(lua_State *L) {
	OneBitLuaRuntime *runtime = static_cast<OneBitLuaRuntime *>(lua_touserdata(L, lua_upvalueindex(1)));
	Dictionary params;
	params["group"] = String(luaL_checkstring(L, 1));
	Variant result = runtime->_query("sense.find_nearest", params);
	if (!runtime->_push_variant(result)) {
		lua_pushnil(L);
	}
	return 1;
}

bool OneBitLuaRuntime::_push_variant(const Variant &value) {
	switch (value.get_type()) {
		case Variant::NIL:
			lua_pushnil(lua);
			return true;
		case Variant::BOOL:
			lua_pushboolean(lua, bool(value));
			return true;
		case Variant::INT:
			lua_pushinteger(lua, int64_t(value));
			return true;
		case Variant::FLOAT:
			lua_pushnumber(lua, double(value));
			return true;
		case Variant::STRING: {
			CharString utf8 = String(value).utf8();
			lua_pushstring(lua, utf8.get_data());
			return true;
		}
		case Variant::VECTOR2I: {
			Vector2i vector = value;
			lua_newtable(lua);
			lua_pushinteger(lua, vector.x);
			lua_setfield(lua, -2, "x");
			lua_pushinteger(lua, vector.y);
			lua_setfield(lua, -2, "y");
			return true;
		}
		case Variant::DICTIONARY: {
			Dictionary dict = value;
			lua_newtable(lua);
			Array keys = dict.keys();
			for (int64_t i = 0; i < keys.size(); i++) {
				_push_variant(keys[i]);
				_push_variant(dict[keys[i]]);
				lua_settable(lua, -3);
			}
			return true;
		}
		case Variant::ARRAY: {
			Array array = value;
			lua_newtable(lua);
			for (int64_t i = 0; i < array.size(); i++) {
				lua_pushinteger(lua, i + 1);
				_push_variant(array[i]);
				lua_settable(lua, -3);
			}
			return true;
		}
		default:
			lua_pushnil(lua);
			return false;
	}
}

Variant OneBitLuaRuntime::_read_lua_value(int index, int depth) {
	if (depth > 16) {
		return Variant();
	}

	int value_type = lua_type(lua, index);
	switch (value_type) {
		case LUA_TNIL:
			return Variant();
		case LUA_TBOOLEAN:
			return bool(lua_toboolean(lua, index));
		case LUA_TNUMBER:
			return double(lua_tonumber(lua, index));
		case LUA_TSTRING:
			return String(lua_tostring(lua, index));
		case LUA_TTABLE:
			return _read_lua_table(index, depth + 1);
		default:
			return Variant();
	}
}

Dictionary OneBitLuaRuntime::_read_lua_table(int index, int depth) {
	Dictionary dict;
	int abs_index = index;
	if (abs_index < 0) {
		abs_index = lua_gettop(lua) + abs_index + 1;
	}

	lua_pushnil(lua);
	while (lua_next(lua, abs_index) != 0) {
		Variant key = _read_lua_value(-2, depth + 1);
		Variant value = _read_lua_value(-1, depth + 1);
		dict[key] = value;
		lua_pop(lua, 1);
	}

	return dict;
}