#include <edit/scene.hpp>

#include <algo/matrix.hpp>
#include <algo/raycast.hpp>
#include <gfx/material.hpp>
#include <gfx/shapes.hpp>

#include <iostream>
#include <ranges>

Scene::Scene()
{
    createScene();
    findCameras(scene);
    findLights(scene);
}

void Scene::createScene()
{
    using namespace shapes;

    /* ADD LIGHT CUBE AND ITS LIGHT */
    Mesh light_cube(lit_cube_vertices, 
                    std::vector<float>{},
                    lit_cube_indices, 
                    glm::vec3(1.0f, 1.0f, 1.0f), 
                    "basic");
    auto sun_node = addNode(glm::vec3(-3.0f, 5.0f, 5.0f));
    sun_node->setScale(glm::vec3(0.1f, 0.1f, 0.1f));
    addMesh(sun_node, light_cube);
    addLight(sun_node, glm::vec4(1.0f), glm::vec4(0.7f), glm::vec4(0.7f), false);

    auto camera_light_node = Node::create(glm::vec3(1.0f, 1.0f, 1.0f));
    camera_light_node->setScale(glm::vec3(0.1f, 0.1f, 0.1f));
    addLight(camera_light_node, glm::vec4(0.0f), glm::vec4(0.5f), glm::vec4(0.0f), false, true);

    const std::vector<std::filesystem::path> skybox_paths = {
        "./data/images/kloofendal_cloudy/right.png",        
        "./data/images/kloofendal_cloudy/left.png",
        "./data/images/kloofendal_cloudy/top.png",
        "./data/images/kloofendal_cloudy/bottom.png",
        "./data/images/kloofendal_cloudy/front.png",
        "./data/images/kloofendal_cloudy/back.png"
        };

    addSkyboxNode("./data/models/cubemap_cube.obj", skybox_paths);

    /* ADD GRID */
    mesh_ptr grid_mesh = std::make_shared<Mesh>(shapes::generate_grid_vertices(100), 
                                               std::vector<float>{},
                                               shapes::generate_grid_indices(100), 
                                               glm::vec3(0.5f, 0.5f, 0.5f), 
                                               "basic");
    grid_mesh->setDrawMode(GL_LINES);     
    grid_mesh->setSelectability(false);   

    grid = addNode(Node::create());
    addMesh(scene.back(), std::move(grid_mesh));

    /* ADD CAMERA */
    auto camera_node = addCameraNode(60.0f, glm::vec3(0.0f, 0.0f, 5.0f), glm::vec3(0.0f, 0.0f, 0.0f));
    auto rot = glm::angleAxis(glm::radians(-45.0f), glm::vec3(0.0f, 1.0f, 0.0f));
    camera_node->translatePosVec(glm::toMat4(rot));
    camera_node->rotateRotationQuat(rot);
    auto rot2 = glm::angleAxis(glm::radians(-20.0f), camera_node->getRotationAxisXLocal());
    camera_node->translatePosVec(glm::toMat4(rot2));
    camera_node->rotateRotationQuat(rot2);
    camera_node->addChild(camera_light_node);
}

std::vector<mesh_ptr> Scene::getNodeMesh(node_ptr node) const
{
    std::vector<mesh_ptr> result;

    for (auto &object : node->getObjects())
    {
        auto mesh = std::dynamic_pointer_cast<Mesh>(object);
        if (mesh)
            result.emplace_back(mesh);
    }
    return result;
}

void Scene::removeActiveCamera(camera_ptr camera)
{
    std::erase_if(active_cameras,
                  [&camera](const camera_weak_ptr &cam) { return cam.lock() == camera; });
}

node_ptr Scene::addNode(node_ptr node) 
{
    scene.emplace_back(node);
    return scene.back();
}

node_ptr Scene::addNode(glm::vec3 position, glm::vec3 rotation, glm::vec3 scale) 
{
    scene.emplace_back(Node::create(position, rotation, scale));
    return scene.back();
}

void Scene::removeNode(node_ptr to_remove)
{
    removeNodeRec(to_remove, scene);
}

void Scene::removeNodeRec(node_ptr to_remove, std::vector<node_ptr> &where)
{
    auto it = std::find(where.begin(), where.end(), to_remove);
    if (it != where.end())
        where.erase(it);
    
    for (auto &node : where)
    {
        removeNodeRec(to_remove, node->getChildren());
    }
}

node_ptr Scene::addMeshNode(const std::string &file_name, node_ptr node)
{
    scene.emplace_back(node);
    scene.back()->addObject(Mesh::loadMesh(file_name));
    return scene.back();
}

node_ptr Scene::addMeshNode(const std::string &file_name, glm::vec3 position)
{
    scene.emplace_back(Node::create(position));
    scene.back()->addObject(Mesh::loadMesh(file_name));
    return scene.back();
}

node_ptr Scene::addLightNode(glm::vec3 color, node_ptr node)
{
    scene.emplace_back(node);
    scene.back()->addObject(std::make_shared<Light>(color, scene.back()));
    return scene.back();
}

node_ptr Scene::addLightNode(glm::vec3 color, glm::vec3 position)
{
    scene.emplace_back(Node::create(position));
    scene.back()->addObject(std::make_shared<Light>(color, scene.back()));
    return scene.back();
}

node_ptr Scene::addCameraNode(float FOV, node_ptr node)
{
    scene.emplace_back(node);
    scene.back()->addObject(std::make_shared<Camera>(FOV, scene.back()));
    return scene.back();
}


node_ptr Scene::addCameraNode(float FOV, glm::vec3 position, glm::vec3 rotation)
{
    scene.emplace_back(Node::create(position, rotation));
    scene.back()->addObject(std::make_shared<Camera>(FOV, scene.back()));
    return scene.back();
}

mesh_ptr Scene::addMesh(node_ptr node, const std::string &file_name, const std::string &texture_name)
{
    auto mesh = Mesh::loadMesh(file_name, texture_name);
    node->addObject(mesh);
    return mesh;
}

void Scene::addMesh(node_ptr node, mesh_ptr mesh)
{
    node->addObject(std::move(mesh));
}

void Scene::addMesh(node_ptr node, Mesh mesh)
{
    node->addObject(std::make_shared<Mesh>(mesh));
}

void Scene::addMesh(node_ptr node, std::vector<float> vertices, std::vector<float> normals, std::vector<unsigned int> indices)
{
    node->addObject(std::make_shared<Mesh>(Mesh{vertices, normals, indices}));
}

void Scene::removeMesh(node_ptr node, mesh_ptr mesh)
{
    std::erase_if(node->getObjects(),
                  [&mesh](const objectbase_ptr &object) { return std::dynamic_pointer_cast<Mesh>(object) == mesh; });
}

void Scene::removeAllMesh(node_ptr node)
{
    std::erase_if(node->getObjects(),
                  [](const objectbase_ptr &object) { return std::dynamic_pointer_cast<Mesh>(object); });
}

void Scene::addSkyboxNode(const std::string &file_name, const std::vector<std::filesystem::path> &texture_names)
{
    skybox = addNode(glm::vec3(0.0f));
    auto skybox_mesh = Mesh::loadMesh(file_name, texture_names);
    skybox_mesh->setShaderType("skybox");
    skybox_mesh->setSelectability(false);
    skybox->addObject(std::move(skybox_mesh));
}

void Scene::addLight(node_ptr node, light_ptr light)
{
    node->addObject(std::move(light));
}

void Scene::addLight(node_ptr node, Light light)
{
    node->addObject(std::make_shared<Light>(light));
}

void Scene::addLight(node_ptr node, glm::vec3 color, bool is_point)
{
    node->addObject(std::make_shared<Light>(Light{color, node, is_point}));
}

void Scene::addLight(node_ptr node, glm::vec4 ambient_color, glm::vec4 diffuse_color, glm::vec4 specular_color, bool is_point, bool use_rotation)
{
    node->addObject(std::make_shared<Light>(Light{node, ambient_color, diffuse_color, specular_color, is_point, use_rotation}));
}

void Scene::addLight(node_ptr node, glm::vec4 ambient_color, glm::vec4 diffuse_color, glm::vec4 specular_color, glm::vec3 cone_direction)
{
    node->addObject(std::make_shared<Light>(Light{node, ambient_color, diffuse_color, specular_color, cone_direction}));
}

void Scene::addCamera(node_ptr node, camera_ptr camera)
{
    node->addObject(std::move(camera));
}

void Scene::addCamera(node_ptr node, Camera camera)
{
    node->addObject(std::make_shared<Camera>(camera));
}

void Scene::addCamera(node_ptr node, float FOV)
{
    node->addObject(std::make_shared<Camera>(Camera{FOV, node}));
}
void Scene::addRay(glm::vec3 ray_position, glm::vec3 ray_direction)
{
    mesh_ptr ray = std::make_shared<Mesh>(Mesh{{0, 0, 0, ray_direction.x, ray_direction.y, ray_direction.z},
                                                     {}, 
                                                     {0, 1}, 
                                                     glm::vec3(0.5f, 1, 0.31f), "basic"});
    ray->setDrawMode(GL_LINES);
    ray->setSelectability(false);

    addNode(ray_position);
    addMesh(scene.back(), ray);
}

// TO DO: following two functions almost identical - fix

void Scene::addSelectionAABB(node_ptr node, std::size_t index)
{
    ASSUMPTION(index <= selection_AABBs.size());

    if (index >= selection_AABBs.size())
        selection_AABBs.emplace_back(Node::create());
    else
        selection_AABBs[index] = Node::create();

    for (auto &object : node->getObjects())
    {
        auto obj = dynamic_pointer_cast<Mesh>(object);
        if (!obj)
            continue;
        
        auto color = (index == 0) || node == selections[0] ? glm::vec3(0.5f, 1, 0.31f) : glm::vec3(1.0f, 0.5f, 0.31f);
        Mesh obj_bbox(obj->getAABB().getVertices(), std::vector<float>{}, obj->getAABB().getIndices(), 
                      color, "basic");
        obj_bbox.setSelectability(false);
        obj_bbox.setDrawMode(GL_LINES);
        
        addMesh(selection_AABBs[index], std::move(obj_bbox));
    }

    node->addChild(selection_AABBs[index]);
}

void Scene::addPreAddAABB(node_ptr node)
{
    auto new_node = Node::create(node->getPositionLocal(), node->getRotationLocal(), node->getScaleLocal());
    addNode(new_node);
    pre_add_AABB = scene.back();
    for (auto &object : node->getObjects())
    {
        auto obj = dynamic_pointer_cast<Mesh>(object);
        if (!obj)
            continue;
            
        Mesh obj_bbox(obj->getAABB().getVertices(), std::vector<float>{}, obj->getAABB().getIndices(), 
                      glm::vec3(0.5f, 1, 0.31f), "basic");
        obj_bbox.setSelectability(false);
        obj_bbox.setDrawMode(GL_LINES);

        addMesh(new_node, std::move(obj_bbox));
    }
}
void Scene::addSideHighlight(node_ptr node, glm::vec3 side_direction, std::array<glm::vec3, 4> &side_vertices, std::size_t index)
{
    ASSUMPTION(index <= side_highlights.size());

    if (index >= side_highlights.size())
        side_highlights.emplace_back(Node::create());
    else
        side_highlights[index] = Node::create();

    node->addChild(side_highlights[index]);

    std::vector<float> vertices;
    for (auto v : side_vertices)
    {
        v += 0.01f * side_direction; // moves highlight little bit infront
        vertices.emplace_back(v.x);
        vertices.emplace_back(v.y);
        vertices.emplace_back(v.z);
    }

    std::vector<float> normals;
    std::vector<unsigned int> indices = { 0, 1, 2, 1, 2, 3 };

    auto color = (index == 0) || node == selections[0] ? glm::vec3(0.0f, 1.0f, 0.0f) : glm::vec3(1.0f, 0.5f, 0.31f);
    Mesh m(vertices, std::vector<float>{}, indices, color, "basic");
    m.setSelectability(false);
    m.setMaterial(mat::green_highlight);

    addMesh(side_highlights[index], std::move(m));
}

void Scene::addRotationAxis(glm::vec3 axis_direction, node_ptr selected_node)
{
    auto rotation_axis = addNode(Node::create(selected_node->getPositionWorld(), selected_node->getRotationWorld(), selected_node->getScaleWorld()));
    auto fst = axis_direction;
    auto snd = -axis_direction;
    std::vector<float> vertices = {fst.x, fst.y, fst.z, snd.x, snd.y, snd.z};
    std::vector<unsigned int> indices = { 0, 1 };
    Mesh m(vertices, std::vector<float>{}, indices, glm::vec3(0.0f, 0.0f, 1.0f), "basic");
    m.setSelectability(false);
    m.setDrawMode(GL_LINES);

    addMesh(rotation_axis, std::move(m));
}

void Scene::addObjectAxes(node_ptr node)
{
    object_axes = Node::create();
    node->addChild(object_axes);

    // FTO DO: add scaling given the object's AABB(LH)
    mesh_ptr debug_x_axis = std::make_shared<Mesh>(shapes::axis_x_vertices, std::vector<float>{}, shapes::axis_indices);
    mesh_ptr debug_y_axis = std::make_shared<Mesh>(shapes::axis_y_vertices, std::vector<float>{}, shapes::axis_indices);
    mesh_ptr debug_z_axis = std::make_shared<Mesh>(shapes::axis_z_vertices, std::vector<float>{}, shapes::axis_indices);

    debug_x_axis->setMaterial(Material(color::red));
    debug_y_axis->setMaterial(Material(color::green));
    debug_z_axis->setMaterial(Material(color::blue));

    debug_x_axis->setShaderType("basic");
    debug_y_axis->setShaderType("basic");
    debug_z_axis->setShaderType("basic");

    debug_x_axis->setDrawMode(GL_LINES);
    debug_y_axis->setDrawMode(GL_LINES);
    debug_z_axis->setDrawMode(GL_LINES);

    object_axes->addObject(debug_x_axis);
    object_axes->addObject(debug_y_axis);
    object_axes->addObject(debug_z_axis);
}

void Scene::createCameraAxes()
{
    camera_axes = Node::create();

    mesh_ptr rofi_x_axis = std::make_shared<Mesh>(shapes::rofi_axis_x_vertices, std::vector<float>{}, shapes::rofi_axis_indices);
    mesh_ptr rofi_y_axis = std::make_shared<Mesh>(shapes::rofi_axis_y_vertices, std::vector<float>{}, shapes::rofi_axis_indices);
    mesh_ptr rofi_z_axis = std::make_shared<Mesh>(shapes::rofi_axis_z_vertices, std::vector<float>{}, shapes::rofi_axis_indices);

    rofi_x_axis->setMaterial(Material(color::red));
    rofi_y_axis->setMaterial(Material(color::green));
    rofi_z_axis->setMaterial(Material(color::blue));

    rofi_x_axis->setShaderType("interface");
    rofi_y_axis->setShaderType("interface");
    rofi_z_axis->setShaderType("interface");

    rofi_x_axis->setDrawMode(GL_LINES);
    rofi_y_axis->setDrawMode(GL_LINES);
    rofi_z_axis->setDrawMode(GL_LINES);

    // ordered due to disabled depth test
    camera_axes->addObject(rofi_z_axis);
    camera_axes->addObject(rofi_x_axis);
    camera_axes->addObject(rofi_y_axis);

    ASSUMPTION(!active_cameras.empty());
    active_cameras.front().lock()->getNode()->addChild(camera_axes);
}

void Scene::createCameraAxesDebug()
{
     camera_axes_debug = Node::create();

    mesh_ptr debug_x_axis = std::make_shared<Mesh>(shapes::debug_axis_x_vertices, std::vector<float>{}, shapes::debug_axis_indices);
    mesh_ptr debug_y_axis = std::make_shared<Mesh>(shapes::debug_axis_y_vertices, std::vector<float>{}, shapes::debug_axis_indices);
    mesh_ptr debug_z_axis = std::make_shared<Mesh>(shapes::debug_axis_z_vertices, std::vector<float>{}, shapes::debug_axis_indices);

    debug_x_axis->setMaterial(Material(color::red));
    debug_y_axis->setMaterial(Material(color::green));
    debug_z_axis->setMaterial(Material(color::blue));

    debug_x_axis->setShaderType("interface");
    debug_y_axis->setShaderType("interface");
    debug_z_axis->setShaderType("interface");

    debug_x_axis->setDrawMode(GL_LINES);
    debug_y_axis->setDrawMode(GL_LINES);
    debug_z_axis->setDrawMode(GL_LINES);

    // ordered due to disabled depth test
    camera_axes_debug->addObject(debug_y_axis);
    camera_axes_debug->addObject(debug_z_axis);
    camera_axes_debug->addObject(debug_x_axis);

    ASSUMPTION(!active_cameras.empty());
    active_cameras.front().lock()->getNode()->addChild(camera_axes_debug);
}

bool Scene::containsMesh(node_ptr node, mesh_ptr mesh)
{
    auto node_mesh = getNodeMesh(node);
    return std::ranges::find(node_mesh, mesh) != node_mesh.end();
}

void Scene::findLights(const std::vector<node_ptr> &nodes)
{
    lights.clear();
    findLightsRec(nodes);
}

void Scene::findLightsRec(const std::vector<node_ptr> &nodes)
{
    for (auto &node : nodes)
    {
        for (auto &object : node->getObjects())
        {
            auto light = dynamic_pointer_cast<Light>(object);
            if (light)
                lights.emplace_back(light);
        } 
        findLightsRec(node->getChildren());
    }
}

void Scene::findCameras(const std::vector<node_ptr> &nodes)
{
    for (auto &node : nodes)
    {
        for (auto &object : node->getObjects())
        {
            auto camera = dynamic_pointer_cast<Camera>(object);
            if (camera)
                cameras.emplace_back(camera);
        } 
        findCameras(node->getChildren());
    }
}
void Scene::changeDrawMode(node_ptr node, unsigned int mode)
{
    for (auto &object : node->getObjects())
    {
        auto obj = dynamic_pointer_cast<Mesh>(object);
        if (!obj)
            continue;
        obj->setDrawMode(mode);
    }
}

void Scene::manageSelection(node_ptr nearest_node)
{
    removeNode(selection_AABBs[0]);

    if (!nearest_node)
    {
        // Uncomment for Debug
        // manageObjectAxes(nullptr);
        selections[0] = nullptr;
        return;
    }
    selections[0] = nearest_node;
    addSelectionAABB(nearest_node);
    // Uncomment for Debug
    // manageObjectAxes(selections[0]);
}

void Scene::manageHighlight(glm::vec3 hit_coord, node_ptr nearest_node)
{
    removeNode(side_highlights[0]);
    if (!nearest_node || nearest_node != selections[0])
    {
        side_highlights[0] = nullptr;
        return;
    }

    auto [side_direction, side_vertices] = find_cube_side(selection_AABBs[0], hit_coord);    

    highlight_directions[0] = side_direction;
    addSideHighlight(nearest_node, side_direction, side_vertices);
}

void Scene::manageAxis(glm::vec3 axis_direction, node_ptr selected_node)
{
    removeNode(rotation_axis);
    rotation_axis = nullptr;

    if (!selected_node)
        return;

    addRotationAxis(axis_direction, selected_node);

}

void Scene::manageObjectAxes(node_ptr selected_node)
{
    removeNode(object_axes);
    object_axes = nullptr;

    if(!selected_node)
        return;

    addObjectAxes(selected_node);
}

void Scene::manageSecondarySelection(std::vector<node_ptr> nearest_nodes)
{
    for (std::size_t index = 1; index < selection_AABBs.size(); ++index)
    {
        removeNode(selection_AABBs[index]);
        selection_AABBs[index] = nullptr;
    }

    for (std::size_t index = 1; index < selections.size(); ++index)
        selections[index] = nullptr;

    for(std::size_t index = 0; index < nearest_nodes.size(); ++index)
    {
        if (!nearest_nodes[index])
            continue;

        if (index + 1 >= selections.size())
            selections.emplace_back(nearest_nodes[index]);
        else
            selections[index + 1] = nearest_nodes[index];
        
        addSelectionAABB(nearest_nodes[index], index + 1);
    }
}

void Scene::manageSecondaryHighlights(std::vector<std::pair<glm::vec3, node_ptr>> hit_node_pairs)
{
    for (std::size_t index = 1; index < side_highlights.size(); ++index)
    {
        removeNode(side_highlights[index]);
        side_highlights[index] = nullptr;
    }

    for (std::size_t index = 1; index < highlight_directions.size(); ++index)
        highlight_directions[index] = glm::vec3(0.0f);

    for(std::size_t index = 0; index < hit_node_pairs.size(); ++index)
    {
        auto &[hit_coord, node] = hit_node_pairs[index];
        if (!node)
            continue;

        auto [side_direction, side_vertices] = find_cube_side(selection_AABBs[index + 1], hit_coord);

        if (index + 1 >= highlight_directions.size())
            highlight_directions.emplace_back(side_direction);
        else
            highlight_directions[index + 1] = side_direction;

        addSideHighlight(node, side_direction, side_vertices, index + 1);
    }
}

void Scene::manageCameraAxes(bool enabled)
{
    if (!enabled)
    {
        removeNode(camera_axes);
    }
    else
    {
        if (!camera_axes)
            createCameraAxes();

        addNode(camera_axes);
    }
    camera_axes_enabled = enabled;

    // Uncomment for Debug
    // manageCameraAxesDebug(enabled);
}

void Scene::manageCameraAxesDebug(bool enabled)
{
    if (!enabled)
    {
        removeNode(camera_axes_debug);
    }
    else
    {
        if (!camera_axes_debug)
            createCameraAxesDebug();

        addNode(camera_axes_debug);
    }
}

void Scene::updateCameraAxesPosition()
{
    if (!camera_axes_enabled)
        return;

    ASSUMPTION(!active_cameras.empty());

    const auto active_cam = active_cameras.front().lock();
    const auto window_size = active_cam->getWindowSize();
    const auto ratio = window_size.x / window_size.y; 
    const auto screen_offset = 0.06f;
    const auto x_compensation = 0.01f;
    camera_axes->setPosVec(glm::vec3(screen_offset * ratio + (x_compensation * ratio), screen_offset, -0.15f));

    const auto camera_rotation = active_cameras.front().lock()->getNode()->getRotationMatWorld();
    camera_axes->setRotationQuat(glm::inverse(camera_rotation));

    // Uncomment for Debug
    // updateCameraAxesDebugPosition();
}

void Scene::updateCameraAxesDebugPosition()
{
    if (!camera_axes_enabled)
        return;

    ASSUMPTION(!active_cameras.empty());

    const auto active_cam = active_cameras.front().lock();
    const auto window_size = active_cam->getWindowSize();
    const auto ratio = window_size.x / window_size.y; 
    const auto screen_offset = 0.06f;
    const auto x_compensation = -0.01f;
    camera_axes_debug->setPosVec(glm::vec3(screen_offset * ratio + (x_compensation * ratio), screen_offset, -0.15f));

    const auto camera_rotation = active_cameras.front().lock()->getNode()->getRotationMatWorld();
    camera_axes_debug->setRotationQuat(glm::inverse(camera_rotation));
}

void Scene::toggleVisibility(node_ptr node, bool visible)
{
    if (!node)
        return;
        
    for (auto &object : node->getObjects())
    {
        auto mesh = dynamic_pointer_cast<Mesh>(object);
        if (!mesh)
            continue;
        mesh->setVisible(visible);
    }
}
