#include <iostream>
#include <vector>

#include "lib/glad/glad.c"
#include <GLFW/glfw3.h>

#define GLM_ENABLE_EXPERIMENTAL
#include <glm/ext/matrix_transform.hpp>
#include <glm/glm.hpp>
#include <glm/gtx/quaternion.hpp>
#include <glm/gtc/type_ptr.hpp>
#include <glm/gtc/matrix_transform.hpp>

#include "gfx/gfx.cpp"
#include "VoxelSpace.cpp"
#include "SomaSolve.cpp"
#include "lib/djstdlib/core.cpp"

#define STB_IMAGE_IMPLEMENTATION
#include "lib/loaders/stb_image.h"

#include "./tests.cpp"

struct Entity;
struct Polycube;
struct SceneGraphNode;
uint32 new_entity();
Entity *get_entity(int id);
SceneGraphNode *get_scene_graph_node(int id); 
uint32 new_graph_node();

void print_mat(glm::mat4* matrix) {
    glm::mat4 mat = *matrix;
    std::cout << mat[0][0] << mat[0][1] << mat[0][2] << mat[0][3] << std::endl;
    std::cout << mat[1][0] << mat[1][1] << mat[1][2] << mat[1][3] << std::endl;
    std::cout << mat[2][0] << mat[2][1] << mat[2][2] << mat[2][3] << std::endl;
    std::cout << mat[3][0] << mat[3][1] << mat[3][2] << mat[3][3] << std::endl;
    std::cout << std::endl;
}

struct Camera {
    glm::mat4 view;
    glm::mat4 proj;
    glm::vec3 pos;
    glm::vec3 up;
    glm::vec3 target;
};

Camera *createCamera(Arena *arena, real32 aspect_ratio = 800.0f / 600.0f) {
    Camera *result = PushStruct(arena, Camera);
    result->view = glm::mat4();
    result->proj = glm::perspective(glm::radians(45.0f), aspect_ratio, 0.1f, 100.0f);
    result->pos = glm::vec3(0.0f);
    result->up = glm::vec3(0.0f, 1.0f, 0.0f);
    return result;
}

void camera_look_at(Camera *c, float x, float y, float z) {
    c->target = glm::vec3(x, y, z);
    c->view = glm::lookAt(c->pos, c->target, c->up);
}

void camera_set_up(Camera *c, real32 up_x, real32 up_y, real32 up_z) {
    c->up = glm::vec3(up_x, up_y, up_z);
}

struct GlobalAppState {
    uint32 current_polycube;
    uint32 last_polycube_visible;
    Shader *active_shader;
    std::vector<Polycube> polycubes;
};

GlobalAppState app_state = {};

struct WindowDims {
    uint32 width;
    uint32 height;
};

struct Entity {
    Mesh *mesh;
    Texture *tex;
    bool visible;
    uint32 scene_graph_node;
};

struct SceneGraphNode {
    glm::mat4 local;
    glm::mat4 world;
    glm::vec3 translation;
    glm::quat rotation;
    glm::vec3 scale;
    std::vector<uint32> children;
    uint32 entity;
};

void init_sg_node(SceneGraphNode *n) {
    n->scale = glm::vec3(1.0f, 1.0f, 1.0f);
    n->translation = glm::vec3(0.0f, 0.0f, 0.0f);
    n->rotation = glm::quat(0.0f, 0.0f, 0.0f, 0.0f);
    n->local = glm::mat4(1.0f);
    n->world = n->local;
}

void recalculate_sg_node(SceneGraphNode *n) {
    n->local = glm::scale(
        glm::translate(
            glm::mat4(1.0f), 
            n->translation
        ) * glm::toMat4(n->rotation), 
        n->scale
    );
}

struct Polycube {
    int graph_node;
    glm::vec3 color;
};

void show_polycube(Polycube *p) {
    SceneGraphNode *node = get_scene_graph_node(p->graph_node);
    for (uint32 &child : node->children) {
        SceneGraphNode *subNode = get_scene_graph_node(child);
        if (subNode->entity) {
            get_entity(subNode->entity)->visible = true;
        }
    }
}

void hide_polycube(Polycube *p) {
    SceneGraphNode *node = get_scene_graph_node(p->graph_node);
    for (uint32 &child : node->children) {
        SceneGraphNode *subNode = get_scene_graph_node(child);
        if (subNode->entity) {
            get_entity(subNode->entity)->visible = false;
        }
    }
}

glm::vec3 centreFromPolycube(Polycube *p) {
    glm::vec3 centre = glm::vec3(0.0f);
    for (uint32 &child : get_scene_graph_node(p->graph_node)->children) {
        centre += get_scene_graph_node(child)->translation;
    }
    centre /= get_scene_graph_node(p->graph_node)->children.size();
    return centre;
}

struct Frame {
    uint32 width;
    uint32 height;
    int32 x;
    int32 y;
    Camera* cam;
};

Frame createFrame(Arena *arena, uint32 width, uint32 height, uint32 x, uint32 y) {
    Frame result = {0};
    result.width = width;
    result.height = height;
    result.x = x;
    result.y = y;
    result.cam = createCamera(arena, (real32)result.width / (real32)result.height);
    return result;
}

void framebuffer_size_callback(GLFWwindow *window, int width, int height) {
    glViewport(0, 0, width, height);
}

GLFWwindow *init_window_and_gl(WindowDims *window_dims) {
    glfwInit();
    glfwWindowHint(GLFW_CONTEXT_VERSION_MAJOR, 4);
    glfwWindowHint(GLFW_CONTEXT_VERSION_MINOR, 6);
    glfwWindowHint(GLFW_OPENGL_PROFILE, GLFW_OPENGL_CORE_PROFILE);
    GLFWwindow *window = glfwCreateWindow(window_dims->width, window_dims->height, "Somaesque", NULL, NULL);
    if (window == NULL) {
        std::cout << "Failed to create GLFW window" << std::endl;
        glfwTerminate();
        return nullptr;
    }
    glfwMakeContextCurrent(window);

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

    glViewport(0, 0, 800, 600);
    glfwSetFramebufferSizeCallback(window, framebuffer_size_callback);
    glEnable(GL_DEPTH_TEST);
    return window;
}

void gl_update_viewport(WindowDims* window_dims, Frame* frame) {
    glViewport(frame->x, window_dims->height - frame->y - frame->height, frame->width, frame->height);
}

Mesh cube_mesh = {0};
Texture wall_tex = {0};
std::vector<Entity> entities = std::vector<Entity>();
std::vector<SceneGraphNode> scene_graph_nodes = std::vector<SceneGraphNode>();

void process_input(GLFWwindow *window) {
    local_persist bool wireframe = false;
    local_persist bool last_frame_state_press_enter = false;
    local_persist bool last_frame_state_press = false;

    if (glfwGetKey(window, GLFW_KEY_ESCAPE) == GLFW_PRESS) {
        glfwSetWindowShouldClose(window, true);
    } 

    if (glfwGetKey(window, GLFW_KEY_SPACE) == GLFW_PRESS && !last_frame_state_press) {
        glPolygonMode(GL_FRONT_AND_BACK, !wireframe ? GL_LINE : GL_FILL);
        wireframe = !wireframe;
        last_frame_state_press = true;
    } else if (glfwGetKey(window, GLFW_KEY_SPACE) == GLFW_RELEASE) {
        last_frame_state_press = false;
    } 

    if (glfwGetKey(window, GLFW_KEY_ENTER) == GLFW_PRESS && !last_frame_state_press_enter) {
        if (app_state.current_polycube == 6) {
            app_state.current_polycube = 0;
        } else {
            app_state.current_polycube += 1;
        }
        last_frame_state_press_enter = true;
    } else if (glfwGetKey(window, GLFW_KEY_ENTER) == GLFW_RELEASE) {
        last_frame_state_press_enter = false;
    }
}


uint32 new_entity() {
    entities.emplace_back();
    scene_graph_nodes.emplace_back();
    entities.back().scene_graph_node = (uint32)scene_graph_nodes.size();
    scene_graph_nodes.back().entity = (uint32)entities.size();
    return (uint32)entities.size();
}

Entity *get_entity(int id) {
    return &entities[id - 1];
}

SceneGraphNode *get_scene_graph_node(int id) {
    return &scene_graph_nodes[id - 1];
}

uint32 new_graph_node() {
    scene_graph_nodes.emplace_back();
    return (uint32)scene_graph_nodes.size();
}

Polycube create_polycube_from_repr(Space *repr) {
    uint32 polycube_id = new_graph_node();
    init_sg_node(get_scene_graph_node(polycube_id));
    for (int x = 0; x < repr->dim_x; x++) {
        for (int y = 0; y < repr->dim_y; y++) {
            for (int z = 0; z < repr->dim_z; z++) {
                if (filledAt(repr, x, y, z)) {
                    Entity *polycube_segment = get_entity(new_entity());
                    polycube_segment->mesh=&cube_mesh, 
                    polycube_segment->tex=&wall_tex;
                    SceneGraphNode *graph_node = get_scene_graph_node(polycube_segment->scene_graph_node);
                    init_sg_node(graph_node);
                    graph_node->translation = glm::vec3(
                        -((repr->dim_z - 1)/2.0f) + z,
                        ((repr->dim_x - 1)/2.0f) - x,
                        -((repr->dim_y - 1)/2.0f) + y
                    );
                    recalculate_sg_node(graph_node);
                    get_scene_graph_node(polycube_id)->children.push_back(polycube_segment->scene_graph_node);
                }
            }
        }
    }
    Polycube result = {};
    result.graph_node = polycube_id;
    result.color = glm::vec3(1.0f);
    return result;
}

void recalculate_scene_graph(SceneGraphNode *top) {
    if (top->children.size() == 0) {
        return;
    }
    for (uint32 &node_id : top->children) {
        SceneGraphNode *graph_node = get_scene_graph_node(node_id);
        recalculate_sg_node(graph_node);
        graph_node->world = top->world * graph_node->local;
        recalculate_scene_graph(graph_node);
    }
}

int main_cmd() {
    interactive_cmd_line_solve_soma();
    return 0;
}

int main_gfx() {
    Arena *arena = arenaAlloc(Megabytes(128));

    WindowDims window_dims = { 800, 600 };
    GLFWwindow *window = init_window_and_gl(&window_dims);
    if (!window) {
        return -1;
    }

    app_state.current_polycube = 0;
    app_state.last_polycube_visible = 6;
    app_state.active_shader = 0;
    app_state.polycubes = {};

    Shader phong_shader = createShader("../assets/shaders/phong-solid.vertex.glsl"_s, "../assets/shaders/phong-solid.fragment.glsl"_s);
    app_state.active_shader = &phong_shader;

    cube_mesh = createMesh("../assets/models/c000000.obj");
    wall_tex = createTexture("../assets/textures/brick-wall.jpg");

    Frame little_frame = createFrame(arena, 80, 60, 20, 20);
    Frame big_frame = createFrame(arena, 800, 600, 0, 0);
    Frame *frames[] = { &big_frame, &little_frame };

    SceneGraphNode root_node = {};
    init_sg_node(&root_node);

    for (int i = 0; i < ArrayCount(STD_SOMA); i++) {
        Space voxel_space = { STD_SOMA[i], 3, 3, 3 };
        cullEmptySpace(&voxel_space);
        Polycube polycube = create_polycube_from_repr(&voxel_space);
        polycube.color = color_from_index(i);
        app_state.polycubes.push_back(polycube);
        root_node.children.push_back(app_state.polycubes.back().graph_node);
    }

    big_frame.cam->pos = glm::vec3(0.0f, 4.0f, 4.0f);
    camera_look_at(big_frame.cam, 0.0f, 0.0f, 0.0f);

    glm::vec3 light_pos = glm::vec3(4.0f, 12.0f, 2.0f);

    glUseProgram(app_state.active_shader->prog_id);
    glUniform3fv(
            glGetUniformLocation(app_state.active_shader->prog_id, "light_pos"), 
            1, glm::value_ptr(light_pos));
    glUniformMatrix4fv(
            glGetUniformLocation(app_state.active_shader->prog_id, "projection"), 
            1, GL_FALSE, glm::value_ptr(big_frame.cam->proj)); 
    glUniformMatrix4fv(
            glGetUniformLocation(app_state.active_shader->prog_id, "view"), 
            1, GL_FALSE, glm::value_ptr(big_frame.cam->view));

    real64 last_frame = glfwGetTime();
    real64 time_delta = 1.0f/60.0f;
    while (!glfwWindowShouldClose(window)) {
        time_delta = glfwGetTime() - last_frame;
        process_input(window);

        if (app_state.last_polycube_visible != app_state.current_polycube) {
            hide_polycube(&app_state.polycubes[app_state.last_polycube_visible]);
            show_polycube(&app_state.polycubes[app_state.current_polycube]);
            app_state.last_polycube_visible = app_state.current_polycube;
        }

        glClearColor(0.0f, 0.0f, 0.0f, 1.0f);
        glClear(GL_DEPTH_BUFFER_BIT | GL_COLOR_BUFFER_BIT);

        gl_update_viewport(&window_dims, &big_frame);
        Polycube *current_polycube = &app_state.polycubes[app_state.current_polycube];
        get_scene_graph_node(current_polycube->graph_node)->rotation = glm::quat(glm::vec3(0, glfwGetTime() / 2, 0));

        glBindVertexArray(cube_mesh.vao);
        recalculate_scene_graph(&root_node);
        GLint model_uniform_loc = glGetUniformLocation(app_state.active_shader->prog_id, "model");
        glUniform3fv(
                glGetUniformLocation(app_state.active_shader->prog_id, "solid_color"), 
                1, glm::value_ptr(current_polycube->color));

        for (Entity &entity : entities) {
            if (entity.visible) {
                glUniformMatrix4fv(model_uniform_loc, 1, GL_FALSE, glm::value_ptr(get_scene_graph_node(entity.scene_graph_node)->world));
                glBindTexture(GL_TEXTURE_2D, entity.tex->tex_id);
                glDrawArrays(GL_TRIANGLES, 0, (GLsizei)entity.mesh->num_indices);
                //glDrawElements(GL_TRIANGLES, entity->mesh->num_indices, GL_UNSIGNED_INT, 0);
            }
        }

        glfwSwapBuffers(window);
        glfwPollEvents();
    }

    glfwTerminate();
    return 0;
}

int main() {
    initialiseCore();
    return main_gfx();
}