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

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

#include <iostream>
#include <fstream>
#include <sstream>
#include <stdio.h>
#include <stdlib.h>

#define VERTEX_FILE "./shaders/vertex.glsl"
#define FRAGMENT_FILE "./shaders/fragment.glsl"
#define COMPUTE_FILE "./shaders/compute.glsl"

const float SCR_WIDTH = 1200.0f;
const float SCR_HEIGHT = 1200.0f;

// Shader loader helper
std::string loadShaderSource(const std::string& filepath) {
    std::ifstream file(filepath);
    if (!file.is_open()) {
        std::cerr << "Failed to open shader file: " << filepath << std::endl;
        return "";
    }
    std::stringstream buffer;

    std::cout << "Loaded " << filepath << " successfully!\n";

    buffer << file.rdbuf();
    return buffer.str();
}

GLuint compileShader(GLenum type, const std::string& source) {
    GLuint shader = glCreateShader(type);
    const char* src = source.c_str();
    glShaderSource(shader, 1, &src, nullptr);
    glCompileShader(shader);

    int success;
    glGetShaderiv(shader, GL_COMPILE_STATUS, &success);
    if (!success) {
        char infoLog[1024];
        glGetShaderInfoLog(shader, 1024, nullptr, infoLog);
        std::cerr << "Shader compilation failed:\n" << infoLog << std::endl;
        exit(10);
    }
    return shader;
}

GLuint createShaderProgram(const std::string& vertexPath, const std::string& fragmentPath) {
    std::string vertexSrc = loadShaderSource(vertexPath);
    std::string fragmentSrc = loadShaderSource(fragmentPath);

    GLuint vertexShader = compileShader(GL_VERTEX_SHADER, vertexSrc);
    GLuint fragmentShader = compileShader(GL_FRAGMENT_SHADER, fragmentSrc);

    GLuint program = glCreateProgram();
    glAttachShader(program, vertexShader);
    glAttachShader(program, fragmentShader);
    glLinkProgram(program);

    int success;
    glGetProgramiv(program, GL_LINK_STATUS, &success);
    if (!success) {
        char infoLog[1024];
        glGetProgramInfoLog(program, 1024, nullptr, infoLog);
        std::cerr << "Shader linking failed:\n" << infoLog << std::endl;
        exit(11);
    }

    glDeleteShader(vertexShader);
    glDeleteShader(fragmentShader);
    return program;
}

GLuint createComputeProgram(const std::string& computePath) {
    std::string computeSrc = loadShaderSource(computePath);
    GLuint computeShader = compileShader(GL_COMPUTE_SHADER, computeSrc);

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

    int success;
    glGetProgramiv(program, GL_LINK_STATUS, &success);
    if (!success) {
        char infoLog[1024];
        glGetProgramInfoLog(program, 1024, nullptr, infoLog);
        std::cerr << "Compute shader linking failed:\n" << infoLog << std::endl;;
    }

    glDeleteShader(computeShader);
    return program;
}

int main() {
    // Initialize GLFW
    if (!glfwInit()) return -1;

    glfwWindowHint(GLFW_CONTEXT_VERSION_MAJOR, 4);
    glfwWindowHint(GLFW_CONTEXT_VERSION_MINOR, 1);
    glfwWindowHint(GLFW_OPENGL_FORWARD_COMPAT, GL_TRUE); // REQUIRED on macOS
    glfwWindowHint(GLFW_OPENGL_PROFILE, GLFW_OPENGL_CORE_PROFILE);
    glfwWindowHint(GLFW_SAMPLES, 8); // 8x MSAA

    GLFWwindow* window = glfwCreateWindow((int)SCR_WIDTH, (int)SCR_HEIGHT, "RASTER", nullptr, nullptr);
    if (!window) {
        glfwTerminate();
        return -1;
    }
    glfwMakeContextCurrent(window);

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

    glViewport(0, 0, (int)SCR_WIDTH, (int)SCR_HEIGHT);

    //DEBUG
    std::cout << "OpenGL Version: " << glGetString(GL_VERSION) << std::endl;

    // ----- Vertex Data -----
    float vertices[] = {
        // positions
        -1.0f, -1.0f, 0.0f,
         1.0f, -1.0f, 0.0f,
         1.0f,  1.0f, 0.0f,
        -1.0f,  1.0f, 0.0f
    };

    unsigned int indices[] = {
        0, 1, 2,
        2, 3, 0
    };

    GLuint VAO, VBO, EBO;
    glGenVertexArrays(1,&VAO);
    glGenBuffers(1,&VBO);
    glGenBuffers(1,&EBO);

    glBindVertexArray(VAO);

    glBindBuffer(GL_ARRAY_BUFFER,VBO);
    glBufferData(GL_ARRAY_BUFFER,sizeof(vertices),vertices,GL_STATIC_DRAW);

    glBindBuffer(GL_ELEMENT_ARRAY_BUFFER,EBO);
    glBufferData(GL_ELEMENT_ARRAY_BUFFER,sizeof(indices),indices,GL_STATIC_DRAW);

    // position - VERTEX SHADER
    glVertexAttribPointer(0,3,GL_FLOAT,GL_FALSE,3*sizeof(float),(void*)0);
    glEnableVertexAttribArray(0);

    glBindBuffer(GL_ARRAY_BUFFER,0);

    // Load shaders
    GLuint shaderProgram = createShaderProgram(VERTEX_FILE, FRAGMENT_FILE);

    // Caulculate model, view, projection matrices
    glm::mat4 model = glm::mat4(1.0f);
    glm::mat4 view = glm::translate(glm::mat4(1.0f), glm::vec3(0.0f, 0.0f, -1.0f));
    glm::mat4 projection = glm::perspective(glm::radians(45.0f), SCR_WIDTH/SCR_HEIGHT, 0.1f, 100.0f);

    GLuint modelLoc = (unsigned int) glGetUniformLocation(shaderProgram, "model");
    GLuint viewLoc = (unsigned int) glGetUniformLocation(shaderProgram, "view");
    GLuint projLoc = (unsigned int) glGetUniformLocation(shaderProgram, "projection");

    glUseProgram(shaderProgram);
    glUniformMatrix4fv((int) modelLoc, 1, GL_FALSE, glm::value_ptr(model));
    glUniformMatrix4fv((int) viewLoc, 1, GL_FALSE, glm::value_ptr(view));
    glUniformMatrix4fv((int) projLoc, 1, GL_FALSE, glm::value_ptr(projection));

    // Query important uniform locations (we'll keep them around)
    GLint resLoc = glGetUniformLocation(shaderProgram, "u_resolution");
    GLint timeLoc = glGetUniformLocation(shaderProgram, "u_time");
    GLint scaleLoc = glGetUniformLocation(shaderProgram, "u_scale");
    GLint mouseLoc = glGetUniformLocation(shaderProgram, "u_mouse");

    glEnable(GL_MULTISAMPLE);

    // --- Supersampling FBO ---
    int scaleFactor = 2; // try 2 or 4 for smoother output
    int fbWidth = (int)SCR_WIDTH * scaleFactor;
    int fbHeight = (int)SCR_HEIGHT * scaleFactor;

    unsigned int fbo, colorTex, rbo;
    glGenFramebuffers(1, &fbo);
    glBindFramebuffer(GL_FRAMEBUFFER, fbo);

    // Color texture
    glGenTextures(1, &colorTex);
    glBindTexture(GL_TEXTURE_2D, colorTex);
    glTexImage2D(GL_TEXTURE_2D, 0, GL_RGB, fbWidth, fbHeight, 0, GL_RGB, GL_UNSIGNED_BYTE, NULL);
    // filtering + wrap
    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);

    glFramebufferTexture2D(GL_FRAMEBUFFER, GL_COLOR_ATTACHMENT0, GL_TEXTURE_2D, colorTex, 0);

    // Depth/stencil buffer
    glGenRenderbuffers(1, &rbo);
    glBindRenderbuffer(GL_RENDERBUFFER, rbo);
    glRenderbufferStorage(GL_RENDERBUFFER, GL_DEPTH24_STENCIL8, fbWidth, fbHeight);
    glFramebufferRenderbuffer(GL_FRAMEBUFFER, GL_DEPTH_STENCIL_ATTACHMENT, GL_RENDERBUFFER, rbo);

    if (glCheckFramebufferStatus(GL_FRAMEBUFFER) != GL_FRAMEBUFFER_COMPLETE)
        std::cout << "ERROR: Framebuffer not complete!" << std::endl;

    glBindFramebuffer(GL_FRAMEBUFFER, 0);
    // --------------------------

    // set the shader resolution to the FBO size (important!)
    glUseProgram(shaderProgram);
    glUniform2f(resLoc, SCR_WIDTH, SCR_HEIGHT);

    float scale = 1.0f;
    int frameCount = 0;
    float mousePos[2] = {0.0f, 0.0f};

    // FPS tracking
    double lastTime = glfwGetTime();   // when we last printed FPS
    double fps = 0.0;

    // Render loop
    while (!glfwWindowShouldClose(window)) {
        float timeValue = (float)glfwGetTime();
        double xpos, ypos;

        // Get mouse positions
        glfwGetCursorPos(window, &xpos, &ypos);
        mousePos[0] = (float)xpos;
        mousePos[1] = (float)(SCR_HEIGHT - ypos); // Invert y-axis for OpenGL coordinates

        // Input
        glfwPollEvents();

        // Shader hot-reload
        if (glfwGetKey(window, GLFW_KEY_R) == GLFW_PRESS) {
            glDeleteProgram(shaderProgram);
            shaderProgram = createShaderProgram(VERTEX_FILE, FRAGMENT_FILE);
            modelLoc = (unsigned int) glGetUniformLocation(shaderProgram, "model");
            viewLoc = (unsigned int) glGetUniformLocation(shaderProgram, "view");
            projLoc = (unsigned int) glGetUniformLocation(shaderProgram, "projection");

            glUseProgram(shaderProgram);
            glUniformMatrix4fv((int) modelLoc, 1, GL_FALSE, glm::value_ptr(model));
            glUniformMatrix4fv((int) viewLoc, 1, GL_FALSE, glm::value_ptr(view));
            glUniformMatrix4fv((int) projLoc, 1, GL_FALSE, glm::value_ptr(projection));

            // Re-query uniforms and set resolution to FBO size (IMPORTANT)
            resLoc = glGetUniformLocation(shaderProgram, "u_resolution");
            timeLoc = glGetUniformLocation(shaderProgram, "u_time");
            scaleLoc = glGetUniformLocation(shaderProgram, "u_scale");
            mouseLoc = glGetUniformLocation(shaderProgram, "u_mouse");

            glUseProgram(shaderProgram);
            glUniform2f(resLoc, (float)fbWidth, (float)fbHeight);

            std::cout << "Shaders recompiled!\n";
        }

        // Scroll to zoom in/out
        if (glfwGetKey(window, GLFW_KEY_UP) == GLFW_PRESS) {
            scale += 0.01f;
            if (scale > 10.0f) scale = 10.0f; // Max zoom
        }

        if (glfwGetKey(window, GLFW_KEY_DOWN) == GLFW_PRESS) {
            scale -= 0.01f;
            if (scale < 0.1f) scale = 0.1f; // Min zoom
        }

        // increment frame count (for FPS)
        frameCount++;

        // --- FPS Counter (print to console every second) ---
        double currentTime = glfwGetTime();
        double delta = currentTime - lastTime;
        if (delta >= 1.0) { // print every ~1 second
            fps = double(frameCount) / delta;
            std::cout << "FPS: " << fps
                  << " Time: " << timeValue
                  << " Scale: " << scale
                  << " Mouse: (" << mousePos[0] << ", " << mousePos[1] << ")"
                  << std::endl;

            frameCount = 0;
            lastTime = currentTime;
        }

        // --- Render into high-res FBO ---
        glBindFramebuffer(GL_FRAMEBUFFER, fbo);
        glViewport(0, 0, fbWidth, fbHeight);
        glClearColor(0.0f, 0.0f, 0.0f, 1.0f);
        glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT);

        // Use shader & set uniforms
        glUseProgram(shaderProgram);
        glUniform1f(timeLoc, timeValue);
        glUniform1f(scaleLoc, scale);
        glUniform2f(mouseLoc, mousePos[0], mousePos[1]);
        // (u_resolution already set to fbWidth, fbHeight)

        glBindVertexArray(VAO);
        glDrawElements(GL_TRIANGLES, 6, GL_UNSIGNED_INT, 0);

        // --- Blit back to screen ---
        glBindFramebuffer(GL_READ_FRAMEBUFFER, fbo);
        glBindFramebuffer(GL_DRAW_FRAMEBUFFER, 0);
        glBlitFramebuffer(0, 0, fbWidth, fbHeight,
                          0, 0, (int)SCR_WIDTH, (int)SCR_HEIGHT,
                          GL_COLOR_BUFFER_BIT, GL_LINEAR);

        glBindFramebuffer(GL_FRAMEBUFFER, 0);
        glViewport(0, 0, (int)SCR_WIDTH, (int)SCR_HEIGHT);

        glfwSwapBuffers(window);
    }

    glDeleteVertexArrays(1,&VAO);
    glDeleteBuffers(1,&VBO);
    glDeleteBuffers(1,&EBO);
    glDeleteProgram(shaderProgram);

    glfwDestroyWindow(window);
    glfwTerminate();
    return 0;
}