#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_Backend.h>
#include <RmlUi_Include_Windows.h>
#include <RmlUi_Include_GL3.h>
#include <GLFW/glfw3.h>
#include <png.h>

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

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();

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

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

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

    load_fonts(assets_path);

    Rml::ElementDocument* document = context->LoadDocument(file.string());
    if (document)
    {
        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(context);
        context->Update();
        Backend::BeginFrame();
        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(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
        context->Update();
        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(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 (document)
                context->UnloadDocument(document);
            context->Update();
            std::this_thread::sleep_for(std::chrono::milliseconds(100));
            document = context->LoadDocument(file.string());
            if (document)
            {
                document->ReloadStyleSheet();
                document->Show();
            }
            FindNextChangeNotification(hNotif);
        }
        running = Backend::ProcessEvents(context);
        context->Update();
        Backend::BeginFrame();
        context->Render();
        Backend::PresentFrame();
    }

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

    return EXIT_SUCCESS;
}