// main.c
#include <stdio.h>
#include <stdlib.h>
#include <stdint.h>

#include "glad/glad.h"
#define GLFW_INCLUDE_NONE
#include <GLFW/glfw3.h>

static void glfw_error_callback(int error, const char* description) {
    fprintf(stderr, "GLFW error %d: %s\n", error, description);
}

static GLuint create_compute_program(const char* source) {
    GLint success = 0;
    GLchar log[1024];

    GLuint shader = glCreateShader(GL_COMPUTE_SHADER);
    glShaderSource(shader, 1, &source, NULL);
    glCompileShader(shader);

    glGetShaderiv(shader, GL_COMPILE_STATUS, &success);
    if (!success) {
        glGetShaderInfoLog(shader, sizeof(log), NULL, log);
        fprintf(stderr, "Compute shader compilation failed:\n%s\n", log);
        exit(EXIT_FAILURE);
    }

    GLuint program = glCreateProgram();
    glAttachShader(program, shader);
    glLinkProgram(program);
    glDeleteShader(shader);

    glGetProgramiv(program, GL_LINK_STATUS, &success);
    if (!success) {
        glGetProgramInfoLog(program, sizeof(log), NULL, log);
        fprintf(stderr, "Program link failed:\n%s\n", log);
        exit(EXIT_FAILURE);
    }

    return program;
}

int main(void) {
    // --- Init GLFW & create GL context ---
    glfwSetErrorCallback(glfw_error_callback);

    if (!glfwInit()) {
        fprintf(stderr, "Failed to init GLFW\n");
        return EXIT_FAILURE;
    }

    // Request OpenGL 4.3 core (needed for compute shaders)
    glfwWindowHint(GLFW_CONTEXT_VERSION_MAJOR, 4);
    glfwWindowHint(GLFW_CONTEXT_VERSION_MINOR, 3);
    glfwWindowHint(GLFW_OPENGL_PROFILE, GLFW_OPENGL_CORE_PROFILE);
    // We don't actually *need* to show the window
    GLFWwindow* window = glfwCreateWindow(640, 480, "GPU Hello", NULL, NULL);
    if (!window) {
        fprintf(stderr, "Failed to create GLFW window\n");
        glfwTerminate();
        return EXIT_FAILURE;
    }

    glfwMakeContextCurrent(window);

    // --- Load GL symbols via GLAD ---
    if (!gladLoadGLLoader((GLADloadproc)glfwGetProcAddress)) {
        fprintf(stderr, "Failed to initialize GLAD\n");
        glfwDestroyWindow(window);
        glfwTerminate();
        return EXIT_FAILURE;
    }

    printf("OpenGL version: %s\n", glGetString(GL_VERSION));
    printf("GLSL version:   %s\n", glGetString(GL_SHADING_LANGUAGE_VERSION));

    // --- Prepare data buffer on GPU (SSBO) ---
    const GLuint N = 16;
    GLuint initial[N];
    for (GLuint i = 0; i < N; ++i) {
        initial[i] = i;
    }

    GLuint ssbo = 0;
    glGenBuffers(1, &ssbo);
    glBindBuffer(GL_SHADER_STORAGE_BUFFER, ssbo);
    glBufferData(GL_SHADER_STORAGE_BUFFER, sizeof(initial), initial, GL_DYNAMIC_COPY);
    glBindBufferBase(GL_SHADER_STORAGE_BUFFER, 0, ssbo); // binding = 0 in shader

    // --- Compute shader source: data[i] *= 2 ---
    const char* compute_src =
        "#version 430 core\n"
        "layout(local_size_x = 16) in;\n"
        "layout(std430, binding = 0) buffer Data {\n"
        "    uint data[];\n"
        "};\n"
        "void main() {\n"
        "    uint idx = gl_GlobalInvocationID.x;\n"
        "    // N == 16 and local_size_x == 16, so we know idx < 16\n"
        "    data[idx] *= 2u;\n"
        "}\n";

    GLuint compute_program = create_compute_program(compute_src);

    // --- Run the compute shader ---
    glUseProgram(compute_program);

    // One workgroup of 16 threads (matches local_size_x)
    glDispatchCompute(1, 1, 1);

    // Make sure writes to the SSBO are visible to the CPU
    glMemoryBarrier(GL_SHADER_STORAGE_BARRIER_BIT);

    // --- Read back the result ---
    glBindBuffer(GL_SHADER_STORAGE_BUFFER, ssbo);
    GLuint* ptr = (GLuint*)glMapBufferRange(
        GL_SHADER_STORAGE_BUFFER,
        0,
        sizeof(initial),
        GL_MAP_READ_BIT
    );

    if (!ptr) {
        fprintf(stderr, "Failed to map SSBO\n");
    } else {
        printf("GPU computed values:\n");
        for (GLuint i = 0; i < N; ++i) {
            printf("  %2u -> %2u\n", i, ptr[i]);
        }
        glUnmapBuffer(GL_SHADER_STORAGE_BUFFER);
    }

    // --- Cleanup ---
    glDeleteProgram(compute_program);
    glDeleteBuffers(1, &ssbo);

    glfwDestroyWindow(window);
    glfwTerminate();

    return EXIT_SUCCESS;
}