#include <iostream>
#include <vector>
#include <string>
#include <fstream>
#include <sstream>
#include <glad/glad.h>
#include <GLFW/glfw3.h>
#define STB_IMAGE_IMPLEMENTATION
#include "stb_image.h"
#define STB_IMAGE_WRITE_IMPLEMENTATION
#include "stb_image_write.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_computeshader_simple()
{

    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);
    std::cout << "x,y,z need less than : " << GL_MAX_COMPUTE_WORK_GROUP_SIZE << std::endl;
    // 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();
}

const char* invertVtxShaderSource = R"(
#version 330 core
layout (location = 0) in vec2 aPos;
layout (location = 1) in vec2 aTexCoords;

out vec2 TexCoords;

void main() {
    gl_Position = vec4(aPos, 0.0, 1.0);
    TexCoords = aTexCoords;
})";

const char* invertFrgShaderSource = R"(
#version 330 core
out vec4 FragColor;

in vec2 TexCoords;

uniform sampler2D screenTexture;

void main() {
    FragColor = texture(screenTexture, TexCoords);
})";

const char* invertCptShaderSources = R"(
#version 430 core

// Define the size of the work group. 8x8 is a common size.
layout (local_size_x = 8, local_size_y = 8, local_size_z = 1) in;

// Input texture (the original image)
layout (binding = 0, rgba8) uniform readonly image2D input_image;
// Output texture (where we'll store the inverted image)
layout (binding = 1, rgba8) uniform writeonly image2D output_image;

void main() {
    // Get the pixel coordinates for this shader invocation
    ivec2 pixel_coords = ivec2(gl_GlobalInvocationID.xy);

    // Load the color from the input image at the current coordinates
    vec4 color = imageLoad(input_image, pixel_coords);

    // Invert the color (1.0 - color for each channel)
    // The alpha channel is left unchanged.
    vec4 inverted_color = vec4(1.0 - color.r, 1.0 - color.g, 1.0 - color.b, color.a);

    // Store the inverted color in the output image
    imageStore(output_image, pixel_coords, inverted_color);
})";

GLuint create_shader_program(const char* vert_source, const char* frag_source, const char* comp_source = nullptr) {
    GLuint program = glCreateProgram();

    if (vert_source && frag_source) {
        GLuint vertexShader = glCreateShader(GL_VERTEX_SHADER);
        glShaderSource(vertexShader, 1, &vert_source, NULL);
        glCompileShader(vertexShader);
        // Add error checking...

        GLuint fragmentShader = glCreateShader(GL_FRAGMENT_SHADER);
        glShaderSource(fragmentShader, 1, &frag_source, NULL);
        glCompileShader(fragmentShader);
        // Add error checking...

        glAttachShader(program, vertexShader);
        glAttachShader(program, fragmentShader);
        glDeleteShader(vertexShader);
        glDeleteShader(fragmentShader);
    }

    if (comp_source) {
        GLuint computeShader = glCreateShader(GL_COMPUTE_SHADER);
        glShaderSource(computeShader, 1, &comp_source, NULL);
        glCompileShader(computeShader);
        // Add error checking for compute shader...
        glAttachShader(program, computeShader);
        glDeleteShader(computeShader);
    }

    glLinkProgram(program);
    // Add error checking for linking...
    return program;
}

void test_invert_image() {
    // --- GLFW and GLAD initialization ---
    glfwInit();
    glfwWindowHint(GLFW_CONTEXT_VERSION_MAJOR, 4);
    glfwWindowHint(GLFW_CONTEXT_VERSION_MINOR, 3);
    glfwWindowHint(GLFW_OPENGL_PROFILE, GLFW_OPENGL_CORE_PROFILE);

    constexpr int screen_width = 800;
    constexpr int screen_height = 600;
    GLFWwindow* window = glfwCreateWindow(screen_width, screen_height, "Compute Shader Image Invert", NULL, NULL);
    if (window == NULL) {
        std::cout << "Failed to create GLFW window" << std::endl;
        glfwTerminate();
        
    }
    glfwMakeContextCurrent(window);

    if (!gladLoadGLLoader((GLADloadproc)glfwGetProcAddress)) {
        std::cout << "Failed to initialize GLAD" << std::endl;
        
    }

    // --- Shader Programs ---
    GLuint computeProgram = create_shader_program(nullptr, nullptr, invertCptShaderSources);
    GLuint renderProgram = create_shader_program(invertVtxShaderSource, invertFrgShaderSource);

    // --- Load Image and Create Textures ---
    int width, height, nrChannels;
    stbi_set_flip_vertically_on_load(true);
    unsigned char* data = stbi_load("checkerboard.png", &width, &height, &nrChannels, 0);
    if (!data) {
        std::cout << "Failed to load texture" << std::endl;
        
    }

    GLuint inputTexture, outputTexture;
    glGenTextures(1, &inputTexture);
    glActiveTexture(GL_TEXTURE0);
    glBindTexture(GL_TEXTURE_2D, inputTexture);
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_S, GL_CLAMP_TO_EDGE);
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_T, GL_CLAMP_TO_EDGE);
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_LINEAR);
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_LINEAR);
    glTexImage2D(GL_TEXTURE_2D, 0, GL_RGBA8, width, height, 0, GL_RGB, GL_UNSIGNED_BYTE, data);
    stbi_image_free(data);

    glGenTextures(1, &outputTexture);
    glActiveTexture(GL_TEXTURE1);
    glBindTexture(GL_TEXTURE_2D, outputTexture);
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_S, GL_CLAMP_TO_EDGE);
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_T, GL_CLAMP_TO_EDGE);
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_LINEAR);
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_LINEAR);
    glTexImage2D(GL_TEXTURE_2D, 0, GL_RGBA8, width, height, 0, GL_RGBA, GL_UNSIGNED_BYTE, NULL);

    // --- Dispatch Compute Shader ---
    glUseProgram(computeProgram);
    glBindImageTexture(0, inputTexture, 0, GL_FALSE, 0, GL_READ_ONLY, GL_RGBA8);
    glBindImageTexture(1, outputTexture, 0, GL_FALSE, 0, GL_WRITE_ONLY, GL_RGBA8);
    glDispatchCompute((GLuint)width / 8, (GLuint)height / 8, 1);

    // Block until compute shader is done
    glMemoryBarrier(GL_SHADER_IMAGE_ACCESS_BARRIER_BIT);

    // --- Setup Quad for Rendering ---
    float quadVertices[] = {
        // positions   // texCoords
        -1.0f,  1.0f,  0.0f, 1.0f,
        -1.0f, -1.0f,  0.0f, 0.0f,
         1.0f, -1.0f,  1.0f, 0.0f,

        -1.0f,  1.0f,  0.0f, 1.0f,
         1.0f, -1.0f,  1.0f, 0.0f,
         1.0f,  1.0f,  1.0f, 1.0f
    };
    GLuint quadVAO, quadVBO;
    glGenVertexArrays(1, &quadVAO);
    glGenBuffers(1, &quadVBO);
    glBindVertexArray(quadVAO);
    glBindBuffer(GL_ARRAY_BUFFER, quadVBO);
    glBufferData(GL_ARRAY_BUFFER, sizeof(quadVertices), &quadVertices, GL_STATIC_DRAW);
    glEnableVertexAttribArray(0);
    glVertexAttribPointer(0, 2, GL_FLOAT, GL_FALSE, 4 * sizeof(float), (void*)0);
    glEnableVertexAttribArray(1);
    glVertexAttribPointer(1, 2, GL_FLOAT, GL_FALSE, 4 * sizeof(float), (void*)(2 * sizeof(float)));
    constexpr int half_width = screen_width >> 1;
    constexpr int half_height = screen_height >> 1;

    {// Save Invert Result image
        // 1. Create a buffer on the CPU to hold the pixel data
        const int num_pixels = width * height;
        const int channels = 4; // We are reading RGBA data
        std::vector<unsigned char> pixels(num_pixels * channels);

        // 2. Bind the output texture that we want to read from
        glBindTexture(GL_TEXTURE_2D, outputTexture);

        // 3. Read the texture data from the GPU into our CPU buffer
        glGetTexImage(GL_TEXTURE_2D, 0, GL_RGBA, GL_UNSIGNED_BYTE, pixels.data());

        // 4. Use stb_image_write to save the pixels to a file
        // Most image formats expect the top-left pixel to be first, but OpenGL
        // gives us the bottom-left first. We can tell stb to flip it on write.
        stbi_flip_vertically_on_write(true);
        if (stbi_write_png("invert_img.png", width, height, channels, pixels.data(), width * channels)) {
            std::cout << "Image successfully saved!" << std::endl;
        }
        else {
            std::cerr << "Failed to save image." << std::endl;
        }
    }
    // --- Render Loop ---
    while (!glfwWindowShouldClose(window)) {
        glClear(GL_COLOR_BUFFER_BIT);

        glUseProgram(renderProgram);
        glViewport(0, 0, half_width, half_height);
        glBindVertexArray(quadVAO);
        glActiveTexture(GL_TEXTURE0);
        glBindTexture(GL_TEXTURE_2D, inputTexture);
        glDrawArrays(GL_TRIANGLES, 0, 6);

        glViewport(half_width, half_height, half_width, half_height);
        glBindTexture(GL_TEXTURE_2D, outputTexture); 
        glDrawArrays(GL_TRIANGLES, 0, 6);
        glfwSwapBuffers(window);
        glfwPollEvents();
    }

    // --- Cleanup ---
    glDeleteVertexArrays(1, &quadVAO);
    glDeleteBuffers(1, &quadVBO);
    glDeleteProgram(computeProgram);
    glDeleteProgram(renderProgram);
    glDeleteTextures(1, &inputTexture);
    glDeleteTextures(1, &outputTexture);

    glfwTerminate();
}