import direct
import direct.showbase
import direct.showbase.ShowBaseGlobal
from panda3d.core import (
    Geom,
    GeomLines,
    GeomNode,
    GeomVertexData,
    GeomVertexFormat,
    GeomVertexWriter,
    LColor,
    LPoint3,
    NodePath
)
import panda3d.core as p3d

class CornerBBox:
    def __init__(
        self,
        target: NodePath,
        color: tuple = (0.5, 0, 0.5, 1),
        corner_offset: float = 0.52,
        stub_scale: float = 0.075,
        line_thickness: float = 2.0,
        parent: NodePath | None = None
    ) -> None:
        """
        创建一个“角点风格”的包围盒可视化框。

        :param target: 要描框的 NodePath
        :param color: 线颜色 (r, g, b, a)
        :param corner_offset: 角点在包围盒外扩的偏移，决定框外伸出的长度
        :param stub_scale: 每个角三条小线的比例（实际使用两段：stub 和 3*stub + stub）
        :param line_thickness: 线粗细
        :param parent: 框的挂载节点，默认与 target 同级父节点
        """
        self.world = direct.showbase.ShowBaseGlobal.base
        self.target = target
        self.parent = parent or target
        self.color = LColor(*color)
        self.corner_offset = corner_offset
        self.stub_scale = stub_scale
        self.line_thickness = line_thickness
        self.np = self._build_geom()
        self.update()  # 初始化尺寸/位置

    def _build_geom(self) -> NodePath:
        """构造“八角+三向短线”几何体"""
        fmt = GeomVertexFormat.get_v3c4()
        vdata = GeomVertexData("corner_bbox", fmt, Geom.UH_static)
        pos_w = GeomVertexWriter(vdata, "vertex")
        col_w = GeomVertexWriter(vdata, "color")
        lines = GeomLines(Geom.UH_static)

        # 定义 8 个角（基于偏移量 m）
        m = self.corner_offset
        corners = [(x, y, z) for x in (-m, m) for y in (-m, m) for z in (-m, m)]

        for corner in corners:
            cx, cy, cz = corner
            for axis_idx, coord in enumerate((cx, cy, cz)):
                sign = -1.0 if coord > 0 else 1.0
                # 两段短线：长度分别为 stub_scale 和 3 * stub_scale
                offset1 = self.stub_scale * sign
                offset2 = (3.0 * self.stub_scale + self.stub_scale) * sign

                base = list(corner)

                # 第一段线：从角点到第一个 stub
                p1 = LPoint3(*base)
                pos_w.add_data3(p1)
                col_w.add_data4(self.color)

                base[axis_idx] = coord + offset1
                p2 = LPoint3(*base)
                pos_w.add_data3(p2)
                col_w.add_data4(self.color)
                lines.add_next_vertices(2)

                # 第二段线：从第一个 stub 到更远的 stub
                base[axis_idx] = coord + offset2
                p3 = LPoint3(*base)
                pos_w.add_data3(p2)  # 重复起点
                col_w.add_data4(self.color)
                pos_w.add_data3(p3)
                col_w.add_data4(self.color)
                lines.add_next_vertices(2)

        geom = Geom(vdata)
        geom.add_primitive(lines)
        node = GeomNode("corner_bbox")
        node.add_geom(geom)
        np = self.parent.attach_new_node(node)

        # 渲染设置
        np.set_light_off()
        np.set_texture_off()
        np.set_color_off(False)  # 保留顶点色
        np.set_render_mode_thickness(self.line_thickness)
        
        mat = p3d.Material('unlit')
        np.setMaterial(mat)

        return np

    def update(self) -> None:
        """根据 target 的 tight bounds 更新包围框的位置/缩放。"""
        if not self.world or self.target.is_empty() or self.np is None:
            return
        
        render = self.world.render

        # 统一在 render 坐标系下保存/恢复
        last_hpr = self.target.getHpr(render)

        try:
            # 暂时归零朝向
            self.target.setHpr(render, 0, 0, 0)

            bmin, bmax = self.target.get_tight_bounds(render)
            if not bmin or not bmax:
                return
            
            size = bmax - bmin
            center = (bmin + bmax) * 0.5

            # bbox 挂在 parent 下，但变换以 render 为参照设置，避免坐标系错乱
            self.np.set_pos(render, center)
            self.np.set_scale(render, size)

        finally:
            # 无论成功与否都必须恢复 target 姿态
            self.target.setHpr(render, last_hpr)


    def remove(self) -> None:
        """从场景图中移除该包围框。"""
        if self.np is None:
            return
        self.np.detach_node()
        self.np = None