from direct.task.TaskManagerGlobal import taskMgr


class CustomMouseController:
    def __init__(self, showbase):
        self.showbase = showbase
        # This is used to store which keys are currently pressed.
        self.keyMap = {
            "mouse1": 0,
            "cam-forward": 0,
            "cam-backward": 0,
            "cam-left": 0,
            "cam-right": 0,
            "cam-up": 0,
            "cam-down": 0
        }
        
        # 添加鼠标控制
        # self.showbase.accept("mouse1", self.setKey, ["mouse1", True])
        # self.showbase.accept("mouse1-up", self.setKey, ["mouse1", False])
        
        self.showbase.accept("w", self.setKey, ["cam-forward", True])
        self.showbase.accept("a", self.setKey, ["cam-left", True])
        self.showbase.accept("s", self.setKey, ["cam-backward", True])
        self.showbase.accept("d", self.setKey, ["cam-right", True])
        self.showbase.accept("e", self.setKey, ["cam-up", True])
        self.showbase.accept("q", self.setKey, ["cam-down", True])
        
        self.showbase.accept("w-up", self.setKey, ["cam-forward", False])
        self.showbase.accept("a-up", self.setKey, ["cam-left", False])
        self.showbase.accept("s-up", self.setKey, ["cam-backward", False])
        self.showbase.accept("d-up", self.setKey, ["cam-right", False])
        self.showbase.accept("e-up", self.setKey, ["cam-up", False])
        self.showbase.accept("q-up", self.setKey, ["cam-down", False])
        
        self.last_mouse_x = 0
        self.last_mouse_y = 0
        
        self.camera_speed = 25
        self.move_speed = 20
    
    def setUp(self, mouse_speed: float = 25, move_speed: float = 20):
        taskMgr.add(self.move, "moveTask")
        self.camera_speed = mouse_speed
        self.move_speed = move_speed
    
    def setKey(self, key, value, arg: str = None):
        self.keyMap[key] = value
        if key == "mouse1" and value == True and self.keyMap[key]:
            mouse_pos = self.showbase.mouseWatcherNode.getMouse()
            self.last_mouse_x = mouse_pos.get_x()
            self.last_mouse_y = mouse_pos.get_y()
    
    def move(self, task):
        dt = self.showbase.clock.dt
        
        if self.keyMap["cam-left"]:
            self.showbase.camera.setX(self.showbase.camera, -self.move_speed * dt)
        if self.keyMap["cam-right"]:
            self.showbase.camera.setX(self.showbase.camera, +self.move_speed * dt)
        if self.keyMap["cam-backward"]:
            self.showbase.camera.setY(self.showbase.camera, -self.move_speed * dt)
        if self.keyMap["cam-forward"]:
            self.showbase.camera.setY(self.showbase.camera, +self.move_speed * dt)
        if self.keyMap["cam-up"]:
            self.showbase.camera.setZ(self.showbase.camera, +self.move_speed * dt)
        if self.keyMap["cam-down"]:
            self.showbase.camera.setZ(self.showbase.camera, -self.move_speed * dt)
        if self.keyMap["mouse1"]:
            try:
                mouse_pos = self.showbase.mouseWatcherNode.getMouse()
                current_x = mouse_pos.get_x()
                current_y = mouse_pos.get_y()
                
                # 计算鼠标移动差值
                dx = current_x - self.last_mouse_x
                dy = current_y - self.last_mouse_y
                dy *= -1
                
                # 根据鼠标移动调整摄像机
                if abs(dx) > 0.01 or abs(dy) > 0.01:
                    cam_hpr = self.showbase.camera.get_hpr()
                    self.showbase.camera.set_hpr(
                        cam_hpr.x - dx * self.camera_speed,
                        max(-90, min(90, cam_hpr.y - dy * self.camera_speed)),
                        0
                    )
                
                # 更新鼠标位置
                self.last_mouse_x = current_x
                self.last_mouse_y = current_y
            except Exception as e:
                print(f"旋转相机失败：{e}")
        
        return task.cont