#include <cassert>

#include "internal_api.hpp"

#include "globals.hpp"
#include "internal_primitive_desc.hpp"
#include "internal_structs.hpp"

/* internal global variables for blobtree */
std::vector<blobtree_t, tbb::tbb_allocator<blobtree_t>>                  structures{};
std::vector<primitive_node_t, tbb::tbb_allocator<primitive_node_t>>      primitives{};
std::stack<uint32_t, std::deque<uint32_t, tbb::tbb_allocator<uint32_t>>> free_structure_list{};

/* =============================================================================================
 * basic functionalities
 * ============================================================================================= */

BS_API size_t blobtree_get_node_count(uint32_t index) noexcept { return structures[index].nodes.size(); }

BS_API std::vector<uint32_t, tbb::tbb_allocator<uint32_t>> blobtree_get_leaf_nodes(uint32_t index) noexcept
{
    return structures[index].leaf_index;
}

BS_API node_t& blobtree_get_node(const virtual_node_t& node) noexcept
{
    return structures[node.main_index].nodes[node.inner_index];
}

BS_API size_t get_primitive_count() noexcept { return primitives.size(); }

BS_API const primitive_node_t& get_primitive_node(uint32_t index) noexcept
{
    assert(index < primitives.size());
    return primitives[index];
}

void shrink_primitives() { primitives.shrink_to_fit(); }

virtual_node_t copy(virtual_node_t old_node, virtual_node_t new_node)
{
    assert(old_node.main_index != new_node.main_index);

    // Copy a tree and its subtrees to a temporary tree
    auto temp = structures[old_node.main_index];

    // Update all index
    const auto size = static_cast<uint32_t>(structures[new_node.main_index].nodes.size());
    for (uint32_t i = 0; i < temp.nodes.size(); i++) {
        if (!node_is_parent_null(temp.nodes[i])) node_fetch_parent_index(temp.nodes[i]) += size;
        if (!node_is_left_child_null(temp.nodes[i])) node_fetch_left_child_index(temp.nodes[i]) += size;
        if (!node_is_right_child_null(temp.nodes[i])) node_fetch_right_child_index(temp.nodes[i]) += size;
    }
    for (uint32_t i = 0; i < temp.leaf_index.size(); i++) { temp.leaf_index[i] += size; }

    // Copy the updated index tree to the array at the new location
    structures[new_node.main_index].nodes.insert(structures[new_node.main_index].nodes.end(),
                                                 temp.nodes.begin(),
                                                 temp.nodes.end());
    structures[new_node.main_index].leaf_index.insert(structures[new_node.main_index].leaf_index.end(),
                                                      temp.leaf_index.begin(),
                                                      temp.leaf_index.end());

    return virtual_node_t{new_node.main_index, old_node.inner_index + size};
}

BS_API void free_sub_blobtree(uint32_t index) noexcept
{
    // 这里尽量打标记，延迟修改和删除
    free_structure_list.push(index);
}

BS_API void clear_blobtree() noexcept
{
    structures.clear();
    primitives.clear();
    while (!free_structure_list.empty()) free_structure_list.pop();
}

/* =============================================================================================
 * geometry generation
 * ============================================================================================= */

virtual_node_t push_primitive_node(primitive_node_t&& primitive_node)
{
    primitives.emplace_back(primitive_node);
    primitive_node.desc =
        nullptr; // NOTE: primitives.back().desc and primitive_node.desc share same memory, so when this
                 // function returns, primitive_node.desc will be invalidated, so as primitives.back().desc, which outputs
                 // errors. Thus, we manually set primitive_node.desc to nullptr to avoid this problem.

    node_t node{standard_new_node};
    node_fetch_primitive_index(node) = static_cast<uint32_t>(primitives.size() - 1);

    auto& tree = structures.emplace_back();
    tree.nodes.emplace_back(std::move(node));
    tree.leaf_index.emplace_back(0);

    return virtual_node_t{static_cast<uint32_t>(structures.size() - 1), 0};
}

BS_API virtual_node_t blobtree_new_virtual_node(const legal_primitive_descriptor_t& desc)
{
    return push_primitive_node(primitive_node_t{desc});
}

BS_API virtual_node_t blobtree_new_virtual_node(legal_primitive_descriptor_t&& desc)
{
    return push_primitive_node(primitive_node_t{std::move(desc)});
}

BS_API void blobtree_free_virtual_node(const virtual_node_t& node) { free_sub_blobtree(node.main_index); }

/* =============================================================================================
 * tree node operations
 * ============================================================================================= */

BS_API bool virtual_node_set_parent(const virtual_node_t& node, const virtual_node_t& parent)
{
    auto& node_in_tree = structures[node.main_index].nodes[node.inner_index];
    // The node's parent is not empty
    if (!node_is_parent_null(node_in_tree)) { return false; }

    auto& parent_in_tree     = structures[parent.main_index].nodes[parent.inner_index];
    auto  parent_inner_index = parent.inner_index;
    // not on the same tree
    if (node.main_index != parent.main_index) {
        auto new_parent    = copy(parent, node);
        parent_in_tree     = structures[new_parent.main_index].nodes[new_parent.inner_index];
        parent_inner_index = new_parent.inner_index;
    }
    auto parent_left_child  = node_fetch_left_child_index(parent_in_tree);
    auto parent_right_child = node_fetch_right_child_index(parent_in_tree);

    // set parent index
    node_fetch_parent_index(node_in_tree) = parent_inner_index;

    // The parent's left child is empty
    if (node_is_left_child_null(node_in_tree)) {
        parent_left_child = node.inner_index;
        return true;
    }
    // The parent's right child is empty
    else if (node_is_right_child_null(node_in_tree)) {
        parent_right_child = node.inner_index;
        return true;
    }

    return false;
}

BS_API bool virtual_node_set_left_child(const virtual_node_t& node, const virtual_node_t& child)
{
    auto& node_in_tree = structures[node.main_index].nodes[node.inner_index];

    // The child's parent is not empty
    if (!node_is_parent_null(node_in_tree)) { return false; }

    // The node's left child is not empty
    if (!node_is_left_child_null(node_in_tree)) { return false; }

    auto node_left_child  = node_fetch_left_child_index(node_in_tree);
    auto node_right_child = node_fetch_right_child_index(node_in_tree);

    // On the same tree
    if (node.main_index == child.main_index) {
        auto& child_in_tree                    = structures[child.main_index].nodes[child.inner_index];
        node_fetch_parent_index(child_in_tree) = node.inner_index;
        node_left_child                        = child.inner_index;
    } else {
        auto  new_child                        = copy(child, node);
        auto& child_in_tree                    = structures[new_child.main_index].nodes[new_child.inner_index];
        node_fetch_parent_index(child_in_tree) = node.inner_index;
        node_left_child                        = new_child.inner_index;
    }

    return true;
}

BS_API bool virtual_node_set_right_child(const virtual_node_t& node, const virtual_node_t& child)
{
    auto& node_in_tree = structures[node.main_index].nodes[node.inner_index];

    // The child's parent is not empty
    if (!node_is_parent_null(node_in_tree)) { return false; }

    // The node's right child is not empty
    if (!node_is_right_child_null(node_in_tree)) { return false; }

    auto node_left_child  = node_fetch_left_child_index(node_in_tree);
    auto node_right_child = node_fetch_right_child_index(node_in_tree);

    // On the same tree
    if (node.main_index == child.main_index) {
        auto& child_in_tree                    = structures[child.main_index].nodes[child.inner_index];
        node_fetch_parent_index(child_in_tree) = node.inner_index;
        node_right_child                       = child.inner_index;
    } else {
        auto  new_child                        = copy(child, node);
        auto& child_in_tree                    = structures[new_child.main_index].nodes[new_child.inner_index];
        node_fetch_parent_index(child_in_tree) = node.inner_index;
        node_right_child                       = new_child.inner_index;
    }

    return true;
}

BS_API bool virtual_node_add_child(const virtual_node_t& node, const virtual_node_t& child)
{
    if (virtual_node_set_left_child(node, child)) {
        return true;
    } else if (virtual_node_set_right_child(node, child)) {
        return true;
    }

    return false;
}

BS_API bool virtual_node_remove_child(const virtual_node_t& node, const virtual_node_t& child)
{
    if (node.main_index != child.main_index) { return false; }

    auto& node_in_tree     = structures[node.main_index].nodes[node.inner_index];
    auto  node_left_child  = node_fetch_left_child_index(node_in_tree);
    auto  node_right_child = node_fetch_right_child_index(node_in_tree);

    if (node_left_child == child.inner_index) {
        node_left_child = 0xFFFFFFFFu;
        blobtree_free_virtual_node(child);
        return true;
    } else if (node_right_child == child.inner_index) {
        node_right_child = 0xFFFFFFFFu;
        blobtree_free_virtual_node(child);
        return true;
    }

    return false;
}

/* =============================================================================================
 * geometry operations
 * ============================================================================================= */

static inline void virtual_node_boolean_op(virtual_node_t& node1, const virtual_node_t& node2, eNodeOperation op)
{
    auto new_node2 = copy(node2, node1);

    auto& inserted_node                         = structures[node1.main_index].nodes.emplace_back(standard_new_node);
    node_fetch_is_primitive(inserted_node)      = false;
    // weird bug: need to force cast, or it will be treated as uint32_t instead of eNodeOperation
    node_fetch_operation(inserted_node)         = (eNodeOperation)op;
    node_fetch_left_child_index(inserted_node)  = node1.inner_index;
    node_fetch_right_child_index(inserted_node) = new_node2.inner_index;

    uint32_t parent_index = structures[node1.main_index].nodes.size() - 1;
    node_fetch_parent_index(structures[node1.main_index].nodes[node1.inner_index])         = parent_index;
    node_fetch_parent_index(structures[new_node2.main_index].nodes[new_node2.inner_index]) = parent_index;

    node1.inner_index = parent_index;
}

BS_API void virtual_node_boolean_union(virtual_node_t& node1, const virtual_node_t& node2)
{
    virtual_node_boolean_op(node1, node2, eNodeOperation::unionOp);
}

BS_API void virtual_node_boolean_intersect(virtual_node_t& node1, const virtual_node_t& node2)
{
    virtual_node_boolean_op(node1, node2, eNodeOperation::intersectionOp);
}

BS_API void virtual_node_boolean_difference(virtual_node_t& node1, const virtual_node_t& node2)
{
    virtual_node_boolean_op(node1, node2, eNodeOperation::differenceOp);
}

void offset_primitive(primitive_node_t& node, const Eigen::Vector3d& offset)
{
    auto offset_point = [](Eigen::Ref<Eigen::Vector3d> point, const Eigen::Ref<const Eigen::Vector3d>& offset) {
        point += offset;
    };

    auto type = node.type;
    switch (type) {
        case PRIMITIVE_TYPE_CONSTANT: {
            break;
        }
        case PRIMITIVE_TYPE_PLANE: {
            auto desc = static_cast<internal::plane*>(node.desc);
            offset_point(desc->point, offset);
            break;
        }
        case PRIMITIVE_TYPE_SPHERE: {
            auto desc = static_cast<internal::sphere*>(node.desc);
            offset_point(desc->center, offset);
            break;
        }
        case PRIMITIVE_TYPE_CYLINDER: {
            auto desc = static_cast<internal::cylinder*>(node.desc);
            offset_point(desc->bottom_center, offset);
            break;
        }
        case PRIMITIVE_TYPE_CONE: {
            auto desc = static_cast<internal::cone*>(node.desc);
            offset_point(desc->top_center, offset);
            offset_point(desc->bottom_center, offset);
            break;
        }
        case PRIMITIVE_TYPE_BOX: {
            auto desc = static_cast<internal::box*>(node.desc);
            offset_point(desc->center, offset);

            break;
        }
        case PRIMITIVE_TYPE_MESH: {
            auto desc = static_cast<internal::mesh*>(node.desc);
            for (auto& v : desc->vertices) v += offset;
            break;
        }
        case PRIMITIVE_TYPE_EXTRUDE_POLYLINE: {
            auto  desc            = static_cast<internal::extrude_polyline*>(node.desc);
            auto& matrix_handle   = desc->axis_to_world.matrix();
            matrix_handle.col(3) += offset;
            break;
        }
        case PRIMITIVE_TYPE_EXTRUDE_HELIXLINE: {
            auto  desc            = static_cast<internal::extrude_helixline*>(node.desc);
            auto& matrix_handle   = desc->axis_to_world.matrix();
            matrix_handle.col(3) += offset;
            break;
        }
    }

    node.aabb.min() += offset;
    node.aabb.max() += offset;
}

BS_API void virtual_node_offset(virtual_node_t& node, const raw_vector3d_t& direction, const double length)
{
    raw_vector3d_t offset = {direction.x * length, direction.y * length, direction.z * length};
    virtual_node_offset(node, offset);
}

BS_API void virtual_node_offset(virtual_node_t& node, const raw_vector3d_t& offset)
{
    Eigen::Map<const Eigen::Vector3d> offset_(&offset.x);

    auto& all_leaf = structures[node.main_index].leaf_index;
    for (const auto& leaf_index : structures[node.main_index].leaf_index) {
        auto&          primitive_node  = structures[node.main_index].nodes[leaf_index];
        const uint32_t primitive_index = node_fetch_primitive_index(primitive_node);

        offset_primitive(primitives[primitive_index], offset_);
    }
}

BS_API void virtual_node_split(virtual_node_t& node, raw_vector3d_t base_point, raw_vector3d_t normal)
{
    plane_descriptor_t descriptor;
    descriptor.normal = raw_vector3d_t{normal.x * -1, normal.y * -1, normal.z * -1};
    descriptor.point  = base_point;
    auto plane        = blobtree_new_virtual_node(descriptor);

    virtual_node_boolean_intersect(node, plane);
}

/* =============================================================================================
 * node replacement operations
 * ============================================================================================= */

bool replace_primitive_node(const virtual_node_t& node, primitive_node_t&& new_primitive_node)
{
    auto& node_in_tree = structures[node.main_index].nodes[node.inner_index];

    if (!node_fetch_is_primitive(node_in_tree)) { return false; }

    const uint32_t primitive_index = node_fetch_primitive_index(node_in_tree);
    auto&          primitive       = primitives[primitive_index];
    primitive.~primitive_node_t();
    primitive = std::move(new_primitive_node);

    return true;
}

BS_API bool virtual_node_replace_primitive(const virtual_node_t& node, const legal_primitive_descriptor_t& desc)
{
    return replace_primitive_node(node, primitive_node_t{desc});
}

BS_API bool virtual_node_replace_primitive(const virtual_node_t& node, legal_primitive_descriptor_t&& desc)
{
    return replace_primitive_node(node, primitive_node_t{std::move(desc)});
}

/* =============================================================================================
 * evaluation proxy operations
 * ============================================================================================= */

BS_API double evaluate_sdf(uint32_t index, const Eigen::Ref<const Eigen::Vector3d>& point)
{
    const auto& primitive = primitives[index];
    return primitive.evaluate_sdf(point);
}