#include "lua_sandbox.h"

#include <lua.hpp>
#include <fstream>
#include <sstream>
#include <iostream>
#include <cstring>
#include <algorithm>
#include <cctype>

namespace mosis {

//=============================================================================
// ALLOCATOR
//=============================================================================

void* LuaSandbox::SandboxAlloc(void* ud, void* ptr, size_t osize, size_t nsize) {
    auto* sandbox = static_cast<LuaSandbox*>(ud);

    // Calculate new usage
    // osize is the old size (0 for new allocations)
    // nsize is the new size (0 for frees)
    size_t new_usage = sandbox->m_memory_used - osize + nsize;

    // Check limit (only when allocating, not freeing)
    if (nsize > 0 && new_usage > sandbox->m_limits.memory_bytes) {
        // Allocation would exceed limit - return nullptr to signal failure
        // Lua will raise a memory error
        return nullptr;
    }

    // Update tracking
    sandbox->m_memory_used = new_usage;

    // Free operation
    if (nsize == 0) {
        free(ptr);
        return nullptr;
    }

    // Alloc or realloc
    return realloc(ptr, nsize);
}

//=============================================================================
// INSTRUCTION HOOK
//=============================================================================

void LuaSandbox::InstructionHook(lua_State* L, lua_Debug* ar) {
    (void)ar;  // Unused

    // Get sandbox pointer from registry
    lua_getfield(L, LUA_REGISTRYINDEX, "__mosis_sandbox");
    auto* sandbox = static_cast<LuaSandbox*>(lua_touserdata(L, -1));
    lua_pop(L, 1);

    if (!sandbox) return;

    // Increment by hook interval (called every 1000 instructions)
    sandbox->m_instructions_used += 1000;

    // Check limit
    if (sandbox->m_instructions_used > sandbox->m_limits.instructions_per_call) {
        luaL_error(L, "instruction limit exceeded");
    }
}

//=============================================================================
// SAFE PRINT
//=============================================================================

int LuaSandbox::SafePrint(lua_State* L) {
    int n = lua_gettop(L);  // number of arguments
    lua_getglobal(L, "tostring");

    for (int i = 1; i <= n; i++) {
        if (i > 1) std::cout << "\t";

        lua_pushvalue(L, -1);  // push tostring
        lua_pushvalue(L, i);   // push argument
        lua_call(L, 1, 1);     // call tostring

        const char* s = lua_tostring(L, -1);
        if (s) {
            std::cout << s;
        }
        lua_pop(L, 1);
    }
    std::cout << std::endl;

    return 0;
}

//=============================================================================
// CONSTRUCTOR / DESTRUCTOR
//=============================================================================

LuaSandbox::LuaSandbox(const SandboxContext& context, const SandboxLimits& limits)
    : m_context(context), m_limits(limits) {

    // Create Lua state with custom allocator
    m_L = lua_newstate(SandboxAlloc, this);
    if (!m_L) {
        m_last_error = "Failed to create Lua state";
        return;
    }

    // Store sandbox pointer in registry for hooks to access
    lua_pushlightuserdata(m_L, this);
    lua_setfield(m_L, LUA_REGISTRYINDEX, "__mosis_sandbox");

    // Setup the sandbox
    SetupSandbox();
}

LuaSandbox::~LuaSandbox() {
    if (m_L) {
        lua_close(m_L);
        m_L = nullptr;
    }
}

//=============================================================================
// SETUP
//=============================================================================

void LuaSandbox::SetupSandbox() {
    // Open safe standard libraries
    luaL_openlibs(m_L);

    // Remove dangerous globals FIRST
    RemoveDangerousGlobals();

    // Setup safe replacements
    SetupSafeGlobals();

    // Protect metatables
    ProtectBuiltinTables();

    // Setup instruction hook for CPU limiting
    SetupInstructionHook();
}

void LuaSandbox::RemoveDangerousGlobals() {
    // List of dangerous globals to remove
    const char* dangerous_globals[] = {
        // Code execution from files/strings
        "dofile",
        "loadfile",
        "load",
        "loadstring",  // Lua 5.1 compat

        // Raw access (bypasses metatables)
        "rawget",
        "rawset",
        "rawequal",
        "rawlen",

        // Metatable manipulation
        // Note: We keep getmetatable but protect the actual metatables
        // setmetatable is removed to prevent modifications
        "setmetatable",

        // GC manipulation
        "collectgarbage",

        // Dangerous libraries
        "os",
        "io",
        "debug",
        "package",

        // LuaJIT / FFI (if present)
        "ffi",
        "jit",
        "newproxy",

        // Module system (we'll add safe version later)
        "require",

        nullptr
    };

    for (const char** p = dangerous_globals; *p; ++p) {
        lua_pushnil(m_L);
        lua_setglobal(m_L, *p);
    }

    // Remove string.dump (can create bytecode from functions)
    lua_getglobal(m_L, "string");
    if (lua_istable(m_L, -1)) {
        lua_pushnil(m_L);
        lua_setfield(m_L, -2, "dump");
    }
    lua_pop(m_L, 1);
}

void LuaSandbox::SetupSafeGlobals() {
    // Replace print with safe version
    lua_pushcfunction(m_L, SafePrint);
    lua_setglobal(m_L, "print");

    // Setup safe require if app_path is set
    if (!m_context.app_path.empty()) {
        SetupSafeRequire();
    }
}

//=============================================================================
// SAFE REQUIRE
//=============================================================================

// Registry key for loaded modules cache
static const char* LOADED_KEY = "mosis.loaded_modules";

// Validate module name for require() - alphanumeric, underscore, dots only
static bool IsValidModuleName(const std::string& name) {
    if (name.empty()) return false;

    for (size_t i = 0; i < name.length(); i++) {
        char c = name[i];
        if (std::isalnum(static_cast<unsigned char>(c))) continue;
        if (c == '_') continue;
        if (c == '.') {
            if (i == 0 || i == name.length() - 1) return false;
            if (i > 0 && name[i-1] == '.') return false;
            continue;
        }
        return false;
    }

    if (name.find("..") != std::string::npos) return false;
    return true;
}

// Convert module name to path (e.g., "ui.button" -> "scripts/ui/button.lua")
static std::string ModuleToPath(const std::string& module_name) {
    std::string path = module_name;
    std::replace(path.begin(), path.end(), '.', '/');
    return "scripts/" + path + ".lua";
}

int LuaSandbox::SafeRequire(lua_State* L) {
    // Get module name
    const char* module_name = luaL_checkstring(L, 1);

    // Validate module name
    if (!IsValidModuleName(module_name)) {
        return luaL_error(L, "invalid module name: %s", module_name);
    }

    // Check cache first
    lua_getfield(L, LUA_REGISTRYINDEX, LOADED_KEY);
    if (lua_istable(L, -1)) {
        lua_getfield(L, -1, module_name);
        if (!lua_isnil(L, -1)) {
            return 1;  // Return cached module
        }
        lua_pop(L, 1);
    }
    lua_pop(L, 1);

    // Get sandbox pointer from registry
    lua_getfield(L, LUA_REGISTRYINDEX, "__mosis_sandbox");
    auto* sandbox = static_cast<LuaSandbox*>(lua_touserdata(L, -1));
    lua_pop(L, 1);

    if (!sandbox) {
        return luaL_error(L, "require not properly initialized");
    }

    // Build full path
    std::string relative_path = ModuleToPath(module_name);
    std::string full_path = sandbox->m_context.app_path;
    if (!full_path.empty() && full_path.back() != '/' && full_path.back() != '\\') {
        full_path += '/';
    }
    full_path += relative_path;

    // Read the file
    std::ifstream file(full_path);
    if (!file.is_open()) {
        return luaL_error(L, "module '%s' not found at '%s'", module_name, full_path.c_str());
    }

    std::stringstream buffer;
    buffer << file.rdbuf();
    std::string source = buffer.str();
    file.close();

    // Load as text only (no bytecode)
    std::string chunk_name = "@" + std::string(module_name);
    int status = luaL_loadbufferx(L, source.c_str(), source.size(),
                                   chunk_name.c_str(), "t");

    if (status != LUA_OK) {
        return lua_error(L);
    }

    // Execute the chunk
    lua_call(L, 0, 1);

    // If chunk returned nil, use true as the module value
    if (lua_isnil(L, -1)) {
        lua_pop(L, 1);
        lua_pushboolean(L, 1);
    }

    // Cache the result
    lua_getfield(L, LUA_REGISTRYINDEX, LOADED_KEY);
    if (!lua_istable(L, -1)) {
        lua_pop(L, 1);
        lua_newtable(L);
        lua_pushvalue(L, -1);
        lua_setfield(L, LUA_REGISTRYINDEX, LOADED_KEY);
    }
    lua_pushvalue(L, -2);
    lua_setfield(L, -2, module_name);
    lua_pop(L, 1);

    return 1;
}

void LuaSandbox::SetupSafeRequire() {
    // Create loaded modules cache
    lua_newtable(m_L);
    lua_setfield(m_L, LUA_REGISTRYINDEX, LOADED_KEY);

    // Register require function
    lua_pushcfunction(m_L, SafeRequire);
    lua_setglobal(m_L, "require");
}

void LuaSandbox::ProtectBuiltinTables() {
    // Protect string metatable
    // When someone calls getmetatable(""), they get "string" instead of the real metatable
    lua_pushstring(m_L, "");
    if (lua_getmetatable(m_L, -1)) {
        lua_pushstring(m_L, "string");
        lua_setfield(m_L, -2, "__metatable");
        lua_pop(m_L, 1);  // pop metatable
    }
    lua_pop(m_L, 1);  // pop string

    // Freeze _G using a proxy pattern
    // This is needed because __newindex only fires for NEW keys, not existing ones
    // We create: empty_proxy -> metatable { __index = real_G, __newindex = error }

    // Get the current _G (with all our safe functions)
    lua_pushglobaltable(m_L);  // stack: real_G

    // Create a new empty table to be the proxy
    lua_newtable(m_L);  // stack: real_G, proxy

    // Create metatable for proxy
    lua_newtable(m_L);  // stack: real_G, proxy, mt

    // __metatable - prevent access to real metatable
    lua_pushstring(m_L, "globals");
    lua_setfield(m_L, -2, "__metatable");

    // __index - read from real_G
    lua_pushvalue(m_L, -3);  // push real_G
    lua_setfield(m_L, -2, "__index");

    // __newindex - block all writes
    lua_pushcfunction(m_L, [](lua_State* L) -> int {
        const char* key = lua_tostring(L, 2);
        return luaL_error(L, "cannot modify global '%s'", key ? key : "(unknown)");
    });
    lua_setfield(m_L, -2, "__newindex");

    // Set metatable on proxy: setmetatable(proxy, mt)
    lua_setmetatable(m_L, -2);  // stack: real_G, proxy

    // Now we need to replace _G with proxy
    // In Lua 5.2+, we use lua_rawseti on the registry
    lua_pushvalue(m_L, -1);  // stack: real_G, proxy, proxy
    lua_rawseti(m_L, LUA_REGISTRYINDEX, LUA_RIDX_GLOBALS);  // stack: real_G, proxy

    // Also update _G variable in real_G to point to proxy
    // This is critical: when code does _G.foo = bar, it accesses _G variable
    lua_setfield(m_L, -2, "_G");  // stack: real_G  (sets real_G["_G"] = proxy)

    lua_pop(m_L, 1);  // pop real_G
}

void LuaSandbox::SetupInstructionHook() {
    // Set hook to fire every 1000 VM instructions
    lua_sethook(m_L, InstructionHook, LUA_MASKCOUNT, 1000);
}

//=============================================================================
// LOAD AND EXECUTE
//=============================================================================

bool LuaSandbox::LoadString(const std::string& code, const std::string& chunk_name) {
    if (!m_L) {
        m_last_error = "Lua state not initialized";
        return false;
    }

    // Reset instruction count for this execution
    ResetInstructionCount();

    // Load as TEXT ONLY - "t" mode rejects bytecode (starts with \x1bLua)
    int result = luaL_loadbufferx(m_L, code.c_str(), code.size(),
                                  chunk_name.c_str(), "t");

    if (result != LUA_OK) {
        m_last_error = lua_tostring(m_L, -1);
        lua_pop(m_L, 1);
        return false;
    }

    // Execute the loaded chunk
    result = lua_pcall(m_L, 0, 0, 0);
    if (result != LUA_OK) {
        m_last_error = lua_tostring(m_L, -1);
        lua_pop(m_L, 1);
        return false;
    }

    m_last_error.clear();
    return true;
}

bool LuaSandbox::LoadFile(const std::string& path) {
    // Read file contents
    std::ifstream f(path);
    if (!f) {
        m_last_error = "Cannot open file: " + path;
        return false;
    }

    std::stringstream ss;
    ss << f.rdbuf();
    std::string code = ss.str();

    // Load as string
    return LoadString(code, "@" + path);
}

void LuaSandbox::ResetInstructionCount() {
    m_instructions_used = 0;
}

} // namespace mosis