MosisDesigner/main.cpp

#include <print>
#include <thread>
#include <chrono>
#include <filesystem>
#include <iostream>
#include <fstream>
#include <vector>
#include <cstring>
#include <RmlUi/Core.h>
#include <RmlUi/Debugger.h>
#include <RmlUi/Lua.h>
#include <RmlUi/Lua/Interpreter.h>
#include <RmlUi_Backend.h>
#include <RmlUi_Include_Windows.h>
#include <RmlUi_Include_GL3.h>
#include <GLFW/glfw3.h>
#include <png.h>

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

// Global app state for Lua access
struct AppState {
    Rml::Context* context = nullptr;
    Rml::ElementDocument* document = nullptr;
    std::filesystem::path assets_path;
} g_app;

void load_fonts(const std::filesystem::path& dir)
{
    for (const auto& file : std::filesystem::directory_iterator(dir))
    {
        if (file.path().extension() == ".ttf")
        {
            Rml::LoadFontFace(file.path().string());
        }
    }
}

// Lua function: loadScreen(path) - loads a new RML document
// Path is relative to assets folder
int lua_loadScreen(lua_State* L)
{
    const char* path = luaL_checkstring(L, 1);

    std::filesystem::path full_path = g_app.assets_path / path;

    if (!std::filesystem::exists(full_path))
    {
        std::println("Screen not found: {}", full_path.string());
        lua_pushboolean(L, false);
        return 1;
    }

    // Unload current document
    if (g_app.document)
    {
        g_app.context->UnloadDocument(g_app.document);
        g_app.document = nullptr;
    }

    // Load new document
    g_app.document = g_app.context->LoadDocument(full_path.string());
    if (g_app.document)
    {
        g_app.document->Show();
        std::println("Loaded screen: {}", path);
        lua_pushboolean(L, true);
    }
    else
    {
        std::println("Failed to load screen: {}", path);
        lua_pushboolean(L, false);
    }

    return 1;
}

// Register custom Lua functions
void registerLuaFunctions()
{
    lua_State* L = Rml::Lua::Interpreter::GetLuaState();

    // Register global function
    lua_pushcfunction(L, lua_loadScreen);
    lua_setglobal(L, "loadScreen");

    std::println("Registered Lua functions");
}

void DumpElementTree(Rml::Element* element, std::ofstream& out, int depth = 0)
{
    if (!element) return;

    std::string indent(depth * 2, ' ');
    std::string id = element->GetId().empty() ? "" : "#" + element->GetId();

    out << indent << element->GetTagName() << id
        << "  @ (" << element->GetAbsoluteLeft() << ", " << element->GetAbsoluteTop() << ")  "
        << element->GetOffsetWidth() << "x" << element->GetOffsetHeight() << "\n";

    for (int i = 0; i < element->GetNumChildren(); i++)
        DumpElementTree(element->GetChild(i), out, depth + 1);
}

// Offscreen framebuffer for screenshot capture
struct OffscreenFBO {
    GLuint fbo = 0;
    GLuint texture = 0;
    GLuint depth_rbo = 0;
    int width = 0;
    int height = 0;

    bool create(int w, int h)
    {
        width = w;
        height = h;

        // Create framebuffer
        glGenFramebuffers(1, &fbo);
        glBindFramebuffer(GL_FRAMEBUFFER, fbo);

        // Create color texture
        glGenTextures(1, &texture);
        glBindTexture(GL_TEXTURE_2D, texture);
        glTexImage2D(GL_TEXTURE_2D, 0, GL_RGBA8, width, height, 0, GL_RGBA, GL_UNSIGNED_BYTE, nullptr);
        glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_LINEAR);
        glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_LINEAR);
        glFramebufferTexture2D(GL_FRAMEBUFFER, GL_COLOR_ATTACHMENT0, GL_TEXTURE_2D, texture, 0);

        // Create depth/stencil renderbuffer
        glGenRenderbuffers(1, &depth_rbo);
        glBindRenderbuffer(GL_RENDERBUFFER, depth_rbo);
        glRenderbufferStorage(GL_RENDERBUFFER, GL_DEPTH24_STENCIL8, width, height);
        glFramebufferRenderbuffer(GL_FRAMEBUFFER, GL_DEPTH_STENCIL_ATTACHMENT, GL_RENDERBUFFER, depth_rbo);

        // Check completeness
        if (glCheckFramebufferStatus(GL_FRAMEBUFFER) != GL_FRAMEBUFFER_COMPLETE)
        {
            std::println("Failed to create offscreen framebuffer");
            destroy();
            return false;
        }

        glBindFramebuffer(GL_FRAMEBUFFER, 0);
        return true;
    }

    void destroy()
    {
        if (texture) glDeleteTextures(1, &texture);
        if (depth_rbo) glDeleteRenderbuffers(1, &depth_rbo);
        if (fbo) glDeleteFramebuffers(1, &fbo);
        fbo = texture = depth_rbo = 0;
    }

    void bind()
    {
        glBindFramebuffer(GL_FRAMEBUFFER, fbo);
        glViewport(0, 0, width, height);
    }

    std::vector<unsigned char> readPixels()
    {
        std::vector<unsigned char> pixels(width * height * 3);
        glBindFramebuffer(GL_FRAMEBUFFER, fbo);
        glReadPixels(0, 0, width, height, GL_RGB, GL_UNSIGNED_BYTE, pixels.data());
        return pixels;
    }
};

void SavePNG(const std::filesystem::path& path, const std::vector<unsigned char>& pixels, int width, int height)
{
    // Flip vertically (OpenGL origin is bottom-left)
    std::vector<unsigned char> flipped(width * height * 3);
    for (int y = 0; y < height; y++)
    {
        std::memcpy(&flipped[y * width * 3],
                    &pixels[(height - 1 - y) * width * 3],
                    width * 3);
    }

    FILE* fp = fopen(path.string().c_str(), "wb");
    if (!fp)
    {
        std::println("Failed to open {} for writing", path.string());
        return;
    }

    png_structp png = png_create_write_struct(PNG_LIBPNG_VER_STRING, nullptr, nullptr, nullptr);
    png_infop info = png_create_info_struct(png);
    png_init_io(png, fp);
    png_set_IHDR(png, info, width, height, 8, PNG_COLOR_TYPE_RGB,
                 PNG_INTERLACE_NONE, PNG_COMPRESSION_TYPE_DEFAULT, PNG_FILTER_TYPE_DEFAULT);
    png_write_info(png, info);

    std::vector<png_bytep> rows(height);
    for (int y = 0; y < height; y++)
        rows[y] = const_cast<png_bytep>(&flipped[y * width * 3]);
    png_write_image(png, rows.data());
    png_write_end(png, nullptr);

    png_destroy_write_struct(&png, &info);
    fclose(fp);

    std::println("Screenshot saved to: {}", path.string());
}

int main(const int argc, const char* argv[])
{
    constexpr int window_width = 540;
    constexpr int window_height = 960;

    if (argc < 2)
    {
        std::println("Usage: mosis-designer file.rml [--dump]");
        return EXIT_FAILURE;
    }

    const std::filesystem::path file = argv[1];
    const bool dump_mode = (argc >= 3 && std::string(argv[2]) == "--dump");
    if (!std::filesystem::exists(file))
    {
        std::println("File does not exist");
        return EXIT_FAILURE;
    }

    if (!Backend::Initialize("Load Document Sample", window_width, window_height, true))
    {
        return EXIT_FAILURE;
    }

    Rml::SetSystemInterface(Backend::GetSystemInterface());
    Rml::SetRenderInterface(Backend::GetRenderInterface());
    Rml::Initialise();
    Rml::Lua::Initialise();

    g_app.context = Rml::CreateContext("main", Rml::Vector2i(window_width, window_height));
    if (!g_app.context)
    {
        Rml::Shutdown();
        Backend::Shutdown();
        return EXIT_FAILURE;
    }

    // Register custom Lua functions
    registerLuaFunctions();

    // Find the assets folder by checking for fonts
    g_app.assets_path = std::filesystem::absolute(file.parent_path());

    // Walk up the directory tree to find a folder containing .ttf fonts
    while (!g_app.assets_path.empty() && g_app.assets_path.has_parent_path())
    {
        bool has_fonts = false;
        for (const auto& entry : std::filesystem::directory_iterator(g_app.assets_path))
        {
            if (entry.path().extension() == ".ttf")
            {
                has_fonts = true;
                break;
            }
        }
        if (has_fonts) break;
        g_app.assets_path = g_app.assets_path.parent_path();
    }

    load_fonts(g_app.assets_path);

    g_app.document = g_app.context->LoadDocument(file.string());
    if (g_app.document)
    {
        g_app.document->Show();
    }

    // Dump mode: render and capture screenshot
    if (dump_mode)
    {
        // Create offscreen FBO at window dimensions
        OffscreenFBO fbo;
        if (!fbo.create(window_width, window_height))
        {
            std::println("Failed to create FBO");
            Rml::Shutdown();
            Backend::Shutdown();
            return EXIT_FAILURE;
        }

        // First render to default framebuffer to initialize everything
        Backend::ProcessEvents(g_app.context);
        g_app.context->Update();
        Backend::BeginFrame();
        g_app.context->Render();
        Backend::PresentFrame();

        // Create dump folder if it doesn't exist
        auto dump_folder = std::filesystem::absolute("dump");
        std::filesystem::create_directories(dump_folder);

        // Write element dump
        std::ofstream dump_file((dump_folder / "element_dump.txt").string());
        if (dump_file)
        {
            dump_file << "Window: " << window_width << "x" << window_height << "\n";
            dump_file << "\n=== Element Layout ===\n";
            DumpElementTree(g_app.document, dump_file);
            dump_file.close();
        }

        // Bind FBO and set up OpenGL state for rendering
        fbo.bind();
        glClearColor(0.0f, 0.0f, 0.0f, 1.0f);
        glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT | GL_STENCIL_BUFFER_BIT);

        // Enable blending for proper text rendering
        glEnable(GL_BLEND);
        glBlendFunc(GL_SRC_ALPHA, GL_ONE_MINUS_SRC_ALPHA);

        // Render to FBO
        g_app.context->Update();
        g_app.context->Render();

        // Read pixels
        auto pixels = fbo.readPixels();
        fbo.destroy();

        SavePNG(dump_folder / "screenshot.png", pixels, window_width, window_height);

        // Restore default framebuffer
        glBindFramebuffer(GL_FRAMEBUFFER, 0);

        Rml::Shutdown();
        Backend::Shutdown();
        return EXIT_SUCCESS;
    }

    const HANDLE hNotif = FindFirstChangeNotification(g_app.assets_path.c_str(),
        TRUE, FILE_NOTIFY_CHANGE_LAST_WRITE);

    bool running = true;
    while (running)
    {
        if (const DWORD wait_result = WaitForSingleObject(hNotif, 100);
            wait_result == WAIT_OBJECT_0)
        {
            if (g_app.document)
                g_app.context->UnloadDocument(g_app.document);
            g_app.context->Update();
            std::this_thread::sleep_for(std::chrono::milliseconds(100));
            g_app.document = g_app.context->LoadDocument(file.string());
            if (g_app.document)
            {
                g_app.document->ReloadStyleSheet();
                g_app.document->Show();
            }
            FindNextChangeNotification(hNotif);
        }
        running = Backend::ProcessEvents(g_app.context);
        g_app.context->Update();
        Backend::BeginFrame();
        g_app.context->Render();
        Backend::PresentFrame();
    }

    Rml::Shutdown();
    Backend::Shutdown();

    return EXIT_SUCCESS;
}