#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;

// Custom system interface to enable logging
class LoggingSystemInterface : public Rml::SystemInterface {
public:
    LoggingSystemInterface(Rml::SystemInterface* backend) : backend_(backend) {}

    double GetElapsedTime() override { return backend_->GetElapsedTime(); }

    bool LogMessage(Rml::Log::Type type, const Rml::String& message) override {
        const char* type_str = "";
        switch (type) {
            case Rml::Log::LT_ERROR: type_str = "ERROR"; break;
            case Rml::Log::LT_WARNING: type_str = "WARNING"; break;
            case Rml::Log::LT_INFO: type_str = "INFO"; break;
            case Rml::Log::LT_DEBUG: type_str = "DEBUG"; break;
            default: type_str = "LOG"; break;
        }
        std::println("[RmlUi {}] {}", type_str, message);
        return true;
    }

    // Forward JoinPath to fix Windows path issues
    void JoinPath(Rml::String& translated_path, const Rml::String& document_path, const Rml::String& path) override {
        // Fix paths where colon was converted to pipe (D| -> D:)
        std::string fixed_path = path;
        if (fixed_path.length() >= 2 && fixed_path[1] == '|') {
            fixed_path[1] = ':';
        }

        std::string fixed_doc = document_path;
        if (fixed_doc.length() >= 2 && fixed_doc[1] == '|') {
            fixed_doc[1] = ':';
        }

        // Use std::filesystem to join paths properly and normalize (resolve ..)
        std::filesystem::path doc_dir = std::filesystem::path(fixed_doc).parent_path();
        std::filesystem::path result = (doc_dir / fixed_path).lexically_normal();
        translated_path = result.generic_string();
        std::println("JoinPath: {} + {} -> {}", fixed_doc, fixed_path, translated_path);
    }

private:
    Rml::SystemInterface* backend_;
};

static LoggingSystemInterface* g_logging_interface = nullptr;

// Custom file interface to fix Windows path issues (D| -> D:)
class WindowsFileInterface : public Rml::FileInterface {
public:
    Rml::FileHandle Open(const Rml::String& path) override {
        // Fix paths where colon was converted to pipe (D| -> D:)
        std::string fixed_path = path;
        if (fixed_path.length() >= 2 && fixed_path[1] == '|') {
            fixed_path[1] = ':';
        }
        std::println("FileInterface::Open: {} -> {}", path, fixed_path);

        FILE* fp = fopen(fixed_path.c_str(), "rb");
        return reinterpret_cast<Rml::FileHandle>(fp);
    }

    void Close(Rml::FileHandle file) override {
        fclose(reinterpret_cast<FILE*>(file));
    }

    size_t Read(void* buffer, size_t size, Rml::FileHandle file) override {
        return fread(buffer, 1, size, reinterpret_cast<FILE*>(file));
    }

    bool Seek(Rml::FileHandle file, long offset, int origin) override {
        return fseek(reinterpret_cast<FILE*>(file), offset, origin) == 0;
    }

    size_t Tell(Rml::FileHandle file) override {
        return ftell(reinterpret_cast<FILE*>(file));
    }
};

static WindowsFileInterface* g_file_interface = nullptr;

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)
{
    std::println("lua_loadScreen called!");
    const char* path = luaL_checkstring(L, 1);
    std::println("Loading: {}", path);

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

    if (!std::filesystem::exists(full_path))
    {
        std::println("Screen not found: {}", full_path.generic_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 with absolute path
    std::string full_path_str = full_path.generic_string();
    std::println("Full path: {}", full_path_str);
    g_app.document = g_app.context->LoadDocument(full_path_str);
    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;
    }

    // Use custom interfaces
    g_logging_interface = new LoggingSystemInterface(Backend::GetSystemInterface());
    g_file_interface = new WindowsFileInterface();
    Rml::SetSystemInterface(g_logging_interface);
    Rml::SetFileInterface(g_file_interface);
    Rml::SetRenderInterface(Backend::GetRenderInterface());
    Rml::Initialise();
    Rml::Lua::Initialise();
    std::println("RmlUi and Lua initialized");

    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/ subdirectory
    g_app.assets_path = std::filesystem::absolute(file.parent_path());

    // Walk up the directory tree to find a folder containing a fonts/ subdirectory with .ttf files
    std::filesystem::path fonts_path;
    while (!g_app.assets_path.empty() && g_app.assets_path.has_parent_path())
    {
        fonts_path = g_app.assets_path / "fonts";
        if (std::filesystem::exists(fonts_path) && std::filesystem::is_directory(fonts_path))
        {
            bool has_fonts = false;
            for (const auto& entry : std::filesystem::directory_iterator(fonts_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(fonts_path);

    // Load document with absolute path
    std::filesystem::path abs_file = std::filesystem::absolute(file);
    std::string abs_file_str = abs_file.generic_string();
    std::println("Loading document: {}", abs_file_str);
    std::println("Assets path: {}", g_app.assets_path.generic_string());
    g_app.document = g_app.context->LoadDocument(abs_file_str);
    if (g_app.document)
    {
        g_app.document->Show();
        std::println("Document loaded successfully");
    }
    else
    {
        std::println("Failed to load document!");
    }

    // 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(abs_file_str);
            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;
}