#pragma once
#include "vde/brep/brep.h"
#include "vde/core/transform.h"
#include <string>
#include <vector>
#include <optional>
#include <memory>

namespace vde::brep {

using core::Point3D;
using core::Vector3D;
using core::Transform3D;

/**
 * @brief 装配体节点
 *
 * 表示装配体树中的一个节点，可以是零件（含 B-Rep 模型）或子装配体。
 * 每个节点有局部变换矩阵和名称。
 *
 * @ingroup brep
 */
struct AssemblyNode {
    /// 节点名称
    std::string name;

    /// 局部（相对父节点）变换矩阵
    Transform3D local_transform = Transform3D::Identity();

    /// 此节点的 B-Rep 模型（零件），可选
    std::optional<BrepModel> model;

    /// 子节点列表
    std::vector<std::unique_ptr<AssemblyNode>> children;

    /// 是否是零件节点（有 model）还是子装配体（有 children）
    [[nodiscard]] bool is_part() const { return model.has_value() && children.empty(); }

    /// 添加子节点
    AssemblyNode* add_child(std::unique_ptr<AssemblyNode> child) {
        auto* ptr = child.get();
        children.push_back(std::move(child));
        return ptr;
    }

    /// 添加 B-Rep 零件
    AssemblyNode* add_part(const std::string& part_name, BrepModel part_model,
                           const Transform3D& transform = Transform3D::Identity()) {
        auto node = std::make_unique<AssemblyNode>();
        node->name = part_name;
        node->model = std::move(part_model);
        node->local_transform = transform;
        return add_child(std::move(node));
    }

    /// 添加子装配体
    AssemblyNode* add_subassembly(const std::string& sub_name,
                                  const Transform3D& transform = Transform3D::Identity()) {
        auto node = std::make_unique<AssemblyNode>();
        node->name = sub_name;
        node->local_transform = transform;
        return add_child(std::move(node));
    }
};

/**
 * @brief 装配体
 *
 * 装配体树的根节点。管理零件和子装配体的层级结构。
 *
 * @code{.cpp}
 * Assembly assembly("robot_arm");
 * auto* base = assembly.root().add_part("base", make_box(10,10,2));
 * auto* arm = assembly.root().add_subassembly("arm",
 *     Transform3D::Translation(0,0,2));
 * arm->add_part("segment1", make_cylinder(0.5, 8));
 * @endcode
 *
 * @ingroup brep
 */
struct Assembly {
    /// 装配体名称
    std::string name;

    /// 版本号
    std::string version = "1.0";

    /// 根节点
    AssemblyNode root;

    explicit Assembly(const std::string& n) : name(n) { root.name = n; }

    /// 遍历所有零件节点
    template <typename Visitor>
    void visit_parts(Visitor&& visitor) const {
        visit_parts_impl(&root, Transform3D::Identity(),
                         std::forward<Visitor>(visitor));
    }

    /// 统计零件总数
    [[nodiscard]] size_t part_count() const {
        size_t count = 0;
        visit_parts([&count](const std::string&, const BrepModel&, const Transform3D&) {
            ++count;
        });
        return count;
    }

    /// 统计节点总数（含子装配体）
    [[nodiscard]] size_t node_count() const {
        return node_count_impl(&root);
    }

private:
    template <typename Visitor>
    void visit_parts_impl(const AssemblyNode* node, const Transform3D& parent_xform,
                          Visitor&& visitor) const {
        Transform3D world = parent_xform * node->local_transform;
        if (node->model.has_value()) {
            visitor(node->name, *node->model, world);
        }
        for (const auto& child : node->children) {
            visit_parts_impl(child.get(), world, std::forward<Visitor>(visitor));
        }
    }

    size_t node_count_impl(const AssemblyNode* node) const {
        size_t n = 1; // count this node
        for (const auto& child : node->children)
            n += node_count_impl(child.get());
        return n;
    }
};

} // namespace vde::brep