#include "json_api.h"

#include <lua.hpp>
#include <nlohmann/json.hpp>
#include <unordered_set>
#include <sstream>

using json = nlohmann::json;

namespace mosis {

// Registry key for storing limits
static const char* JSON_LIMITS_KEY = "__mosis_json_limits";

// Get limits from registry
static JsonLimits GetLimits(lua_State* L) {
    lua_getfield(L, LUA_REGISTRYINDEX, JSON_LIMITS_KEY);
    if (lua_islightuserdata(L, -1)) {
        JsonLimits* limits = static_cast<JsonLimits*>(lua_touserdata(L, -1));
        lua_pop(L, 1);
        return *limits;
    }
    lua_pop(L, 1);
    return JsonLimits{};
}

//=============================================================================
// JSON DECODE
//=============================================================================

// Custom exception for JSON errors (thrown instead of luaL_error to allow catching)
class JsonError : public std::runtime_error {
public:
    explicit JsonError(const std::string& msg) : std::runtime_error(msg) {}
};

// Forward declaration
static void JsonToLua(lua_State* L, const json& j, const JsonLimits& limits, int depth);

static void JsonToLua(lua_State* L, const json& j, const JsonLimits& limits, int depth) {
    if (depth > limits.max_depth) {
        throw JsonError("maximum depth exceeded");
    }

    switch (j.type()) {
        case json::value_t::null:
            lua_pushnil(L);
            break;

        case json::value_t::boolean:
            lua_pushboolean(L, j.get<bool>() ? 1 : 0);
            break;

        case json::value_t::number_integer:
        case json::value_t::number_unsigned:
            lua_pushinteger(L, j.get<lua_Integer>());
            break;

        case json::value_t::number_float:
            lua_pushnumber(L, j.get<lua_Number>());
            break;

        case json::value_t::string: {
            const std::string& s = j.get_ref<const std::string&>();
            if (s.size() > limits.max_string_length) {
                throw JsonError("string too large");
            }
            lua_pushlstring(L, s.c_str(), s.size());
            break;
        }

        case json::value_t::array: {
            if (j.size() > limits.max_array_size) {
                throw JsonError("array size limit exceeded");
            }
            lua_createtable(L, static_cast<int>(j.size()), 0);
            int i = 1;
            for (const auto& elem : j) {
                JsonToLua(L, elem, limits, depth + 1);
                lua_rawseti(L, -2, i++);
            }
            break;
        }

        case json::value_t::object: {
            if (j.size() > limits.max_object_size) {
                throw JsonError("object size limit exceeded");
            }
            lua_createtable(L, 0, static_cast<int>(j.size()));
            for (auto it = j.begin(); it != j.end(); ++it) {
                if (it.key().size() > limits.max_string_length) {
                    throw JsonError("key too large");
                }
                lua_pushlstring(L, it.key().c_str(), it.key().size());
                JsonToLua(L, it.value(), limits, depth + 1);
                lua_rawset(L, -3);
            }
            break;
        }

        default:
            lua_pushnil(L);
            break;
    }
}

// json.decode(str) -> table|nil, error
static int lua_json_decode(lua_State* L) {
    size_t len;
    const char* str = luaL_checklstring(L, 1, &len);

    JsonLimits limits = GetLimits(L);

    if (len > limits.max_output_size) {
        lua_pushnil(L);
        lua_pushstring(L, "input too large");
        return 2;
    }

    try {
        json j = json::parse(str, str + len);
        JsonToLua(L, j, limits, 0);
        return 1;
    } catch (const JsonError& e) {
        lua_pushnil(L);
        lua_pushstring(L, e.what());
        return 2;
    } catch (const json::parse_error& e) {
        lua_pushnil(L);
        lua_pushstring(L, e.what());
        return 2;
    } catch (const std::exception& e) {
        lua_pushnil(L);
        lua_pushstring(L, e.what());
        return 2;
    }
}

//=============================================================================
// JSON ENCODE
//=============================================================================

// Forward declaration
static json LuaToJson(lua_State* L, int index, const JsonLimits& limits,
                      int depth, std::unordered_set<const void*>& visited,
                      size_t& output_size);

static json LuaToJson(lua_State* L, int index, const JsonLimits& limits,
                      int depth, std::unordered_set<const void*>& visited,
                      size_t& output_size) {
    if (depth > limits.max_depth) {
        throw JsonError("maximum depth exceeded");
    }

    if (output_size > limits.max_output_size) {
        throw JsonError("output size limit exceeded");
    }

    int type = lua_type(L, index);

    switch (type) {
        case LUA_TNIL:
            return nullptr;

        case LUA_TBOOLEAN:
            return lua_toboolean(L, index) != 0;

        case LUA_TNUMBER:
            if (lua_isinteger(L, index)) {
                return lua_tointeger(L, index);
            }
            return lua_tonumber(L, index);

        case LUA_TSTRING: {
            size_t len;
            const char* s = lua_tolstring(L, index, &len);
            if (len > limits.max_string_length) {
                throw JsonError("string too large");
            }
            output_size += len + 2;  // Approximate: string + quotes
            return std::string(s, len);
        }

        case LUA_TTABLE: {
            // Check for cycles
            const void* ptr = lua_topointer(L, index);
            if (visited.find(ptr) != visited.end()) {
                throw JsonError("circular reference detected");
            }
            visited.insert(ptr);

            // Determine if array or object by checking keys
            bool is_array = true;
            size_t array_len = 0;

            lua_pushnil(L);
            while (lua_next(L, index) != 0) {
                if (lua_type(L, -2) == LUA_TNUMBER && lua_isinteger(L, -2)) {
                    lua_Integer key = lua_tointeger(L, -2);
                    if (key >= 1) {
                        array_len = std::max(array_len, static_cast<size_t>(key));
                    } else {
                        is_array = false;
                    }
                } else {
                    is_array = false;
                }
                lua_pop(L, 1);
            }

            // Verify array is contiguous
            if (is_array && array_len > 0) {
                for (size_t i = 1; i <= array_len; i++) {
                    lua_rawgeti(L, index, static_cast<int>(i));
                    if (lua_isnil(L, -1)) {
                        is_array = false;
                    }
                    lua_pop(L, 1);
                    if (!is_array) break;
                }
            }

            if (is_array && array_len > 0) {
                if (array_len > limits.max_array_size) {
                    throw JsonError("array size limit exceeded");
                }
                json arr = json::array();
                for (size_t i = 1; i <= array_len; i++) {
                    lua_rawgeti(L, index, static_cast<int>(i));
                    arr.push_back(LuaToJson(L, lua_gettop(L), limits, depth + 1, visited, output_size));
                    lua_pop(L, 1);
                }
                visited.erase(ptr);
                return arr;
            } else {
                // Object
                json obj = json::object();
                size_t key_count = 0;

                lua_pushnil(L);
                while (lua_next(L, index) != 0) {
                    key_count++;
                    if (key_count > limits.max_object_size) {
                        throw JsonError("object size limit exceeded");
                    }

                    // Get key as string
                    std::string key;
                    if (lua_type(L, -2) == LUA_TSTRING) {
                        size_t len;
                        const char* s = lua_tolstring(L, -2, &len);
                        key = std::string(s, len);
                    } else if (lua_type(L, -2) == LUA_TNUMBER) {
                        if (lua_isinteger(L, -2)) {
                            key = std::to_string(lua_tointeger(L, -2));
                        } else {
                            key = std::to_string(lua_tonumber(L, -2));
                        }
                    } else {
                        lua_pop(L, 2);
                        throw JsonError("unsupported key type");
                    }

                    output_size += key.size() + 3;  // key + quotes + colon

                    obj[key] = LuaToJson(L, lua_gettop(L), limits, depth + 1, visited, output_size);
                    lua_pop(L, 1);
                }

                visited.erase(ptr);
                return obj;
            }
        }

        case LUA_TFUNCTION:
        case LUA_TUSERDATA:
        case LUA_TTHREAD:
        case LUA_TLIGHTUSERDATA:
            throw JsonError(std::string("unsupported type '") + lua_typename(L, type) + "'");

        default:
            return nullptr;
    }
}

// json.encode(table) -> string|nil, error
static int lua_json_encode(lua_State* L) {
    luaL_checktype(L, 1, LUA_TTABLE);

    JsonLimits limits = GetLimits(L);
    std::unordered_set<const void*> visited;
    size_t output_size = 0;

    try {
        json j = LuaToJson(L, 1, limits, 0, visited, output_size);
        std::string result = j.dump();

        if (result.size() > limits.max_output_size) {
            lua_pushnil(L);
            lua_pushstring(L, "output size limit exceeded");
            return 2;
        }

        lua_pushlstring(L, result.c_str(), result.size());
        return 1;
    } catch (const JsonError& e) {
        lua_pushnil(L);
        lua_pushstring(L, e.what());
        return 2;
    } catch (const std::exception& e) {
        lua_pushnil(L);
        lua_pushstring(L, e.what());
        return 2;
    }
}

//=============================================================================
// REGISTRATION
//=============================================================================

// Helper to set a global in the real _G (bypassing any proxy)
static void SetGlobalInRealG(lua_State* L, const char* name) {
    // Stack: value to set as global

    // Get _G (might be a proxy)
    lua_rawgeti(L, LUA_REGISTRYINDEX, LUA_RIDX_GLOBALS);

    // Check if it has a metatable with __index (proxy pattern)
    if (lua_getmetatable(L, -1)) {
        lua_getfield(L, -1, "__index");
        if (lua_istable(L, -1)) {
            // Found real _G through proxy's __index
            // Stack: value, proxy, mt, real_G
            lua_pushvalue(L, -4);  // Copy value
            lua_setfield(L, -2, name);  // real_G[name] = value
            lua_pop(L, 4);  // pop real_G, mt, proxy, original value
            return;
        }
        lua_pop(L, 2);  // pop __index, metatable
    }

    // No proxy, set directly in _G
    // Stack: value, _G
    lua_pushvalue(L, -2);  // Copy value
    lua_setfield(L, -2, name);  // _G[name] = value
    lua_pop(L, 2);  // pop _G, original value
}

void RegisterJsonAPI(lua_State* L, const JsonLimits& limits) {
    // Store limits in registry (allocate static storage)
    static JsonLimits stored_limits;
    stored_limits = limits;
    lua_pushlightuserdata(L, &stored_limits);
    lua_setfield(L, LUA_REGISTRYINDEX, JSON_LIMITS_KEY);

    // Create json table
    lua_newtable(L);

    lua_pushcfunction(L, lua_json_decode);
    lua_setfield(L, -2, "decode");

    lua_pushcfunction(L, lua_json_encode);
    lua_setfield(L, -2, "encode");

    // Set as global (bypassing proxy)
    SetGlobalInRealG(L, "json");
}

} // namespace mosis
