from direct.showbase.ShowBase import ShowBase
from panda3d.core import CollisionTraverser, CollisionNode, CollisionHandlerQueue, GeomNode, CollisionRay, NodePath, LPoint3f, LVector3f


class PickerRay():
    def __init__(self, root: ShowBase, pick_root: NodePath = None):
        super().__init__()
        self.root: ShowBase = root
        self.pick_root = pick_root or root.render
        self.debug = False
    
    def __setup_picker(self):
        # 创建一个拾取器
        self.picker = CollisionTraverser()
        self.pq = CollisionHandlerQueue()
        
        # 创建一个拾取节点
        self.picker_node = CollisionNode('mouseRay')
        self.picker_np = self.root.cam.attachNewNode(self.picker_node)
        self.picker_node.setFromCollideMask(GeomNode.getDefaultCollideMask())
        
        # 创建射线
        self.picker_ray = CollisionRay()
        self.picker_node.addSolid(self.picker_ray)
        
        # 将拾取器和处理器关联
        self.picker.addCollider(self.picker_np, self.pq)
        
        # 可选：显示拾取射线（调试用）
        # self.picker_np.show()
    
    def debug_log(self, msg: str) -> None:
        if self.debug:
            print(f"[ScaleGizmo] -> {msg}")
    
    def pick_object(self) -> tuple[NodePath | None, LVector3f | None]:
        """根据屏幕坐标拾取3D物体"""
        if not hasattr(self, 'picker'):
            self.__setup_picker()
            
        mpos = self.root.mouseWatcherNode.getMouse()
        screen_x = mpos.x
        screen_y = mpos.y
        self.debug_log(f"屏幕坐标: {screen_x}, {screen_y}")
        
        # 从摄像机发射射线
        self.picker_ray.setFromLens(self.root.camNode, screen_x, screen_y)
        
        # 执行碰撞检测
        self.picker.traverse(self.pick_root)
        
        picked_object = None
        collision_point = LPoint3f(0, 0, 0)
        
        # 检查是否有碰撞
        if self.pq.getNumEntries() > 0:
            # 按距离排序，最近的在前面
            self.pq.sortEntries()
            
            for i in range(self.pq.getNumEntries()):
                # 获取最近的碰撞点
                closest_entry = self.pq.getEntry(i)
                picked_object: NodePath = closest_entry.getIntoNodePath()
                
                if picked_object.isHidden():
                    continue
                else:
                    # 获取碰撞点的世界坐标
                    collision_point = closest_entry.getSurfacePoint(self.root.render)
                    
                    self.debug_log(f"拾取到物体: {picked_object.getName()}")
                    self.debug_log(f"碰撞点: {collision_point}")
                    break
            return picked_object, collision_point
        
        self.debug_log("未拾取到任何物体")
        return None, None