import os
import vtk
import pyvista as pv

# -----------------------------------------------------------------------------
# 1. 配置：要展示的文件夹路径（改成你自己的）
# -----------------------------------------------------------------------------
ROOT_FOLDER = r"./"  # 示例：当前目录
MAX_DEPTH = 3        # 限制深度，防止太密
NODE_SPACING_X = 2.0
NODE_SPACING_Y = 1.5


# -----------------------------------------------------------------------------
# 2. 递归遍历文件夹，生成树形结构数据
# -----------------------------------------------------------------------------
class TreeNode:
    def __init__(self, name, path, depth=0):
        self.name = name
        self.path = path
        self.depth = depth
        self.children = []
        self.x = 0.0
        self.y = 0.0


def build_tree(folder, depth=0, max_depth=MAX_DEPTH):
    node = TreeNode(os.path.basename(folder) or folder, folder, depth)
    if depth >= max_depth:
        return node

    try:
        entries = sorted(os.listdir(folder))
    except PermissionError:
        return node

    for entry in entries:
        full = os.path.join(folder, entry)
        if os.path.isdir(full):
            child = build_tree(full, depth + 1, max_depth)
            node.children.append(child)
    return node


# -----------------------------------------------------------------------------
# 3. 树布局算法（简单层级横向布局）
# -----------------------------------------------------------------------------
def layout_tree(node: TreeNode, x=0.0, y=0.0):
    node.y = y
    node.x = x

    child_x = x - (len(node.children) - 1) * NODE_SPACING_X / 2
    for child in node.children:
        layout_tree(child, child_x, y - NODE_SPACING_Y)
        child_x += NODE_SPACING_X


# -----------------------------------------------------------------------------
# 4. 收集所有节点、边用于 VTK 绘制
# -----------------------------------------------------------------------------
def collect_geometry(node):
    points = []
    lines = []
    labels = []
    stack = [node]

    while stack:
        n = stack.pop()
        points.append([n.x, n.y, 0.0])
        labels.append(n.name)
        for child in n.children:
            points.append([child.x, child.y, 0.0])
            labels.append(child.name)
            lines.append(2)  # 每个线段有两个点
            lines.append(len(points)-2)  # 父节点的索引
            lines.append(len(points)-1)  # 子节点的索引
            stack.append(child)
    return points, lines, labels


# -----------------------------------------------------------------------------
# 5. PyVista + VTK 绘制树
# -----------------------------------------------------------------------------
if __name__ == "__main__":
    tree = build_tree(ROOT_FOLDER)
    layout_tree(tree)
    points, lines, labels = collect_geometry(tree)

    # 构建点线云
    cloud = pv.PolyData(points)
    cloud.lines = lines
    cloud["label"] = labels

    # 绘图
    plotter = pv.Plotter(window_size=(1200, 800))
    plotter.add_mesh(cloud, color="#66ccff", point_size=12, render_lines_as_tubes=True, line_width=3)
    plotter.add_point_labels(cloud, "label", font_size=8, shape_color="white")

    plotter.add_text(
        f"Folder Tree: {ROOT_FOLDER}\nMax Depth: {MAX_DEPTH}",
        font_size=10, color="blue"
    )
    plotter.camera_position = "xy"
    plotter.show()