// © 2026, CGVR MPA Course

#include <imgui.h>
#include <backends/imgui_impl_glfw.h>
#include <backends/imgui_impl_opengl3.h>

#include <glad/glad.h>
#include <GLFW/glfw3.h>

#include <cuda_runtime.h>
#include <cuda_gl_interop.h>

#include <vector>
#include <iostream>
#include <algorithm>

#include "CudaGLImage.h"

#include "JumpFlooding.h"

#define IMAGE_SIZE 1024

static void checkCuda(cudaError_t err){
    if (err != cudaSuccess)
        std::cerr << "CUDA error: " << cudaGetErrorString(err) << std::endl;
}

static CudaGLImage createCudaGLImage(int width = 1024, int height = 1024){
    CudaGLImage img;
    img.width = width;
    img.height = height;

    glGenTextures(1, &img.texture);
    glBindTexture(GL_TEXTURE_2D, img.texture);

    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_LINEAR);
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_LINEAR);
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_S, GL_CLAMP_TO_EDGE);
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_T, GL_CLAMP_TO_EDGE);

    glTexImage2D(GL_TEXTURE_2D, 0, GL_RGBA8, width, height, 0, GL_RGBA, GL_UNSIGNED_BYTE, nullptr);
    glBindTexture(GL_TEXTURE_2D, 0);

    checkCuda(cudaGraphicsGLRegisterImage(&img.cudaResource, img.texture, GL_TEXTURE_2D, cudaGraphicsRegisterFlagsSurfaceLoadStore));

    return img;
}

static void destroyCudaGLImage(CudaGLImage& img){
    if (img.cudaResource)
        cudaGraphicsUnregisterResource(img.cudaResource);

    if (img.texture)
        glDeleteTextures(1, &img.texture);

    img = {};
}

static void clearCudaGLImage(CudaGLImage& img){
    checkCuda(cudaGraphicsMapResources(1, &img.cudaResource));

    cudaArray_t array;
    checkCuda(cudaGraphicsSubResourceGetMappedArray(&array, img.cudaResource, 0, 0));

    std::vector<unsigned char> zero(img.width * img.height * 4, 0);
    checkCuda(cudaMemcpy2DToArray(array, 0, 0, zero.data(), img.width * 4, img.width * 4, img.height, cudaMemcpyHostToDevice));
    checkCuda(cudaGraphicsUnmapResources(1, &img.cudaResource));
}

__global__
void paintLineKernel(cudaSurfaceObject_t surface, int width, int height, int x0, int y0, int x1, int y1, int radius, uchar4 color){
    int x = blockIdx.x * blockDim.x + threadIdx.x;
    int y = blockIdx.y * blockDim.y + threadIdx.y;

    if(x >= width || y >= height){
        return;
    }

    float vx = float(x1 - x0);
    float vy = float(y1 - y0);

    float wx = float(x - x0);
    float wy = float(y - y0);

    float len2 = vx * vx + vy * vy;

    float t = 0.0f;

    if(len2 > 0.0f){
        t = (wx * vx + wy * vy) / len2;
    }

    t = fminf(1.0f, fmaxf(0.0f, t));

    float px = x0 + t * vx;
    float py = y0 + t * vy;

    float dx = x - px;
    float dy = y - py;

    if(dx * dx + dy * dy <= radius * radius){
        surf2Dwrite(color, surface, x * sizeof(uchar4), y);
    }
}

void paintLine(CudaGLImage& img, int x0, int y0, int x1, int y1, int radius, uchar4 color){
    cudaGraphicsMapResources(1, &img.cudaResource);

    cudaArray_t array;
    cudaGraphicsSubResourceGetMappedArray(&array, img.cudaResource, 0, 0);

    cudaResourceDesc desc = {};
    desc.resType = cudaResourceTypeArray;
    desc.res.array.array = array;

    cudaSurfaceObject_t surface = 0;
    cudaCreateSurfaceObject(&surface, &desc);

    dim3 block(16, 16);
    dim3 grid((img.width + block.x - 1) / block.x, (img.height + block.y - 1) / block.y);

    paintLineKernel<<<grid, block>>>(surface, img.width, img.height, x0, y0, x1, y1, radius, color);

    cudaDeviceSynchronize();

    cudaDestroySurfaceObject(surface);
    cudaGraphicsUnmapResources(1, &img.cudaResource);
}

void highlightButton(bool shouldHighlight){
    if(shouldHighlight){
        ImGui::PushStyleColor(ImGuiCol_Button, ImVec4(0.2f, 0.6f, 1.0f, 1.0f));
        ImGui::PushStyleColor(ImGuiCol_ButtonHovered, ImVec4(0.3f, 0.7f, 1.0f, 1.0f));
        ImGui::PushStyleColor(ImGuiCol_ButtonActive, ImVec4(0.1f, 0.5f, 0.9f, 1.0f));
    }
}

void highlightButtonEnd(bool shouldHighlight){
    if(shouldHighlight){
        ImGui::PopStyleColor(3);
    }
}

int main(int, char**){
    glfwSetErrorCallback([](int error, const char* description) {
        std::cerr << "GLFW Error " << error << ": " << description << std::endl;
    });

    if (!glfwInit())
        return 1;

    const char* glsl_version = "#version 330 core";

    glfwWindowHint(GLFW_CONTEXT_VERSION_MAJOR, 3);
    glfwWindowHint(GLFW_CONTEXT_VERSION_MINOR, 3);
    glfwWindowHint(GLFW_OPENGL_PROFILE, GLFW_OPENGL_CORE_PROFILE);
    glfwWindowHint(GLFW_OPENGL_FORWARD_COMPAT, GL_TRUE);

    GLFWwindow* window = glfwCreateWindow(1600, 840, "Jump Flooding Framework", nullptr, nullptr);
    if (!window)
        return 1;

    glfwMakeContextCurrent(window);
    glfwSwapInterval(1);

    if (!gladLoadGLLoader((GLADloadproc)glfwGetProcAddress))
    {
        std::cerr << "Failed to initialize GLAD" << std::endl;
        return -1;
    }

    IMGUI_CHECKVERSION();
    ImGui::CreateContext();

    ImGuiIO& io = ImGui::GetIO();
    io.ConfigFlags |= ImGuiConfigFlags_NavEnableKeyboard;

    float xscale, yscale;
    glfwGetWindowContentScale(window, &xscale, &yscale);

    ImGui::StyleColorsDark();
    ImGui::GetStyle().ScaleAllSizes(xscale);
    io.FontGlobalScale = xscale;

    io.Fonts->AddFontFromFileTTF(CMAKE_SOURCE_DIR "/lib/imgui-1.88/misc/fonts/Roboto-Medium.ttf", 16.0f);

    ImGui_ImplGlfw_InitForOpenGL(window, true);
    ImGui_ImplOpenGL3_Init(glsl_version);

    bool lastPressed = false;
    int lastX = 0, lastY = 0;

    float brushColor[4] = {1.0f, 1.0f, 1.0f, 1.0f};
    int brushRadius = 4;

    CudaGLImage leftImage = createCudaGLImage(IMAGE_SIZE, IMAGE_SIZE);
    CudaGLImage rightImage = createCudaGLImage(IMAGE_SIZE, IMAGE_SIZE);
    JumpFlooding jfa(IMAGE_SIZE, IMAGE_SIZE);

    JFARenderMode renderMode = JFARenderMode::Voronoi;
    bool randomizeBrushColors = false;
    int maxIterations = 10;
    float dilation = 20;

    while (!glfwWindowShouldClose(window))
    {
        glfwPollEvents();

        int display_w, display_h;
        glfwGetFramebufferSize(window, &display_w, &display_h);

        glViewport(0, 0, display_w, display_h);
        glClearColor(0.f, 0.f, 0.f, 1.f);
        glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT);

        ImGui_ImplOpenGL3_NewFrame();
        ImGui_ImplGlfw_NewFrame();
        ImGui::NewFrame();

        ImGui::SetNextWindowPos(ImVec2(0, 0));
        ImGui::SetNextWindowSize(ImVec2((float)display_w, (float)display_h));

        ImGui::Begin("Viewer", nullptr, ImGuiWindowFlags_NoMove | ImGuiWindowFlags_NoResize | ImGuiWindowFlags_NoCollapse | ImGuiWindowFlags_NoTitleBar);

        ImVec2 avail = ImGui::GetContentRegionAvail();

        float topBarHeight = 30.f * xscale;
        float bottomBarHeight = 40.0f * xscale;
        float spacing = ImGui::GetStyle().ItemSpacing.x;
        float imageSize = std::min((avail.x - spacing) * 0.5f, avail.y - bottomBarHeight) - topBarHeight;

        ImGui::BeginChild("Images", ImVec2(avail.x, imageSize + topBarHeight), false);
        ImGui::BeginGroup();
        ImGui::Text("Input / Seeds");
        ImGui::Image((ImTextureID)(intptr_t)leftImage.texture, ImVec2(imageSize, imageSize), ImVec2(0, 1), ImVec2(1, 0));
        ImVec2 imageMin = ImGui::GetItemRectMin();
        ImVec2 imageMax = ImGui::GetItemRectMax();

        bool hovered = ImGui::IsItemHovered();
        bool painting = hovered && ImGui::IsMouseDown(ImGuiMouseButton_Left);

        if (painting)
        {
            ImVec2 mouse = ImGui::GetIO().MousePos;

            float u = (mouse.x - imageMin.x) / (imageMax.x - imageMin.x);
            float v = (mouse.y - imageMin.y) / (imageMax.y - imageMin.y);

            int x = int(u * leftImage.width);
            int y = int(v * leftImage.height);

            y = leftImage.height - 1 - y;

            uchar4 color = make_uchar4(
                brushColor[0] * 255.0f,
                brushColor[1] * 255.0f,
                brushColor[2] * 255.0f,
                brushColor[3] * 255.0f
                );
            paintLine(leftImage, lastPressed ? lastX : x, lastPressed ? lastY : y, x, y, brushRadius, color);

            lastPressed = true;
            lastX = x;
            lastY = y;
        } else {
            if(lastPressed && randomizeBrushColors){
                brushColor[0] = static_cast<float>(std::rand()) / RAND_MAX;
                brushColor[1] = static_cast<float>(std::rand()) / RAND_MAX;
                brushColor[2] = static_cast<float>(std::rand()) / RAND_MAX;
            }

            lastPressed = false;
        }
        ImGui::EndGroup();

        ImGui::SameLine();
        ImGui::BeginGroup();

        JFARenderMode _renderMode = renderMode;

        highlightButton(_renderMode == JFARenderMode::Voronoi);
        if(ImGui::Button("Voronoi")){
            renderMode = JFARenderMode::Voronoi;
        }
        highlightButtonEnd(_renderMode == JFARenderMode::Voronoi);

        ImGui::SameLine();

        highlightButton(_renderMode == JFARenderMode::Distance);
        if(ImGui::Button("Distance")){
            renderMode = JFARenderMode::Distance;
        }
        highlightButtonEnd(_renderMode == JFARenderMode::Distance);

        ImGui::SameLine();

        highlightButton(_renderMode == JFARenderMode::Dilation);
        if(ImGui::Button("Dilation")){
            renderMode = JFARenderMode::Dilation;
        }
        highlightButtonEnd(_renderMode == JFARenderMode::Dilation);

        ImGui::SameLine();

        highlightButton(_renderMode == JFARenderMode::VoronoiBF);
        if(ImGui::Button("Voronoi Brute-Force")){
            renderMode = JFARenderMode::VoronoiBF;
        }
        highlightButtonEnd(_renderMode == JFARenderMode::VoronoiBF);

        ImGui::SameLine();

        highlightButton(_renderMode == JFARenderMode::VoronoiDiff);
        if(ImGui::Button("Voronoi Diff")){
            renderMode = JFARenderMode::VoronoiDiff;
        }
        highlightButtonEnd(_renderMode == JFARenderMode::VoronoiDiff);

        cudaEvent_t start, stop;
        cudaEventCreate(&start);
        cudaEventCreate(&stop);
        cudaEventRecord(start);

        jfa.compute(leftImage, rightImage, renderMode, maxIterations, dilation);

        cudaEventRecord(stop);
        cudaEventSynchronize(stop);
        float jfaMs = 0.0f;
        cudaEventElapsedTime(&jfaMs, start, stop);

        cudaEventDestroy(start);
        cudaEventDestroy(stop);

        ImGui::Image((ImTextureID)(intptr_t)rightImage.texture, ImVec2(imageSize, imageSize), ImVec2(0, 1), ImVec2(1, 0));
        ImGui::EndGroup();

        ImGui::EndChild();
        ImGui::Separator();

        if (ImGui::Button("Clear")) {
            clearCudaGLImage(leftImage);
            clearCudaGLImage(rightImage);
        }

        ImGui::SameLine();
        ImGui::Text("  ");
        ImGui::SameLine();

        ImGui::ColorEdit4("Brush Color", brushColor, ImGuiColorEditFlags_NoInputs);

        ImGui::SameLine();
        ImGui::Text("  ");

        ImGui::SameLine();
        ImGui::SetNextItemWidth(64.0f * xscale);
        ImGui::SliderInt("Rad.", &brushRadius, 1, 64);

        ImGui::SameLine();
        ImGui::Text("  ");
        ImGui::SameLine();
        ImGui::Checkbox("Rand. Colors", &randomizeBrushColors);

        ImGui::SameLine();
        ImGui::Text("  ");
        ImGui::SameLine();
        ImGui::SetNextItemWidth(64.0f * xscale);
        ImGui::SliderInt("Max Itr.", &maxIterations, 1, 12);

        ImGui::SameLine();
        ImGui::Text("   ");
        ImGui::SameLine();
        ImGui::SetNextItemWidth(128.0f * xscale);
        ImGui::SliderFloat("Dilation", &dilation, 1, 512);

        ImGui::SameLine();
        float textWidth = ImGui::CalcTextSize("00.0 ms").x;
        ImGui::SetCursorPosX(ImGui::GetWindowContentRegionMax().x - textWidth);

        ImGui::Text("%.1f ms", jfaMs);

        ImGui::End();

        ImGui::Render();
        ImGui_ImplOpenGL3_RenderDrawData(ImGui::GetDrawData());

        glfwSwapBuffers(window);
    }

    destroyCudaGLImage(leftImage);
    destroyCudaGLImage(rightImage);

    ImGui_ImplOpenGL3_Shutdown();
    ImGui_ImplGlfw_Shutdown();
    ImGui::DestroyContext();

    glfwDestroyWindow(window);
    glfwTerminate();

    return 0;
}
