#include <GL/glew.h>
#include <GLFW/glfw3.h>
#include <cmath>
#include <vector>
#include <iostream>
#include <glm/glm.hpp>
#include <glm/gtc/matrix_transform.hpp>
#include <glm/gtc/type_ptr.hpp>
#include "shader.h"
#include <cstdlib>
#include <imgui.h>
#include <imgui_impl_glfw.h>
#include <imgui_impl_opengl3.h>

const glm::vec3 initialCameraPos = glm::vec3(0, 0, 30);
const glm::vec3 initialCameraTarget = glm::vec3(0, 0, 0);
const glm::vec3 initialCameraUp = glm::vec3(0, 1, 0);
const float initialGravity = 9.81f;
const float initialAirResistance = 0.1f;
const float initialSpeed = 5.0f;
const float initialAngle = 45.0f;
const float ballRadius = 0.5;

const int numberOfPointsOnCircle = 120;

const int windowWidth = 1920;
const int windowHeight = 800;

bool isLaunched = false;

const int initBallsAmount = 3;
const int maxBallsAmount = 30;

const float initBoxEdgeLength = 20.0f;
const glm::vec3 initBoxCenter = glm::vec3(0.0f, 0.0f, 0.0f);

struct Ball {
    float x, y, z;
    float vx, vy, vz;
    float r, g, b;
};

struct Box {
    float edgeLength;
    glm::vec3 center;
};

Box box = {
        .edgeLength = initBoxEdgeLength,
        .center = initBoxCenter
};

int ballsAmount = initBallsAmount;


float gravity = initialGravity;
float airResistance = initialAirResistance;
float speed = initialSpeed;
float angle = initialAngle * M_PI / 180.0f;
float ballsAngle[maxBallsAmount] = {};





glm::vec3 cameraPos = initialCameraPos;
glm::vec3 cameraTarget = initialCameraTarget;
glm::vec3 cameraUp = initialCameraUp;

GLFWwindow* window;
std::vector<Ball> balls(maxBallsAmount);  // Wektor przechowujący ballsAmount kul
GLuint shaderProgram;
GLuint VAO_SPHERES, VBO_SPHERES;
GLuint VAO_AXES, VBO_AXES;
GLuint VAO_SQUARE, VBO_SQUARE, EBO_SQUARE;

void setupAxesBuffer() {
    // Wierzchołki osi (każda oś ma dwa wierzchołki: początek i koniec)
    float vertices[] = {
            // Oś X (czerwona)
            0.0f, 0.0f, 0.0f,
            10.0f, 0.0f, 0.0f,
            // Oś Y (zielona)
            0.0f, 0.0f, 0.0f,
            0.0f, 10.0f, 0.0f,
            // Oś Z (niebieska)
            0.0f, 0.0f, 0.0f,
            0.0f, 0.0f, 10.0f
    };

    glGenVertexArrays(1, &VAO_AXES);
    glGenBuffers(1, &VBO_AXES);

    glBindVertexArray(VAO_AXES);

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

    glVertexAttribPointer(0, 3, GL_FLOAT, GL_FALSE, 3 * sizeof(float), (void*)0);
    glEnableVertexAttribArray(0);

    glBindBuffer(GL_ARRAY_BUFFER, 0);
    glBindVertexArray(0);
}

void setupSquareBuffer() {
    float halfEdgeLength = box.edgeLength / 2.0f;

    float vertices[] = {
            box.center.x - halfEdgeLength, box.center.y - halfEdgeLength, 0.0f,
            box.center.x + halfEdgeLength, box.center.y - halfEdgeLength, 0.0f,
            box.center.x + halfEdgeLength,  box.center.y + halfEdgeLength, 0.0f,
            box.center.x - halfEdgeLength,  box.center.y + halfEdgeLength, 0.0f
    };

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

    glGenVertexArrays(1, &VAO_SQUARE);
    glGenBuffers(1, &VBO_SQUARE);
    glGenBuffers(1, &EBO_SQUARE);

    glBindVertexArray(VAO_SQUARE);

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

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

    glVertexAttribPointer(0, 3, GL_FLOAT, GL_FALSE, 3 * sizeof(float), (void*)0);
    glEnableVertexAttribArray(0);

    glBindBuffer(GL_ARRAY_BUFFER, 0);
    glBindVertexArray(0);
}



void drawAxes() {
    glm::mat4 model = glm::mat4(1.0f);
    GLint modelLoc = glGetUniformLocation(shaderProgram, "model");
    glUniformMatrix4fv(modelLoc, 1, GL_FALSE, glm::value_ptr(model));

    // Rysowanie osi
    glBindVertexArray(VAO_AXES);

    // Oś X
    GLint colorLoc = glGetUniformLocation(shaderProgram, "color");
    glUniform3f(colorLoc, 1.0f, 0.0f, 0.0f);
    glDrawArrays(GL_LINES, 0, 2);

    // Oś Y
    glUniform3f(colorLoc, 0.0f, 1.0f, 0.0f);
    glDrawArrays(GL_LINES, 2, 2);

    // Oś Z
    glUniform3f(colorLoc, 0.0f, 0.0f, 1.0f);
    glDrawArrays(GL_LINES, 4, 2);

    glBindVertexArray(0);
}

void drawSquare() {
    glm::mat4 model = glm::mat4(1.0f);
    GLint modelLoc = glGetUniformLocation(shaderProgram, "model");
    glUniformMatrix4fv(modelLoc, 1, GL_FALSE, glm::value_ptr(model));

    GLint colorLoc = glGetUniformLocation(shaderProgram, "color");
    glUniform3f(colorLoc, 0.5f, 1.0f, 0.1f);

    glBindVertexArray(VAO_SQUARE);
    glDrawElements(GL_LINES, 8, GL_UNSIGNED_INT, 0);
    glBindVertexArray(0);
}



void computeBallsCollision() {
    float boxHalfLength = box.edgeLength / 2.0f;

    for (int i = 0; i < ballsAmount; ++i) {
        Ball &ball1 = balls[i];
        for (int j = i + 1; j < ballsAmount; ++j) {
            Ball &ball2 = balls[j];
            float distance = sqrt(pow(ball1.x - ball2.x, 2) + pow(ball1.y - ball2.y, 2));
            float radiusSum = ballRadius * 2;

            if (distance < radiusSum) {
                // Obliczanie normalnej wektora kolizji
                float nx = (ball2.x - ball1.x) / distance;
                float ny = (ball2.y - ball1.y) / distance;
                float nz = (ball2.z - ball1.z) / distance;

                // Obliczanie względnej prędkości
                float kx = ball1.vx - ball2.vx;
                float ky = ball1.vy - ball2.vy;
                float kz = ball1.vz - ball2.vz;

                // Obliczanie impulsu wymiany
                float p = nx * kx + ny * ky + nz * kz;

                // Aktualizacja prędkości kulek
                ball1.vx -= p * nx;
                ball1.vy -= p * ny;
                ball1.vz -= p * nz;
                ball2.vx += p * nx;
                ball2.vy += p * ny;
                ball2.vz += p * nz;

                // Przesunięcie kulek, aby nie były już w punkcie kolizji
                float overlap = radiusSum - distance;
                ball1.x -= overlap * nx / 2.0f;
                ball1.y -= overlap * ny / 2.0f;
                ball1.z -= overlap * nz / 2.0f;
                ball2.x += overlap * nx / 2.0f;
                ball2.y += overlap * ny / 2.0f;
                ball2.z += overlap * nz / 2.0f;
            }
        }


        if (ball1.x - ballRadius < box.center.x - boxHalfLength) {
            ball1.x = -boxHalfLength + ballRadius;
            ball1.vx = -ball1.vx;
        }
        if (ball1.x + ballRadius > box.center.x + boxHalfLength) {
            ball1.x = boxHalfLength - ballRadius;
            ball1.vx = -ball1.vx;
        }
        if(ball1.y - ballRadius < box.center.y - boxHalfLength){
            ball1.vy = -ball1.vy;
            ball1.y = box.center.y - boxHalfLength + ballRadius;

        }
        if(ball1.y + ballRadius > box.center.y + boxHalfLength){
            ball1.vy = -ball1.vy;
            ball1.y = box.center.y + boxHalfLength - ballRadius;
        }

    }
}


void initializeBalls() {
    srand(static_cast<unsigned>(time(0)));

    for (int i = 0; i < ballsAmount; ++i) {
        float localAngle = angle + static_cast<float>(rand() % 60);  // <angle, angle + 60>
        float localSpeed = speed + static_cast<float>(rand() % 5);

        ballsAngle[i] = localAngle;

        balls[i] = {
                .x = 0.0f,
                .y = 2.0f * i,
                .z = 0.0f,
                .vx = static_cast<float>(localSpeed * cos(localAngle * M_PI / 180.0f)),
                .vy = static_cast<float>(localSpeed * sin(localAngle * M_PI / 180.0f)),
                .vz = 0.0f,
                .r = static_cast<float>(rand()) / RAND_MAX,
                .g = static_cast<float>(rand()) / RAND_MAX,
                .b = static_cast<float>(rand()) / RAND_MAX
        };
    }

    isLaunched = false;
}

void initImGui(){
    IMGUI_CHECKVERSION();
    ImGui::CreateContext();
    ImGuiIO& io = ImGui::GetIO(); (void)io;
    ImGui::StyleColorsDark();
    ImGui_ImplGlfw_InitForOpenGL(window, true);
    ImGui_ImplOpenGL3_Init("#version 330 core");
}

void imGuiRenderBoundingBoxFrame(){
    ImGui::Begin("Bounding Box");

    if(ImGui::SliderFloat("Box Edge Length", &box.edgeLength, 1.0f, 20.0f) || ImGui::SliderFloat3("Box Center", glm::value_ptr(box.center), -20.0f, 20.0f)) {
        setupSquareBuffer();
    };

    ImGui::End();
}

void imGuiRenderCameraControlsFrame(){
    ImGui::Begin("Camera Control");
    ImGui::SliderFloat3("Camera Position", glm::value_ptr(cameraPos), -100.0f, 100.0f);
    ImGui::SliderFloat3("Camera Target", glm::value_ptr(cameraTarget), -100.0f, 100.0f);
    ImGui::SliderFloat3("Camera Up", glm::value_ptr(cameraUp), -1.0f, 1.0f);

    if(ImGui::Button("Reset Camera")){
        cameraPos = initialCameraPos;
        cameraTarget = initialCameraTarget;
        cameraUp = initialCameraUp;
    }

    ImGui::End();
}

void imGuiRenderBallsControlsFrame(){
    ImGui::Begin("Balls Control");

    ImGui::SliderInt("Balls Amount", &ballsAmount, 1, maxBallsAmount);
    ImGui::SliderFloat("Gravity", &gravity, 0.0f, 20.0f);
    ImGui::SliderFloat("Air Resistance", &airResistance, 0.0f, 1.0f);
    ImGui::SliderFloat("Speed", &speed, 0.0f, 20.0f);
    ImGui::SliderFloat("Angle", &angle, 0.0f, 90.0f);

    if(ImGui::Button("Apply ball modifiers")){
        initializeBalls();
    }

    if(ImGui::Button("Reset ball modifiers")){
        gravity = initialGravity;
        airResistance = initialAirResistance;
        speed = initialSpeed;
        angle = initialAngle * M_PI / 180.0f;
    }

    ImGui::End();
}

void imGuiBallsThrowingStateControlsFrame(){
    ImGui::Begin("Balls Throwing State");

    if (ImGui::Button("Reset Balls")) {
        initializeBalls();
    }

    if (ImGui::Button("Launch Balls")) {
        isLaunched = true;
    }

    ImGui::End();
}

void imGuiBallsMetricFrame(){
    ImGui::Begin("Balls Metric");

    for(int i = 0; i < ballsAmount; i++){
        ImGui::Text("Ball %d Position: (%.2f, %.2f, %.2f)", i + 1, balls[i].x, balls[i].y, balls[i].z);
        ImGui::Text("Ball %d Angle: %.2f", i + 1, ballsAngle[i]);
    }

    ImGui::End();
}

void renderImGui(){
    ImGui_ImplOpenGL3_NewFrame();
    ImGui_ImplGlfw_NewFrame();
    ImGui::NewFrame();

    imGuiRenderBallsControlsFrame();
    imGuiRenderCameraControlsFrame();
    imGuiBallsThrowingStateControlsFrame();
    imGuiBallsMetricFrame();
    imGuiRenderBoundingBoxFrame();

    ImGui::Render();
    ImGui_ImplOpenGL3_RenderDrawData(ImGui::GetDrawData());
}

void cleanUpImGui(){
    ImGui_ImplOpenGL3_Shutdown();
    ImGui_ImplGlfw_Shutdown();
    ImGui::DestroyContext();
}

void compileShaders(){
    Shader shader("circle_vs.glsl", "circle_fs.glsl");
    shaderProgram = shader.programID();
}

// Funkcja do integracji Rungego-Kutty
void rungeKuttaStep(Ball &ball, float dt) {
    float k1vx = -airResistance * ball.vx;
    float k1vy = -gravity - airResistance * ball.vy;
    float k2vx = -airResistance * (ball.vx + ballRadius * dt * k1vx);
    float k2vy = -gravity - airResistance * (ball.vy + ballRadius * dt * k1vy);
    float k3vx = -airResistance * (ball.vx + ballRadius * dt * k2vx);
    float k3vy = -gravity - airResistance * (ball.vy + ballRadius * dt * k2vy);
    float k4vx = -airResistance * (ball.vx + dt * k3vx);
    float k4vy = -gravity - airResistance * (ball.vy + dt * k3vy);

    ball.vx += (dt / 6.0f) * (k1vx + 2.0f * k2vx + 2.0f * k3vx + k4vx);
    ball.vy += (dt / 6.0f) * (k1vy + 2.0f * k2vy + 2.0f * k3vy + k4vy);
    ball.x += ball.vx * dt;
    ball.y += ball.vy * dt;
}

void setupSphereBuffer() {

    std::vector<glm::vec3> vertices;

    float angle = 360.0f / numberOfPointsOnCircle;

    int triangleCount = numberOfPointsOnCircle - 2;

    std::vector<glm::vec3> temp;
    // positions
    for (int i = 0; i < numberOfPointsOnCircle; i++)
    {
        float currentAngle = angle * i;
        float x = ballRadius * cos(glm::radians(currentAngle));
        float y = ballRadius * sin(glm::radians(currentAngle));
        float z = 0.0f;

        temp.push_back(glm::vec3(x, y, z));
    }

    for (int i = 0; i < triangleCount; i++)
    {
        vertices.push_back(temp[0]);
        vertices.push_back(temp[i + 1]);
        vertices.push_back(temp[i + 2]);
    }

    glGenVertexArrays(1, &VAO_SPHERES);
    glGenBuffers(1, &VBO_SPHERES);

    glBindVertexArray(VAO_SPHERES);

    glBindBuffer(GL_ARRAY_BUFFER, VBO_SPHERES);
    glBufferData(GL_ARRAY_BUFFER, vertices.size() * sizeof(glm::vec3), vertices.data(), GL_STATIC_DRAW);

    // Współrzędne wierzchołków
    glVertexAttribPointer(0, 3, GL_FLOAT, GL_FALSE, 6 * sizeof(float), (void*)0);
    glEnableVertexAttribArray(0);

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

    glBindBuffer(GL_ARRAY_BUFFER, 0);
    glBindVertexArray(0);
}


void drawSphere() {
    glBindVertexArray(VAO_SPHERES);
    glDrawArrays(GL_TRIANGLE_STRIP, 0, (20 + 1) * (20 + 1));
    glBindVertexArray(0);
}

void display() {
    glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT);
    glUseProgram(shaderProgram);

    // Ustawienie widoku kamery
    glm::mat4 view = glm::lookAt(cameraPos, cameraTarget, cameraUp);
    glm::mat4 projection = glm::perspective(glm::radians(45.0f), (float)windowWidth / (float)windowHeight, 0.1f, 100.0f);

    // Ustawienie macierzy widoku i projekcji w shaderze
    GLint viewLoc = glGetUniformLocation(shaderProgram, "view");
    GLint projLoc = glGetUniformLocation(shaderProgram, "projection");
    glUniformMatrix4fv(viewLoc, 1, GL_FALSE, glm::value_ptr(view));
    glUniformMatrix4fv(projLoc, 1, GL_FALSE, glm::value_ptr(projection));

    // Renderowanie osi
    drawAxes();

    // Renderowanie każdej z ballsAmount kul
    for (int i = 0; i < ballsAmount; ++i) {
        Ball &ball = balls[i];  // Uzyskanie dostępu do danej kuli
        glm::mat4 model = glm::translate(glm::mat4(1.0f), glm::vec3(ball.x, ball.y, ball.z));

        // Ustawienie macierzy modelu w shaderze
        GLint modelLoc = glGetUniformLocation(shaderProgram, "model");
        glUniformMatrix4fv(modelLoc, 1, GL_FALSE, glm::value_ptr(model));

        // Ustawienie koloru kuli w shaderze
        GLint colorLoc = glGetUniformLocation(shaderProgram, "color");
        glUniform3f(colorLoc, ball.r, ball.g, ball.b);

        drawSphere();  // Rysowanie kuli
    }

    drawSquare();

    renderImGui();
    glfwSwapBuffers(glfwGetCurrentContext());  // Przełączanie buforów
}



void update() {
    float dt = 0.01f;

    if (isLaunched) {
        for (auto &ball : balls) {
            rungeKuttaStep(ball, dt);
        }
        computeBallsCollision();
    }
}

void setupOpenGL() {
    compileShaders();
    setupSphereBuffer();
    setupAxesBuffer();
    setupSquareBuffer(); // Inicjalizowanie buforów kwadratu

    glEnable(GL_DEPTH_TEST);
    glClearColor(0.1f, 0.1f, 0.1f, 1.0f);
}



void keyCallback(GLFWwindow* window, int key, int scancode, int action, int mods) {
    if (action == GLFW_PRESS) {
        switch (key) {
            case GLFW_KEY_ESCAPE:
                glfwSetWindowShouldClose(window, GLFW_TRUE);
                break;
        }
    }
}

int main() {
    if (!glfwInit()) {
        std::cerr << "Failed to initialize GLFW" << std::endl;
        return -1;
    }

    glfwWindowHint(GLFW_SAMPLES, 4);
    glfwWindowHint(GLFW_CONTEXT_VERSION_MAJOR, 3);
    glfwWindowHint(GLFW_CONTEXT_VERSION_MINOR, 3);
    glfwWindowHint(GLFW_OPENGL_PROFILE, GLFW_OPENGL_CORE_PROFILE);

    window = glfwCreateWindow(windowWidth, windowHeight, "Animacja 10 kul - Rzut ukośny", nullptr, nullptr);
    if (!window) {
        std::cerr << "Failed to create GLFW window" << std::endl;
        glfwTerminate();
        return -1;
    }

    glfwMakeContextCurrent(window);
    glfwSetKeyCallback(window, keyCallback); // Rejestracja callbacku klawiatury

    glewExperimental = GL_TRUE;
    if (glewInit() != GLEW_OK) {
        std::cerr << "Failed to initialize GLEW" << std::endl;
        return -1;
    }

    initImGui();
    initializeBalls();
    setupOpenGL();

    while (!glfwWindowShouldClose(window)) {
        display();
        update();
        glfwPollEvents();
    }

    cleanUpImGui();

    glDeleteVertexArrays(1, &VAO_SPHERES);
    glDeleteBuffers(1, &VBO_SPHERES);
    glDeleteProgram(shaderProgram);

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