#include <iostream>
#include <unordered_set>

#include <glad/glad.h> // Must be included before GLFW
#include <GLFW/glfw3.h>
#include <glm/glm.hpp>

#include "Shader.h"
#include "Camera.h"
#include "Object3D.h"
#include "World.h"

// Window dimensions
unsigned int SCR_WIDTH = 1920;
unsigned int SCR_HEIGHT = 1080;

// Jump
const float JUMP_COOLDOWN = 0.0f;
const float JUMP_POWER = 10.0f; // Initial Jump Speed
float last_jump = 0.0f;

// Camera
const float CAMERA_SPEED = 1.0f;
const float MAX_CAMERA_SPEED = 2.0f;
const float LOOK_SENSITIVITY = 0.1f;

const float Z_FAR = 1000.0f;
const float Z_NEAR = 0.1f;
bool WIREFRAME_MODE = !true;
const float FOV = 110.0f;

bool W_pressed, S_pressed, A_pressed, D_pressed = false;

glm::mat4 projection = glm::mat4(1.0f);
Camera camera = Camera(CAMERA_SPEED, glm::vec3(0.0f, 0.0f, 3.0f));

// Callback function for when the window is resized
// glfw: whenever the window size changed (by OS or user resize) this callback function executes
void framebuffer_size_callback(GLFWwindow* window, int width, int height)
{
    // make sure the viewport matches the new window dimensions; note that width and
    // height will be significantly larger than specified on retina displays.
    glViewport(0, 0, width, height);
    SCR_WIDTH = width;
    SCR_HEIGHT = height;

    // Perspective needs to be recalculated on window size change
    projection = glm::perspective(glm::radians(FOV), (float)SCR_WIDTH / (float)SCR_HEIGHT, Z_NEAR, Z_FAR);
}

float yaw, pitch = 0;
bool first_mouse = true;
bool mouse_lock = true;
float last_x = 400, last_y = 300;

// Mouse Look
void mouse_callback(GLFWwindow * window, double xpos, double ypos)
{
    float x_offset = xpos - last_x;
    float y_offset = ypos - last_y;
    last_x = xpos;
    last_y = ypos;

    if (first_mouse)
    {
        first_mouse = false;
        return;
    }

    x_offset *= LOOK_SENSITIVITY;
    y_offset *= LOOK_SENSITIVITY;

    yaw += x_offset;
    pitch += y_offset;

    if (pitch > 89.0f)
        pitch = 89.0f;
    if (pitch < -89.0f)
        pitch = -89.0f;

    glm::vec3 direction;
    direction.x = cos(glm::radians(yaw)) * cos(glm::radians(pitch));
    direction.y = -sin(glm::radians(pitch));
    direction.z = sin(glm::radians(yaw)) * cos(glm::radians(pitch));
    camera.target = glm::normalize(direction);
}

// process all input: query GLFW whether relevant keys are pressed/released this frame and react accordingly
void processInput(GLFWwindow *window)
{
    float current_time = glfwGetTime();

    if(glfwGetKey(window, GLFW_KEY_ESCAPE) == GLFW_PRESS) {
        // if (mouse_lock) {
        //     glfwSetInputMode(window, GLFW_CURSOR, GLFW_CURSOR_NORMAL);
        //     mouse_lock = false;
        // } else {
        //     glfwSetInputMode(window, GLFW_CURSOR, GLFW_CURSOR_DISABLED);
        //     mouse_lock = true;
        // }
        glfwSetWindowShouldClose(window, true);
    }

    // ---- Movement ----
    // Forward / Backward
    if(glfwGetKey(window, GLFW_KEY_W) == GLFW_PRESS)
    {
        camera.forward_velocity = camera.speed * glm::vec3(camera.target.x, 0.0f, camera.target.z);
        W_pressed = true;
    }

    if(glfwGetKey(window, GLFW_KEY_S) == GLFW_PRESS)
    {
        camera.forward_velocity = -camera.speed * glm::vec3(camera.target.x, 0.0f, camera.target.z);
        S_pressed = true;
    }


    if(glfwGetKey(window, GLFW_KEY_S) == GLFW_RELEASE && S_pressed)
    {
        camera.forward_velocity = glm::vec3(0.0f);
        S_pressed = false;
    }
    if(glfwGetKey(window, GLFW_KEY_W) == GLFW_RELEASE && W_pressed)
    {
        camera.forward_velocity = glm::vec3(0.0f);
        W_pressed = false;
    }


    // Horizontal
    if(glfwGetKey(window, GLFW_KEY_A) == GLFW_PRESS)
    {
        camera.horizontal_velocity = -glm::normalize(glm::cross(glm::vec3(camera.target.x, 0.0f, camera.target.z), camera.worldUp)) * camera.speed;
        A_pressed = true;
    }
    if(glfwGetKey(window, GLFW_KEY_D) == GLFW_PRESS)
    {
        camera.horizontal_velocity = glm::normalize(glm::cross(glm::vec3(camera.target.x, 0.0f, camera.target.z), camera.worldUp)) * camera.speed;
        D_pressed = true;
    }


    if(glfwGetKey(window, GLFW_KEY_A) == GLFW_RELEASE && A_pressed)
    {
        camera.horizontal_velocity = glm::vec3(0.0f);
        A_pressed = false;
    }
    if(glfwGetKey(window, GLFW_KEY_D) == GLFW_RELEASE && D_pressed)
    {
        camera.horizontal_velocity = glm::vec3(0.0f);
        D_pressed = false;
    }


    // Jump
    if(glfwGetKey(window, GLFW_KEY_SPACE) == GLFW_PRESS)
        // camera.velocity -= glm::normalize(glm::cross(camera.target, glm::normalize(glm::cross(camera.target, camera.worldUp)))) * camera.speed * 0.5f;
    {
        if (current_time - last_jump > JUMP_COOLDOWN && camera.position.y == 0)
        {
            camera.vertical_velocity.y = JUMP_POWER;
            last_jump = current_time;
        }
    }

    // Sprint
    if(glfwGetKey(window, GLFW_KEY_LEFT_SHIFT) == GLFW_PRESS)
        camera.speed = CAMERA_SPEED * 7.5;
    if(glfwGetKey(window, GLFW_KEY_LEFT_SHIFT) == GLFW_RELEASE)
        camera.speed = CAMERA_SPEED * 20;
}

// Draw a cube at some x, y, z
// [UPDATED] Added a new 'color' parameter
void draw_cube(glm::ivec3 coords, Shader& shader, unsigned int VAO, size_t indexCount, const glm::mat4& view, const glm::mat4& projection, const glm::vec3& color) {
    // Create model matrix for positioning the cube
    glm::mat4 model = glm::mat4(1.0f);
    model = glm::translate(model, glm::vec3(coords));

    // Calculate MVP matrix
    glm::mat4 mvp = projection * view * model;

    // Set the uniform and draw
    shader.use();
    shader.setMat4("u_mvp", mvp);
    // [NEW] Set the color uniform in the fragment shader
    shader.setVec3("u_Color", color);

    glBindVertexArray(VAO);
    glDrawElements(GL_TRIANGLES, indexCount, GL_UNSIGNED_INT, nullptr);
}
int main() {
    // 1. --- Initialize GLFW ---
    if (!glfwInit()) {
        std::cerr << "Failed to initialize GLFW" << std::endl;
        return -1;
    }
    glfwWindowHint(GLFW_CONTEXT_VERSION_MAJOR, 3);
    glfwWindowHint(GLFW_CONTEXT_VERSION_MINOR, 3);
    glfwWindowHint(GLFW_OPENGL_PROFILE, GLFW_OPENGL_CORE_PROFILE);

    // 2. --- Create a Window ---
    GLFWwindow* window = glfwCreateWindow(SCR_WIDTH, SCR_HEIGHT, "HelloTriangle", NULL, NULL);
    if (window == NULL) {
        std::cerr << "Failed to create GLFW window" << std::endl;
        glfwTerminate();
        return -1;
    }
    glfwMakeContextCurrent(window);
    glfwSetFramebufferSizeCallback(window, framebuffer_size_callback);

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

    glViewport(0, 0, SCR_WIDTH, SCR_HEIGHT);

    // 4. Shader Setup
    Shader shader("shaders/cube_vert.glsl", "shaders/cube_frag.glsl");

    // 5. --- Set up Vertex Data and Buffers ---

    // Load Cube OBJ
    World world;
    Object3D cube = Object3D("objs/cube.obj");

    // Place cubes in the world
    for (int z = 0; z <= 32; z++) {
        for (int x = 0; x <= 32; x++) {
            world.voxels.insert(glm::ivec3(x, -1, z));
        }
    }

    world.voxels.insert(glm::ivec3(7, 0, 13));

    unsigned int VBO, VAO, EBO;
    glGenVertexArrays(1, &VAO); // 1. Create Vertex Array Object (VAO)
    glGenBuffers(1, &VBO);      // 2. Create Vertex Buffer Object (VBO)
    glGenBuffers(1, &EBO);

    glBindVertexArray(VAO);     // 3. Bind the VAO

    glBindBuffer(GL_ARRAY_BUFFER, VBO); // 4. Bind the VBO
    // 5. Copy vertex data into VBO's memory
    glBufferData(GL_ARRAY_BUFFER, cube.VBO_buffer.size() * sizeof(float), &cube.VBO_buffer.front(), GL_STATIC_DRAW);

    // 6. Tell OpenGL how to interpret the vertex data

    // Position Attribute
    glVertexAttribPointer(0, 3, GL_FLOAT, GL_FALSE, 6 * sizeof(float), (void*)0);
    glEnableVertexAttribArray(0); // Enable the vertex attribute (location 0)

    // Normal Attribute
    glVertexAttribPointer(1, 3, GL_FLOAT, GL_FALSE, 6 * sizeof(float), (void*)(3 * sizeof(float)));
    glEnableVertexAttribArray(1);

    // 6.5 Setup EBO
    glBindBuffer(GL_ELEMENT_ARRAY_BUFFER, EBO);
    glBufferData(GL_ELEMENT_ARRAY_BUFFER, cube.EBO_buffer.size() * sizeof(unsigned int), &cube.EBO_buffer.front(), GL_STATIC_DRAW);

    // Unbind VBO and VAO (optional, but good practice)
    glBindBuffer(GL_ARRAY_BUFFER, 0);
    glBindVertexArray(0);

    if (WIREFRAME_MODE)
        glPolygonMode(GL_FRONT_AND_BACK, GL_LINE);

    // Setup mouse look
    glfwSetInputMode(window, GLFW_CURSOR, GLFW_CURSOR_DISABLED);
    glfwSetCursorPosCallback(window, mouse_callback);

    // Setup projection matrix
    projection = glm::perspective(glm::radians(FOV), (float)SCR_WIDTH / (float)SCR_HEIGHT, Z_NEAR, Z_FAR); // This doesn't need to be recalculated every frame

    glm::mat4 view;
    double start_time = glfwGetTime();
    double move_time = glfwGetTime();

    glEnable(GL_DEPTH_TEST);

    // 6. --- The Render Loop ---
    while (!glfwWindowShouldClose(window)) {
        // Input (e.g., close window on ESC)
        processInput(window);

        glClearColor(0.0f, 0.0f, 0.0f, 1.0f); // Clear screen
        glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT); // Clear Z-Buffer

        view = camera.getView();

        // ---- PASS 1: Draw solid grey cubes ----
        glEnable(GL_POLYGON_OFFSET_FILL);
        glPolygonOffset(1.0, 1.0); // Push solid faces "back"
        glPolygonMode(GL_FRONT_AND_BACK, GL_FILL);

        glm::vec3 solidColor = glm::vec3(0.5, 0.5, 0.5); // Grey
        for (glm::ivec3 voxel_pos : world.voxels) {
            // [UPDATED] Pass the grey color
            draw_cube(voxel_pos, shader, VAO, cube.EBO_buffer.size(), view, projection, solidColor);
        }

        glDisable(GL_POLYGON_OFFSET_FILL);

        // ---- PASS 2: Draw wireframe on top ----
        glPolygonMode(GL_FRONT_AND_BACK, GL_LINE); // Just the lines

        glm::vec3 wireframeColor = glm::vec3(1.0, 1.0, 1.0); // White
        for (glm::ivec3 voxel_pos : world.voxels) {
            // [UPDATED] Pass the white wireframe color
            draw_cube(voxel_pos, shader, VAO, cube.EBO_buffer.size(), view, projection, wireframeColor);
        }

        // --- Reset polygon mode for next frame (good practice) ---
        glPolygonMode(GL_FRONT_AND_BACK, GL_FILL);

        double end_time = glfwGetTime();
        float time_since_last_move = end_time - move_time;
        if (time_since_last_move  >= 0.017)
        {
            camera.move(time_since_last_move);
            move_time = end_time;
        }

        // Swap buffers and poll for events
        glfwSwapBuffers(window);
        glfwPollEvents();
    }

    // 7. --- Cleanup ---
    glDeleteVertexArrays(1, &VAO);
    glDeleteBuffers(1, &VBO);

    glfwTerminate(); // Clean up GLFW
    return 0;
}