#include <rofi/rofiworld.hpp>

#include <algo/misc.hpp>
#include <rofi/connector.hpp>

#include <algorithm>
#include <ranges>

namespace rofi
{

const module_ptr RofiWorld::getModuleByID(const uint64_t id) const
{
    const auto it = std::find_if(modules_map.begin(), modules_map.end(), [&id] (const auto pair) { return pair.second == id; } );

    ASSUMPTION(it != modules_map.end());

    return it->first;
}

void RofiWorld::changeModuleID(module_ptr module, uint64_t new_id)
{
    ASSUMPTION(modules_map.contains(module));

    const auto old_id = modules_map[module];
    modules_map[module] = new_id;
    ids_set.erase(old_id);
    ids_set.emplace(new_id);
}

module_ptr RofiWorld::getModuleWithNode(const node_ptr node) const
{
    auto curr_node = node;
    while (!nodes_map.contains(curr_node))
    {
        auto parent_weak = curr_node->getParent();
        if (parent_weak.expired())
            return nullptr;

        curr_node = parent_weak.lock();
    }

    return nodes_map.at(curr_node);
}

std::vector<module_ptr> RofiWorld::getNeighborModules(module_ptr module) const
{
    std::vector<module_ptr> neighbor_modules;

    for (auto &connector : module->getParts()->getConnections())
    {
        if (!connector->hasModuleLink())
            continue;

        auto opposite_connector = connector->getModuleLink().lock()->getOppositeConnector(connector).lock();
        neighbor_modules.emplace_back(getModuleWithNode(opposite_connector->getNode()));
    }

    return neighbor_modules;
}

std::vector<node_ptr> RofiWorld::getDependentModulesNodes(node_ptr module_node) const
{
    std::vector<node_ptr> module_type_children;
    for (auto child : module_node->getChildren())
    {
        // if child is a component or other scene element
        if (!modules_nodes_set.contains(child))
            continue;
        module_type_children.emplace_back(child);
    }

    return module_type_children;
}

void RofiWorld::addModule(module_ptr module)
{
    voxel_graph->extend(module->getParts());

    const auto id = findAvailableID(ids_set);
    ids_set.emplace(id);
    modules_map.emplace(module, id);
    nodes_map.emplace(module->getNode(), module);

    modules_set.emplace(module);
    modules_nodes_set.emplace(module->getNode());

    addNodeToNodeTree(module->getNode());
}

void RofiWorld::removeModule(module_ptr module)
{
    disconnectModule(module);
    voxel_graph->erase(module->getParts());

    removeNodeFromNodeTree(module->getNode());
    ids_set.erase(modules_map[module]);
    modules_map.erase(module);
    nodes_map.erase(module->getNode());

    modules_set.erase(module);
    modules_nodes_set.erase(module->getNode());
}

bool RofiWorld::reachable(module_ptr fst_module, module_ptr snd_module) const
{
    std::set<module_ptr> found_modules;
    bool has_cycle = false; // Not being detected here
    traverseConnectedModules(fst_module, nullptr, nullptr, found_modules, has_cycle);

    return found_modules.contains(snd_module);
}

void RofiWorld::rotateVoxel(voxel_ptr voxel, float angle)
{
    voxel->getNode()->rotateRotationQuat(glm::angleAxis(glm::radians(angle), glm::vec3(1.0f, 0.0f, 0.0f)));
}

void RofiWorld::rotateConnector(connector_ptr connector, float angle)
{
    using namespace rofi;
    auto rot_con = std::dynamic_pointer_cast<RotationConnector>(connector);
    ASSUMPTION(rot_con);

    float current_rotation = rot_con->getCurrentRotation();

    angle = glm::clamp(angle, 
                       rot_con->getNegRotationLimit() - current_rotation, 
                       rot_con->getPosRotationLimit() - current_rotation);

    rot_con->setCurrentRotation(current_rotation + angle);

    auto connector_node = rot_con->getNode();
    auto rotation_mat = glm::toMat4(glm::angleAxis(glm::radians(-angle), glm::vec3(1.0f, 0.0f, 0.0f)));
    connector_node->translatePosVec(rotation_mat);
    connector_node->rotateRotationQuat(glm::toQuat(rotation_mat));
}

void RofiWorld::teleportModulesByConnectors(const connector_ptr source_connector, const connector_ptr destination_connector, const cardinal mutual_orientation)
{
    const auto source_node = source_connector->getNode();
    const auto destination_node = destination_connector->getNode();
    const auto source_parent = source_node->getRootParent().lock();

    //Rotation

    // Rotate to equal plane
    const auto &source_facing = source_node->getRotationAxisZWorld();
    const auto &destination_facing = destination_node->getRotationAxisZWorld();
    source_parent->rotateRotationQuat(glm::rotation(source_facing, -destination_facing));
    // Rotate North to North
    // auto source_north = source_connector->getNode()->getRotationAxisXWorld();
    // auto destination_north = destination_connector->getNode()->getRotationAxisXWorld();
    auto source_rot_con = std::dynamic_pointer_cast<RotationConnector>(source_connector);
    bool source_is_head = source_rot_con && source_rot_con->isHead();
    auto destination_rot_con = std::dynamic_pointer_cast<RotationConnector>(destination_connector);
    bool destination_is_head = destination_rot_con && destination_rot_con->isHead();
    const auto source_north = source_is_head 
                               ? -source_connector->getNode()->getRotationAxisXWorld()
                               : source_connector->getNode()->getRotationAxisXWorld();
    const auto destination_north = destination_is_head 
                                    ? -destination_connector->getNode()->getRotationAxisXWorld()
                                    : destination_connector->getNode()->getRotationAxisXWorld();
                                    
    const auto oriented_angle = glm::orientedAngle(source_north, destination_north, destination_facing);

    source_parent->rotateRotationQuat(glm::angleAxis(oriented_angle, destination_facing));

    // Rotate to selected cardinality
    auto orientation = glm::radians(cardinal_to_degree(mutual_orientation));

    source_parent->rotateRotationQuat(glm::angleAxis(orientation, destination_facing));

    // Translation

    const auto &destination_world_pos = destination_node->getPositionWorld();
    const auto &source_world_pos = source_node->getPositionWorld();
    source_parent->translatePosVec(destination_world_pos - source_world_pos);
}

std::set<module_ptr> RofiWorld::reachableModules(objectbase_ptr component) const
{
    node_ptr start_node;
    if (auto voxel = std::dynamic_pointer_cast<Voxel>(component))
        start_node = voxel->getNode();
    else if (auto connector = std::dynamic_pointer_cast<Connector>(component))
        start_node = connector->getNode();

    auto start_module = getModuleWithNode(start_node);

    bool has_cycle = false; // Not being detected here
    std::set<module_ptr> found_modules;
    traverseConnectedModules(start_module, nullptr, nullptr, found_modules, has_cycle);
    
    return found_modules;
}

bool RofiWorld::rebuildNodeTree(module_ptr module, node_ptr root_node) // NOTE: Maybe could use views if I knew how to.
{
    std::set<module_ptr> found_modules;
    bool has_cycle = false; // Not being detected here
    traverseConnectedModules(module, nullptr, nullptr, found_modules, has_cycle);
    found_modules.erase(module);

    // Erase all module children in node_tree, children frames are automatically recalculated to world space
    destroyNodeTree(modules_nodes_set, node_tree);
    node_tree.clear();

    // found modules are set as children of selected module, their frame is automatically recalculated into selected module base
    std::vector<node_ptr> found_modules_nodes;
    std::ranges::transform(found_modules, std::back_inserter(found_modules_nodes), [](const auto m) { return m->getNode(); } );
    root_node->extendChildren(found_modules_nodes, true);
    
    // Remaining modules are placed into the tree vector, including the selected module
    std::set<module_ptr> independent_modules;
    std::set_difference(modules_set.begin(), modules_set.end(), found_modules.begin(), found_modules.end(), std::inserter(independent_modules, independent_modules.begin()));

    std::vector<node_ptr> independent_modules_nodes;
    std::ranges::transform(independent_modules, std::back_inserter(independent_modules_nodes), [](const auto m) { return m->getNode(); } ); 
    extendNodeTree(independent_modules_nodes);

    return true;
}

bool RofiWorld::rebuildNodeTree(voxel_ptr voxel, node_ptr root_node)
{
    if (!voxel->hasRotationConnector())
        return true;

    auto voxel_node = voxel->getNode();
    std::set<module_ptr> found_modules;
    bool has_cycle = false;
    auto module = getModuleWithNode(voxel_node);
    found_modules.emplace(module);

    auto rot_con = std::dynamic_pointer_cast<rofi::RotationConnector>(voxel->getRotationConnector().lock());
    std::vector<connector_ptr> connectors = { rot_con, rot_con->getBoundConnectors().first.lock(), rot_con->getBoundConnectors().second.lock() };
    for (auto & connector : connectors)
    {
        if (!connector->hasModuleLink())
            continue;

        auto opposite_connector = connector->getModuleLink().lock()->getOppositeConnector(connector).lock();
        traverseConnectedModules(getModuleWithNode(opposite_connector->getNode()), module, rot_con, found_modules, has_cycle);
    }

    if (has_cycle)
        return false;

    found_modules.erase(module);

    // Erase all module children in node_tree, children frames are automatically recalculated to world space
    destroyNodeTree(modules_nodes_set, node_tree);
    node_tree.clear();

    // found modules are set as children of selected module, their frame is automatically recalculated into selected module base
    std::vector<node_ptr> found_modules_nodes;
    std::ranges::transform(found_modules, std::back_inserter(found_modules_nodes), [](const auto m) { return m->getNode(); } );
    root_node->extendChildren(found_modules_nodes, true);
    
    // Remaining modules are placed into the tree vector, including the selected module
    std::set<module_ptr> independent_modules;
    std::set_difference(modules_set.begin(), modules_set.end(), found_modules.begin(), found_modules.end(), std::inserter(independent_modules, independent_modules.begin()));

    std::vector<node_ptr> independent_modules_nodes;
    std::ranges::transform(independent_modules, std::back_inserter(independent_modules_nodes), [](const auto m) { return m->getNode(); } ); 
    extendNodeTree(independent_modules_nodes);

    return true;
}

bool RofiWorld::detectCollision(std::set<module_ptr> &collided_modules) const
{
    bool retval = false;

    for (const auto &[module_node, module] : nodes_map)
    {
        for (const auto &voxel : module->getParts()->getVoxels())
        {
            
            for (const auto &[module_node2, module2] : nodes_map)
            {
                if (module_node == module_node2)
                    continue;

                for (const auto &voxel2 : module2->getParts()->getVoxels())
                {
                    // 0.86 is the length of the longest side of UM body mesh
                    if (glm::distance(voxel->getNode()->getPositionWorld(), voxel2->getNode()->getPositionWorld()) < 0.86f)
                    {
                        collided_modules.emplace(module);
                        collided_modules.emplace(module2);
                        retval = true;
                    }
                }
            }
        }
    }

    return retval;
}

/* PRIVATE */

uint64_t RofiWorld::findAvailableID(const std::set<uint64_t> &ids_set)
{
    uint64_t id = 0;

    while (true)
    {
        if (!ids_set.contains(id))
            return id;
        ++id;
    }

    ASSUMPTION(false);
    return id;
}

void RofiWorld::disconnectModule(module_ptr module)
{
    for (auto &connector: module->getParts()->getConnections())
    {
        if (!connector->hasModuleLink())
            continue;

        auto opposite_connector = connector->getModuleLink().lock()->getOppositeConnector(connector).lock();

        voxel_graph->disjoinConnectors(connector, opposite_connector);
    }
}

void RofiWorld::addNodeToNodeTree(node_ptr node)
{
    node_tree.emplace_back(node);
}
void RofiWorld::removeNodeFromNodeTree(const node_ptr node_to_remove)
{
    removeNodeFromNodeTreeRec(node_to_remove, node_tree);
}

void RofiWorld::removeNodeFromNodeTreeRec(const node_ptr node_to_remove, std::vector<node_ptr> &node_tree)
{
    if (node_tree.empty())
        return;

    if (std::erase(node_tree, node_to_remove) != 0)
    {
        auto module_type_children = getDependentModulesNodes(node_to_remove);
        node_to_remove->removeChildren(module_type_children);
        extendNodeTree(module_type_children);
    }

    for (const auto &node : node_tree)
        removeNodeFromNodeTreeRec(node_to_remove, node->getChildren());
}

void RofiWorld::extendNodeTree(const std::vector<node_ptr> &nodes)
{
    node_tree.reserve(node_tree.size() + nodes.size());
    node_tree.insert(node_tree.end(), nodes.begin(), nodes.end());
}

void RofiWorld::traverseConnectedModules(module_ptr module, module_ptr initial_module, connector_ptr initial_shoe, std::set<module_ptr> &found_modules, bool &has_cycle) const
{
    auto module_vg = module->getParts();
    found_modules.emplace(module);

    for (auto &connector : module_vg->getConnections())
    {
        if (!connector->hasModuleLink())
            continue;
        
        auto opposite_connector = connector->getModuleLink().lock()->getOppositeConnector(connector).lock();
        auto next_module = getModuleWithNode(opposite_connector->getNode());

        if (initial_module)
            detectCycle(next_module, initial_module, opposite_connector, initial_shoe, has_cycle);

        if (!found_modules.contains(next_module))
            traverseConnectedModules(next_module, initial_module, initial_shoe, found_modules, has_cycle);
    }
}

void RofiWorld::detectCycle(module_ptr next_module, module_ptr initial_module, connector_ptr opposite_connector, connector_ptr initial_shoe, bool &has_cycle) const
{
    auto opposite_rot_con = std::dynamic_pointer_cast<RotationConnector>(opposite_connector);
    connector_ptr opposite_head_con;

    if (opposite_rot_con)
        opposite_head_con = opposite_rot_con->isHead() ? opposite_rot_con : opposite_rot_con->getBoundConnectors().first.lock();

    if (next_module == initial_module && opposite_rot_con && opposite_head_con != initial_shoe)
        has_cycle = true;
}

void RofiWorld::destroyNodeTree(const std::set<node_ptr> &module_set_nodes, std::vector<node_ptr> &node_tree)
{
    for (auto &node : node_tree)
    {
        auto &children = node->getChildren();
        destroyNodeTree(module_set_nodes, children);

        const std::vector<node_ptr> children_to_remove = getDependentModulesNodes(node);
        node->removeChildren(children_to_remove);
    }
}

}