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

#include <glm/glm.hpp>
#include <glm/gtc/matrix_transform.hpp>
#include <glm/gtc/type_ptr.hpp>

#include <vector>
#include <cstdlib>
#include <ctime>
#include <iostream>

// Window dimensions
const int WINDOW_WIDTH  = 800;
const int WINDOW_HEIGHT = 600;

// We might have 1K images per frame.
const int NUM_IMAGES = 1000;

// Define a texture atlas size (must be large enough to pack all images)
const int ATLAS_WIDTH  = 2048;
const int ATLAS_HEIGHT = 2048;

// A simple structure to hold an RGB image (3 channels, no alpha) with variable dimensions.
struct Image {
    int width;
    int height;
    std::vector<unsigned char> data; // packed as RGB (3 bytes per pixel)
};

// Generate a random image with given dimensions.
Image generateRandomImage(int width, int height) {
    Image img;
    img.width  = width;
    img.height = height;
    img.data.resize(width * height * 3);
    for (int i = 0; i < width * height * 3; i++) {
        img.data[i] = static_cast<unsigned char>(std::rand() % 256);
    }
    return img;
}

// This structure holds per-image instance data (destination and atlas coordinates).
struct InstanceData {
    // Where on screen to draw this image quad (in pixels)
    float destX, destY, destW, destH;
    // Texture coordinates (normalized: 0.0–1.0) for this image within the atlas.
    float atlasX, atlasY, atlasW, atlasH;
};

//
// Main function
//
int main()
{
    std::srand(static_cast<unsigned int>(std::time(nullptr)));
    if (!glfwInit()){
        std::cerr << "Failed to init GLFW" << std::endl;
        return -1;
    }
    // Request OpenGL 3.3 Core Profile.
    glfwWindowHint(GLFW_CONTEXT_VERSION_MAJOR, 3);
    glfwWindowHint(GLFW_CONTEXT_VERSION_MINOR, 3);
    glfwWindowHint(GLFW_OPENGL_PROFILE, GLFW_OPENGL_CORE_PROFILE);

    GLFWwindow* window = glfwCreateWindow(WINDOW_WIDTH, WINDOW_HEIGHT, "Atlas Rendering with Instancing", nullptr, nullptr);
    if (!window){
        std::cerr << "Failed to create GLFW window" << std::endl;
        glfwTerminate();
        return -1;
    }
    glfwMakeContextCurrent(window);

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

    // -------------------------------------------------------------------
    // Create shaders (using instanced attributes)
    const char* vertexShaderSrc = R"(
        #version 330 core
        // Per-vertex attributes for a unit quad.
        layout (location = 0) in vec2 aPos;   // positions in [0,1]
        layout (location = 1) in vec2 aTex;   // texture coordinates in [0,1]
        // Instance attributes:
        layout (location = 2) in vec4 aDest;  // destination rectangle: x, y, w, h (in screen pixels)
        layout (location = 3) in vec4 aAtlas; // atlas coordinates: x, y, w, h (normalized 0–1)

        uniform mat4 uProjection; // Orthographic projection.

        out vec2 TexCoord;

        void main(){
            // Scale the unit quad (0..1) to the destination rectangle.
            vec2 pos = aPos * aDest.zw + aDest.xy;
            gl_Position = uProjection * vec4(pos, 0.0, 1.0);
            // Compute per-vertex texture coordinate from atlas region.
            TexCoord = aTex * aAtlas.zw + aAtlas.xy;
        }
    )";

    const char* fragmentShaderSrc = R"(
        #version 330 core
        in vec2 TexCoord;
        out vec4 FragColor;

        uniform sampler2D uAtlas;

        void main(){
            FragColor = texture(uAtlas, TexCoord);
        }
    )";

    // Simple shader compilation utility.
    auto compile_shader = [](GLenum type, const char* src) -> GLuint {
        GLuint shader = glCreateShader(type);
        glShaderSource(shader, 1, &src, nullptr);
        glCompileShader(shader);
        int success;
        glGetShaderiv(shader, GL_COMPILE_STATUS, &success);
        if (!success){
            char infoLog[512];
            glGetShaderInfoLog(shader, 512, nullptr, infoLog);
            std::cerr << "Shader compilation error: " << infoLog << std::endl;
        }
        return shader;
    };

    GLuint vs = compile_shader(GL_VERTEX_SHADER, vertexShaderSrc);
    GLuint fs = compile_shader(GL_FRAGMENT_SHADER, fragmentShaderSrc);
    GLuint shaderProgram = glCreateProgram();
    glAttachShader(shaderProgram, vs);
    glAttachShader(shaderProgram, fs);
    glLinkProgram(shaderProgram);
    glDeleteShader(vs);
    glDeleteShader(fs);

    // -------------------------------------------------------------------
    // Create a unit quad (two triangles) VBO/VAO.
    float quadVertices[] = {
        // positions    // texCoords
         0.0f, 0.0f,    0.0f, 0.0f,  // bottom-left
         1.0f, 0.0f,    1.0f, 0.0f,  // bottom-right
         1.0f, 1.0f,    1.0f, 1.0f,  // top-right

         0.0f, 0.0f,    0.0f, 0.0f,  // bottom-left
         1.0f, 1.0f,    1.0f, 1.0f,  // top-right
         0.0f, 1.0f,    0.0f, 1.0f   // top-left
    };
    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);
      glVertexAttribPointer(0, 2, GL_FLOAT, GL_FALSE, 4 * sizeof(float), (void*)0);
      glEnableVertexAttribArray(0);
      glVertexAttribPointer(1, 2, GL_FLOAT, GL_FALSE, 4 * sizeof(float), (void*)(2 * sizeof(float)));
      glEnableVertexAttribArray(1);
    glBindVertexArray(0);

    // -------------------------------------------------------------------
    // Create a texture atlas and pack all small images into it.
    // We'll allocate an empty atlas buffer and then copy each image's data into it.
    std::vector<unsigned char> atlasData(ATLAS_WIDTH * ATLAS_HEIGHT * 3, 0);
    std::vector<Image> images;
    std::vector<InstanceData> instances;

    // For a simple atlas packer, we use a "shelf" algorithm.
    int curX = 0, curY = 0, rowHeight = 0;
    // Also arrange destination positions in a grid.
    int gridCols = 32; // number of images per row on screen
    int cellSize  = 40;

    for (int i = 0; i < NUM_IMAGES; i++){
        // Random image size: for demo, vary between 16 and 40 pixels.
        int w = 16 + std::rand() % 25;
        int h = 16 + std::rand() % 25;
        Image img = generateRandomImage(w, h);
        images.push_back(img);

        // Pack into atlas: if not enough room in current row, move to next row.
        if (curX + w > ATLAS_WIDTH) {
            curX = 0;
            curY += rowHeight;
            rowHeight = 0;
        }
        if (curY + h > ATLAS_HEIGHT) {
            std::cerr << "Atlas overflow!" << std::endl;
            break;
        }

        // Copy pixel data into atlasData at position (curX, curY)
        for (int row = 0; row < h; row++){
            int destRow = curY + row;
            int srcIndex  = row * w * 3;
            int destIndex = (destRow * ATLAS_WIDTH + curX) * 3;
            std::copy(images[i].data.begin() + srcIndex,
                      images[i].data.begin() + srcIndex + w * 3,
                      atlasData.begin() + destIndex);
        }

        // Compute normalized atlas coordinates for this image.
        float atlasX = float(curX) / ATLAS_WIDTH;
        float atlasY = float(curY) / ATLAS_HEIGHT;
        float atlasW = float(w) / ATLAS_WIDTH;
        float atlasH = float(h) / ATLAS_HEIGHT;

        // Set destination rectangle (for demonstration, arrange in a grid).
        int col = i % gridCols;
        int row = i / gridCols;
        float destX = col * cellSize + 10.0f;
        float destY = row * cellSize + 10.0f;
        float destW = float(w); // you might choose to scale these to a fixed cell size
        float destH = float(h);

        InstanceData inst;
        inst.destX = destX; inst.destY = destY; inst.destW = destW; inst.destH = destH;
        inst.atlasX = atlasX; inst.atlasY = atlasY; inst.atlasW = atlasW; inst.atlasH = atlasH;
        instances.push_back(inst);

        curX += w; // advance current x in atlas
        if (h > rowHeight) rowHeight = h;
    }

    // Create the atlas texture.
    GLuint atlasTexture;
    glGenTextures(1, &atlasTexture);
    glBindTexture(GL_TEXTURE_2D, atlasTexture);
    glTexImage2D(GL_TEXTURE_2D, 0, GL_RGB, ATLAS_WIDTH, ATLAS_HEIGHT, 0, GL_RGB, GL_UNSIGNED_BYTE, atlasData.data());
    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_MIN_FILTER, GL_LINEAR);
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_LINEAR);

    // -------------------------------------------------------------------
    // Create and fill an instance VBO with per-image data.
    GLuint instanceVBO;
    glGenBuffers(1, &instanceVBO);
    glBindBuffer(GL_ARRAY_BUFFER, instanceVBO);
    glBufferData(GL_ARRAY_BUFFER, instances.size() * sizeof(InstanceData), instances.data(), GL_DYNAMIC_DRAW);

    // Attach instance data to the quad VAO.
    glBindVertexArray(quadVAO);
      // Location 2: destination rectangle (vec4)
      glVertexAttribPointer(2, 4, GL_FLOAT, GL_FALSE, sizeof(InstanceData), (void*)0);
      glEnableVertexAttribArray(2);
      glVertexAttribDivisor(2, 1); // update per instance
      // Location 3: atlas rectangle (vec4)
      glVertexAttribPointer(3, 4, GL_FLOAT, GL_FALSE, sizeof(InstanceData), (void*)(4 * sizeof(float)));
      glEnableVertexAttribArray(3);
      glVertexAttribDivisor(3, 1);
    glBindVertexArray(0);

    // Create an orthographic projection matrix.
    glm::mat4 projection = glm::ortho(0.0f, float(WINDOW_WIDTH), 0.0f, float(WINDOW_HEIGHT), -1.0f, 1.0f);

    // -------------------------------------------------------------------
    // Render loop.
    while (!glfwWindowShouldClose(window))
    {
        glClearColor(0.2f, 0.2f, 0.2f, 1.0f);
        glClear(GL_COLOR_BUFFER_BIT);

        glUseProgram(shaderProgram);
        glUniformMatrix4fv(glGetUniformLocation(shaderProgram, "uProjection"), 1, GL_FALSE, glm::value_ptr(projection));

        // Bind the atlas texture once.
        glActiveTexture(GL_TEXTURE0);
        glBindTexture(GL_TEXTURE_2D, atlasTexture);
        glUniform1i(glGetUniformLocation(shaderProgram, "uAtlas"), 0);

        glBindVertexArray(quadVAO);
        // Draw all image quads with instancing.
        glDrawArraysInstanced(GL_TRIANGLES, 0, 6, instances.size());

        glfwSwapBuffers(window);
        glfwPollEvents();
    }

    // Cleanup.
    glDeleteVertexArrays(1, &quadVAO);
    glDeleteBuffers(1, &quadVBO);
    glDeleteBuffers(1, &instanceVBO);
    glDeleteTextures(1, &atlasTexture);
    glDeleteProgram(shaderProgram);
    glfwDestroyWindow(window);
    glfwTerminate();

    return 0;
}