#include "path_sandbox.h"

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

namespace mosis {

//=============================================================================
// CONSTRUCTOR
//=============================================================================

PathSandbox::PathSandbox(const std::string& app_path)
    : m_app_path(app_path)
{
    // Normalize the app path
    if (!m_app_path.empty()) {
        // Ensure trailing separator for prefix matching
        if (m_app_path.back() != '/' && m_app_path.back() != '\\') {
            m_app_path += '/';
        }
        // Normalize separators to forward slash
        std::replace(m_app_path.begin(), m_app_path.end(), '\\', '/');
    }
}

//=============================================================================
// PATH VALIDATION
//=============================================================================

bool PathSandbox::ContainsTraversal(const std::string& path) {
    std::string normalized = NormalizePath(path);

    // Check for .. anywhere in the path
    size_t pos = 0;
    while ((pos = normalized.find("..", pos)) != std::string::npos) {
        // Make sure it's actually a parent directory reference, not part of a filename
        bool at_start = (pos == 0);
        bool before_is_sep = (pos > 0 && (normalized[pos-1] == '/' || normalized[pos-1] == '\\'));

        size_t after_pos = pos + 2;
        bool at_end = (after_pos >= normalized.size());
        bool after_is_sep = (!at_end && (normalized[after_pos] == '/' || normalized[after_pos] == '\\'));

        if ((at_start || before_is_sep) && (at_end || after_is_sep)) {
            return true;
        }
        pos++;
    }

    return false;
}

bool PathSandbox::IsAbsolutePath(const std::string& path) {
    if (path.empty()) return false;

    // Unix absolute path
    if (path[0] == '/') return true;

    // Windows absolute path (C:\ or C:/)
    if (path.length() >= 2) {
        char first = path[0];
        if (std::isalpha(static_cast<unsigned char>(first)) && path[1] == ':') {
            return true;
        }
    }

    // UNC path (\\server\share or //server/share)
    if (path.length() >= 2) {
        if ((path[0] == '\\' && path[1] == '\\') ||
            (path[0] == '/' && path[1] == '/')) {
            return true;
        }
    }

    return false;
}

std::string PathSandbox::NormalizePath(const std::string& path) {
    std::string result = path;

    // Convert backslashes to forward slashes
    std::replace(result.begin(), result.end(), '\\', '/');

    // Remove leading ./
    while (result.length() >= 2 && result[0] == '.' && result[1] == '/') {
        result = result.substr(2);
    }

    // Remove duplicate slashes
    std::string cleaned;
    bool last_was_slash = false;
    for (char c : result) {
        if (c == '/') {
            if (!last_was_slash) {
                cleaned += c;
            }
            last_was_slash = true;
        } else {
            cleaned += c;
            last_was_slash = false;
        }
    }

    return cleaned;
}

bool PathSandbox::ValidatePath(const std::string& path, std::string& out_canonical) {
    // Reject empty paths
    if (path.empty()) {
        return false;
    }

    // Reject absolute paths
    if (IsAbsolutePath(path)) {
        return false;
    }

    // Reject traversal attempts
    if (ContainsTraversal(path)) {
        return false;
    }

    // Normalize and resolve the path
    std::string normalized = NormalizePath(path);
    std::string resolved = ResolvePath(normalized);

    // Use filesystem to get canonical path (resolves any remaining .)
    try {
        std::filesystem::path fs_path(resolved);

        // If the file exists, use canonical path for strict checking
        if (std::filesystem::exists(fs_path)) {
            std::filesystem::path canonical = std::filesystem::canonical(fs_path);
            std::string canonical_str = canonical.string();
            std::replace(canonical_str.begin(), canonical_str.end(), '\\', '/');

            // Verify the canonical path is still within app_path
            std::string app_canonical = std::filesystem::canonical(
                std::filesystem::path(m_app_path)).string();
            std::replace(app_canonical.begin(), app_canonical.end(), '\\', '/');
            if (!app_canonical.empty() && app_canonical.back() != '/') {
                app_canonical += '/';
            }

            if (canonical_str.rfind(app_canonical, 0) != 0) {
                return false;  // Path escaped sandbox via symlink
            }

            out_canonical = canonical_str;
        } else {
            // File doesn't exist, just use the resolved path
            out_canonical = resolved;
        }
    } catch (const std::filesystem::filesystem_error&) {
        // Filesystem error, use the resolved path as-is
        out_canonical = resolved;
    }

    return true;
}

std::string PathSandbox::ResolvePath(const std::string& relative_path) {
    std::string normalized = NormalizePath(relative_path);

    // Combine with app path
    std::string result = m_app_path + normalized;

    return result;
}

//=============================================================================
// MODULE NAME VALIDATION
//=============================================================================

bool PathSandbox::IsValidModuleName(const std::string& name) {
    if (name.empty()) {
        return false;
    }

    // Check each character
    for (size_t i = 0; i < name.length(); i++) {
        char c = name[i];

        // Allow alphanumeric
        if (std::isalnum(static_cast<unsigned char>(c))) {
            continue;
        }

        // Allow underscore
        if (c == '_') {
            continue;
        }

        // Allow dot for submodules, but not at start/end or consecutive
        if (c == '.') {
            if (i == 0 || i == name.length() - 1) {
                return false;  // Dot at start or end
            }
            if (i > 0 && name[i-1] == '.') {
                return false;  // Consecutive dots
            }
            continue;
        }

        // Any other character is invalid
        return false;
    }

    // Reject names that look like traversal
    if (name.find("..") != std::string::npos) {
        return false;
    }

    return true;
}

std::string PathSandbox::ModuleToPath(const std::string& module_name) {
    // Convert dots to path separators
    std::string path = module_name;
    std::replace(path.begin(), path.end(), '.', '/');

    // Add scripts/ prefix and .lua suffix
    return "scripts/" + path + ".lua";
}

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

// Registry key for PathSandbox pointer
static const char* SANDBOX_KEY = "mosis.path_sandbox";

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

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

    // Validate module name
    if (!PathSandbox::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)) {
            // Module already loaded, return cached value
            return 1;
        }
        lua_pop(L, 1);  // Pop nil
    }
    lua_pop(L, 1);  // Pop cache table (or nil if not exists)

    // Get PathSandbox from registry
    lua_getfield(L, LUA_REGISTRYINDEX, SANDBOX_KEY);
    if (!lua_islightuserdata(L, -1)) {
        return luaL_error(L, "require not properly initialized");
    }
    PathSandbox* sandbox = static_cast<PathSandbox*>(lua_touserdata(L, -1));
    lua_pop(L, 1);

    // Convert module name to path
    std::string relative_path = PathSandbox::ModuleToPath(module_name);

    // Validate the path
    std::string canonical;
    if (!sandbox->ValidatePath(relative_path, canonical)) {
        return luaL_error(L, "cannot load module '%s': path validation failed", module_name);
    }

    // Read the file
    std::ifstream file(canonical);
    if (!file.is_open()) {
        // Try with the resolved path directly (in case canonical check failed)
        std::string resolved = sandbox->ResolvePath(relative_path);
        file.open(resolved);
        if (!file.is_open()) {
            return luaL_error(L, "module '%s' not found", module_name);
        }
    }

    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);  // Propagate error
    }

    // 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)) {
        // Create cache table if it doesn't exist
        lua_pop(L, 1);
        lua_newtable(L);
        lua_pushvalue(L, -1);
        lua_setfield(L, LUA_REGISTRYINDEX, LOADED_KEY);
    }

    // cache[module_name] = result
    lua_pushvalue(L, -2);  // Push the result
    lua_setfield(L, -2, module_name);
    lua_pop(L, 1);  // Pop cache table

    // Return the module
    return 1;
}

void RegisterSafeRequire(lua_State* L, PathSandbox* sandbox) {
    // Store PathSandbox pointer in registry
    lua_pushlightuserdata(L, sandbox);
    lua_setfield(L, LUA_REGISTRYINDEX, SANDBOX_KEY);

    // Create loaded modules cache
    lua_newtable(L);
    lua_setfield(L, LUA_REGISTRYINDEX, LOADED_KEY);

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

} // namespace mosis
