#include <iostream>
#include <vector>
#include <string>
#include <fstream>
#include <sstream>
#include <glad/glad.h>
#include <GLFW/glfw3.h>
const char* computeShaderSource = R"(
#version 430 core

layout (local_size_x = 64, local_size_y = 1, local_size_z = 1) in;

layout (std430, binding = 0) buffer Data {
    float numbers[];
};

void main() {
    uint index = gl_GlobalInvocationID.x;
    numbers[index] = numbers[index] * 2.0f;
}
)";

void test_compute()
{

    if (!glfwInit()) {
        std::cerr << "Failed to initialize GLFW" << std::endl;
        return;
    }

    glfwWindowHint(GLFW_CONTEXT_VERSION_MAJOR, 4);
    glfwWindowHint(GLFW_CONTEXT_VERSION_MINOR, 3);
    glfwWindowHint(GLFW_OPENGL_PROFILE, GLFW_OPENGL_CORE_PROFILE);
    glfwWindowHint(GLFW_VISIBLE, GLFW_FALSE); // We don't need a visible window

    GLFWwindow* window = glfwCreateWindow(1, 1, "", NULL, NULL);
    if (!window) {
        std::cerr << "Failed to create GLFW window" << std::endl;
        glfwTerminate();
        return ;
    }
    glfwMakeContextCurrent(window);

    // --- 2. Load OpenGL functions with GLAD ---
    if (!gladLoadGLLoader((GLADloadproc)glfwGetProcAddress)) {
        std::cerr << "Failed to initialize GLAD" << std::endl;
        return;
    }

    // --- 3. Create and compile the compute shader ---
    GLuint computeShader = glCreateShader(GL_COMPUTE_SHADER);
    glShaderSource(computeShader, 1, &computeShaderSource, NULL);
    glCompileShader(computeShader);

    // Check for compilation errors
    int success;
    char infoLog[512];
    glGetShaderiv(computeShader, GL_COMPILE_STATUS, &success);
    if (!success) {
        glGetShaderInfoLog(computeShader, 512, NULL, infoLog);
        std::cerr << "ERROR::SHADER::COMPUTE::COMPILATION_FAILED\n" << infoLog << std::endl;
    }

    GLuint shaderProgram = glCreateProgram();
    glAttachShader(shaderProgram, computeShader);
    glLinkProgram(shaderProgram);

    // Check for linking errors
    glGetProgramiv(shaderProgram, GL_LINK_STATUS, &success);
    if (!success) {
        glGetProgramInfoLog(shaderProgram, 512, NULL, infoLog);
        std::cerr << "ERROR::PROGRAM::LINKING_FAILED\n" << infoLog << std::endl;
    }
    glDeleteShader(computeShader);

    // --- 4. Prepare the data and create the SSBO ---
    const int dataSize = 256;
    std::vector<float> data(dataSize);
    for (int i = 0; i < dataSize; ++i) {
        data[i] = static_cast<float>(i);
    }

    GLuint ssbo;
    glGenBuffers(1, &ssbo);
    glBindBuffer(GL_SHADER_STORAGE_BUFFER, ssbo);
    glBufferData(GL_SHADER_STORAGE_BUFFER, dataSize * sizeof(float), data.data(), GL_DYNAMIC_COPY);
    glBindBufferBase(GL_SHADER_STORAGE_BUFFER, 0, ssbo);

    // --- 5. Dispatch the compute shader ---
    glUseProgram(shaderProgram);
    glDispatchCompute(dataSize / 64, 1, 1);

    // Make sure the computation is finished before reading back the data
    glMemoryBarrier(GL_SHADER_STORAGE_BARRIER_BIT);

    // --- 6. Retrieve the results ---
    glBindBuffer(GL_SHADER_STORAGE_BUFFER, ssbo);
    float* ptr = (float*)glMapBufferRange(GL_SHADER_STORAGE_BUFFER, 0, dataSize * sizeof(float), GL_MAP_READ_BIT);

    std::cout << "Original Data | Processed Data" << std::endl;
    std::cout << "---------------------------------" << std::endl;
    for (int i = 0; i < dataSize; ++i) { // Print first 10 results
        std::cout << i << " | " << ptr[i] << std::endl;
    }

    glUnmapBuffer(GL_SHADER_STORAGE_BUFFER);

    // --- 7. Cleanup ---
    glDeleteProgram(shaderProgram);
    glDeleteBuffers(1, &ssbo);
    glfwDestroyWindow(window);
    glfwTerminate();
}