from __future__ import annotations

import importlib
import importlib.util
from typing import Callable, Optional, List, Dict, Any

from direct.task import Task
from panda3d.core import NodePath
from direct.showbase.DirectObject import DirectObject
from direct.showbase.ShowBase import ShowBase

HAS_IMGUI = False
if importlib.util.find_spec("imgui_bundle"):
    from imgui_bundle import imgui
    HAS_IMGUI = True

from .move_gizmo import MoveGizmo
from .rotate_gizmo import RotateGizmo
from .scale_gizmo import ScaleGizmo
from .events import GizmoEvent
import panda3d.core as p3d


class TransformGizmoMode:
    """Simple string constants for gizmo modes."""

    NONE = "none"
    MOVE = "move"
    ROTATE = "rotate"
    SCALE = "scale"  # reserved for future implementation
    ALL = "all"      # move + rotate together (matches existing UI semantics)


class TransformGizmo(DirectObject):
    """
    A unified Transform Gizmo wrapper for Panda3D.

    It internally manages:
    - MoveGizmo  (translation)
    - RotateGizmo (rotation)

    and exposes a single API to:
    - attach / detach to a target NodePath
    - switch transform mode: move / rotate / all / none
    - forward undo operations for move / rotate

    Typical usage:

        world = Panda3DWorld()
        gizmo = TransformGizmo(world)
        gizmo.set_mode(TransformGizmoMode.MOVE)
        gizmo.attach(some_node)

        # Later:
        gizmo.set_mode(TransformGizmoMode.ROTATE)
        gizmo.attach(other_node)
    """

    def __init__(
        self,
        world: ShowBase = None,
        camera_np: Optional[NodePath] = None,
        on_action_handler_changed: Optional[Callable[[str], None]] = None,
        event_hooks: Optional[Dict[str, Dict[str, List[Callable[[Dict[str, Any]], None]]]]] = None
    ):
        super().__init__()

        if not world:
            from direct.showbase.ShowBaseGlobal import base
            world = base
        else: self.world = world
        self.__debug: bool = False
        self.enabled: bool = True
        self.on_action_handler_changed = on_action_handler_changed

        # Try to reuse the same camera for all internal gizmos
        self.camera: Optional[NodePath] = camera_np or getattr(
            world, "cam", None)
        if not self.camera and hasattr(world, "base"):
            self.camera = world.base.cam

        # Global history stack of transform actions pushed by sub gizmos.
        # Each entry is a dict, e.g.:
        #   {"kind": "move", "node": node, "old_pos": Vec3, "new_pos": Vec3}
        #   {"kind": "rotate", "node": node, "old_hpr": Vec3, "new_hpr": Vec3}
        self._history: List[Dict[str, Any]] = []
        # Redo stack for transform actions (cleared on new action commit).
        self._redo_history: List[Dict[str, Any]] = []
        self._event_hooks = event_hooks or {}

        # Internal gizmos – they report actions back via a callback so that
        # TransformGizmo can build a unified, time-ordered undo stack.
        self.overlay_cam = None
        self.last_cam_transform: Optional[p3d.LMatrix4f] = None
        self._setup_overlay_rendering()
        
        self.move_gizmo = MoveGizmo(
            self.world,
            self.overlay_cam_np,
            on_action_committed=self._on_subgizmo_action,
            event_hooks=self._extract_event_hooks(TransformGizmoMode.MOVE)
        )
        self.rotate_gizmo = RotateGizmo(
            self.world,
            self.overlay_cam_np,
            on_action_committed=self._on_subgizmo_action,
            event_hooks=self._extract_event_hooks(TransformGizmoMode.ROTATE)
        )
        self.scale_gizmo = ScaleGizmo(
            self.world,
            self.overlay_cam_np,
            on_action_committed=self._on_subgizmo_action,
            event_hooks=self._extract_event_hooks(TransformGizmoMode.SCALE)
        )

        self.move_gizmo.root.reparentTo(self.gizmo_render)
        self.rotate_gizmo.root.reparentTo(self.gizmo_render)
        self.scale_gizmo.root.reparentTo(self.gizmo_render)

        # Keep debug flag consistent across all gizmos
        self.move_gizmo.debug = self.__debug
        self.rotate_gizmo.debug = self.__debug
        self.scale_gizmo.debug = self.__debug

        # Current attached target and mode
        self.target_node: Optional[NodePath] = None
        self.__mode: str = TransformGizmoMode.MOVE
        # Fine-grained switch for each handle type; combined with mode below.
        self._handle_enabled: Dict[str, bool] = {
            TransformGizmoMode.MOVE: True,
            TransformGizmoMode.ROTATE: True,
            TransformGizmoMode.SCALE: True,
        }
        # Track right mouse button so that when RMB is held
        # (used by CameraOrbitController for free-look + WASD fly),
        # we ignore W/E/R transform hotkeys, matching Unity behaviour.
        self._rmb_down: bool = False

        # Undo/redo shortcuts are handled by app-level action routing.
        # Avoid double-trigger with main.py bindings.
        # Mode switch hotkeys (Unity style), but disabled while RMB is pressed.
        self.accept("w", self._on_key_w)
        self.accept("e", self._on_key_e)
        self.accept("r", self._on_key_r)
        # Track right mouse button state.
        self.accept("mouse2", self._on_mouse2_down)
        self.accept("mouse2-up", self._on_mouse2_up)
        self.world.task_mgr.add(self._update_task, "transform_gizmo_update")
        
    def _setup_overlay_rendering(self):
        """
        创建独立的3D渲染场景用于Gizmo，完全绕过RenderPipeline后处理。
        
        原理：
        - 创建一个独立的render场景(gizmo_render)，不属于主render场景
        - 创建一个Sort值极高的DisplayRegion，在RP完成后渲染
        - 由于Gizmo不在主render下，RenderPipeline完全不会处理它
        - overlay相机与主相机同步变换，保持透视效果一致
        """
        if not hasattr(self.world, 'win') or self.world.win is None:
            raise RuntimeError("No world showbase.win found")
        # 创建独立的Gizmo渲染场景（不在主render下，RP不会处理）
        self.gizmo_render = NodePath("gizmo_render")
        # 创建高Sort值的DisplayRegion
        self.overlay_dr = self.world.win.makeDisplayRegion()
        self.overlay_dr.setSort(8)  # 极高值，确保在RP后处理之后
        
        self.overlay_dr.setClearColorActive(False)  # 透明背景，保留3D画面
        self.overlay_dr.setClearDepthActive(True)   # 清除深度，Gizmo始终在前
        # 创建overlay专用相机
        self.overlay_cam = p3d.Camera("gizmo_overlay_cam")
        self.overlay_cam.setLens(self.world.camLens)  # 共享主相机镜头
        self.overlay_cam.setScene(self.gizmo_render)  # 渲染gizmo_render场景
        # 创建相机节点（独立于主相机，但会同步变换）
        self.overlay_cam_np = self.gizmo_render.attachNewNode(self.overlay_cam)
        self.overlay_dr.setCamera(self.overlay_cam_np)
        # 每帧同步overlay相机位置到主相机
        from direct.task.TaskManagerGlobal import taskMgr
        taskMgr.add(self._sync_overlay_camera, "GizmoOverlayCameraSync")
        
    def _sync_overlay_camera(self, task:Task):
        """每帧同步overlay相机变换到主相机。"""
        curr = self.world.camera.getTransform(self.world.render)
        if self.overlay_cam_np and self.last_cam_transform != curr:
            # 从主相机获取世界坐标变换，应用到overlay相机
            self.overlay_cam_np.setTransform(curr)
            self.last_cam_transform = curr
        return task.cont
    
    def _check_node_empty(self,task:Task):
        if self.target_node is None or self.target_node.parent is None or self.target_node.isEmpty():
            self.move_gizmo.detach()
            self.rotate_gizmo.detach()
            self.scale_gizmo.detach()
            self.target_node = None
            return task.done
        return task.cont
    
    def _update_task(self,task:Task):
        # if get_state().hover_imgui: return task.cont
        self.move_gizmo.update()
        self.rotate_gizmo.update()
        self.scale_gizmo.update()
        return task.cont

    @property
    def is_hovering(self)->bool:
        return self.move_gizmo.is_hovering or self.rotate_gizmo.is_hovering or self.scale_gizmo.is_hovering
    
    # ------------------------------------------------------------------ #
    # Public API
    # ------------------------------------------------------------------ #
    def set_debug(self, enabled: bool) -> None:
        """Enable / disable debug logs for all internal gizmos."""
        self.__debug = bool(enabled)
        self.move_gizmo.debug = self.__debug
        self.rotate_gizmo.debug = self.__debug
        self.scale_gizmo.debug = self.__debug

    def attach(self, node_path: NodePath) -> None:
        """
        Attach the transform gizmo to a target node.

        Only the gizmos corresponding to the current mode will be attached.
        """
        if not node_path or not self.enabled:
            return

        self.target_node = node_path
        self._apply_mode_to_subgizmos()
        self.world.task_mgr.remove("CheckGizmoNodeEmpty")
        self.world.task_mgr.add(self._check_node_empty, "CheckGizmoNodeEmpty")

    def detach(self) -> None:
        """
        Detach from the current target and hide all gizmos.
        """
        self.target_node = None
        # Detach all sub gizmos to ensure they stop listening to events
        self.move_gizmo.detach()
        self.rotate_gizmo.detach()
        self.scale_gizmo.detach()

    def set_handles_enabled(
        self,
        *,
        move: Optional[bool] = False,
        rotate: Optional[bool] = False,
        scale: Optional[bool] = False,
    ) -> None:
        """
        Enable / disable each handle type independently of the current mode.

        Args:
            move:   True/False to enable/disable move handle; None keeps current.
            rotate: True/False to enable/disable rotate handle; None keeps current.
            scale:  True/False to enable/disable scale handle; None keeps current.

        Example:
            gizmo.set_mode(TransformGizmoMode.ALL)
            gizmo.set_handles_enabled(move=True, rotate=False, scale=False)
            # => only move handle stays active even though mode is ALL.
        """
        changed = False
        if move is not None:
            new_val = bool(move)
            if self._handle_enabled[TransformGizmoMode.MOVE] != new_val:
                self._handle_enabled[TransformGizmoMode.MOVE] = new_val
                changed = True
        if rotate is not None:
            new_val = bool(rotate)
            if self._handle_enabled[TransformGizmoMode.ROTATE] != new_val:
                self._handle_enabled[TransformGizmoMode.ROTATE] = new_val
                changed = True
        if scale is not None:
            new_val = bool(scale)
            if self._handle_enabled[TransformGizmoMode.SCALE] != new_val:
                self._handle_enabled[TransformGizmoMode.SCALE] = new_val
                changed = True

        # Re-attach handles to reflect the new mask if we already have a target.
        if changed and self.target_node is not None:
            self._apply_mode_to_subgizmos()

    def get_handles_enabled(self) -> Dict[str, bool]:
        """Return current per-handle enable states."""
        return dict(self._handle_enabled)

    def set_event_hooks(self, event_hooks: Dict[str, Dict[str, List[Callable[[Dict[str, Any]], None]]]], replace: bool = False) -> None:
        """Set or merge runtime drag event hooks for move/rotate/scale sub gizmos."""
        if not isinstance(event_hooks, dict):
            return
        if replace or not isinstance(self._event_hooks, dict):
            self._event_hooks = {}
        for key, sub in event_hooks.items():
            if not isinstance(sub, dict):
                continue
            key_l = str(key).lower()
            if replace or key_l not in self._event_hooks:
                self._event_hooks[key_l] = {}
            for event_name, callbacks in sub.items():
                if not callbacks:
                    continue
                if replace:
                    self._event_hooks[key_l][event_name] = list(callbacks)
                else:
                    existing = self._event_hooks[key_l].setdefault(event_name, [])
                    for cb in callbacks:
                        if cb not in existing:
                            existing.append(cb)

        # Apply updated hooks to existing sub gizmos immediately.
        self.move_gizmo._event_hooks = self._extract_event_hooks(TransformGizmoMode.MOVE)
        self.rotate_gizmo._event_hooks = self._extract_event_hooks(TransformGizmoMode.ROTATE)
        self.scale_gizmo._event_hooks = self._extract_event_hooks(TransformGizmoMode.SCALE)

    def set_mode(self, mode: str) -> None:
        """
        Set current transform mode.

        Supported values (see TransformGizmoMode):
            - "move"
            - "rotate"
            - "all"    (move + rotate)
            - "none"
            - "scale"  (reserved, not implemented yet)
        """
        mode = (mode or "").lower()
        if mode not in {
            TransformGizmoMode.NONE,
            TransformGizmoMode.MOVE,
            TransformGizmoMode.ROTATE,
            TransformGizmoMode.SCALE,
            TransformGizmoMode.ALL,
        }:
            raise ValueError(f"Unsupported gizmo mode: {mode}")

        if mode == self.__mode:
            return

        self.__mode = mode
        self._apply_mode_to_subgizmos()
        if self.on_action_handler_changed is not None:
            self.on_action_handler_changed(self.__mode)

    def get_mode(self) -> str:
        """Return current mode as a string."""
        return self.__mode

    def undo_last(self) -> bool:
        """
        Undo the most recent transform operation (move or rotate),
        regardless of current mode.
        """
        if not self._history:
            return False

        action = self._history.pop()
        kind = action.get("kind")
        node: NodePath = action.get("node")
        if node is None or node.isEmpty():
            return False

        if kind == "move":
            old_pos = action.get("old_pos")
            if old_pos is not None:
                node.setPos(self.world.render, old_pos)
        elif kind == "rotate":
            old_hpr = action.get("old_hpr")
            if old_hpr is not None:
                node.setHpr(self.world.render, old_hpr)
        elif kind == "scale":
            old_scale = action.get("old_scale")
            if old_scale is not None:
                node.setScale(old_scale)
        self._sync_light_position_if_needed(node)
        self._redo_history.append(action)
        return True

    def redo_last(self) -> bool:
        """
        Redo the most recently undone transform operation.
        """
        if not self._redo_history:
            return False

        action = self._redo_history.pop()
        kind = action.get("kind")
        node: NodePath = action.get("node")
        if node is None or node.isEmpty():
            return False

        if kind == "move":
            new_pos = action.get("new_pos")
            if new_pos is not None:
                node.setPos(self.world.render, new_pos)
        elif kind == "rotate":
            new_hpr = action.get("new_hpr")
            if new_hpr is not None:
                node.setHpr(self.world.render, new_hpr)
        elif kind == "scale":
            new_scale = action.get("new_scale")
            if new_scale is not None:
                node.setScale(new_scale)

        self._sync_light_position_if_needed(node)
        self._history.append(action)
        return True

    def update(self) -> None:
        """
        Forward update to internal gizmos.

        Both MoveGizmo and RotateGizmo already have their own Panda3D
        tasks to keep screen-size scaling; this is provided mainly for
        explicit/manual calls if needed.
        """
        self.move_gizmo.update()
        self.rotate_gizmo.update()
        self.scale_gizmo.update()

    # ------------------------------------------------------------------ #
    # Internal helpers
    # ------------------------------------------------------------------ #
    def _extract_event_hooks(self, key: str) -> Dict[str, List[Callable[[Dict[str, Any]], None]]]:
        """Return per-gizmo event hooks (drag_start / drag_move / drag_end)."""
        base = {name: [] for name in GizmoEvent.ALL}
        if not self._event_hooks:
            return base
        sub = self._event_hooks.get(key) or self._event_hooks.get(key.lower())
        if not isinstance(sub, dict):
            return base
        for name in list(base.keys()):
            cbs = sub.get(name)
            if cbs:
                base[name] = list(cbs)
        return base

    def _apply_mode_to_subgizmos(self) -> None:
        """
        Attach / detach internal gizmos based on current mode and target.
        """
        # Always detach first to ensure we don't keep stale event listeners.
        self.move_gizmo.detach()
        self.rotate_gizmo.detach()
        self.scale_gizmo.detach()

        if self.target_node is None:
            return

        if self.__mode in (TransformGizmoMode.MOVE, TransformGizmoMode.ALL) and self._handle_enabled.get(TransformGizmoMode.MOVE, True):
            self.move_gizmo.attach(self.target_node)

        if self.__mode in (TransformGizmoMode.ROTATE, TransformGizmoMode.ALL) and self._handle_enabled.get(TransformGizmoMode.ROTATE, True):
            self.rotate_gizmo.attach(self.target_node)

        if self.__mode in (TransformGizmoMode.SCALE, TransformGizmoMode.ALL) and self._handle_enabled.get(TransformGizmoMode.SCALE, True):
            self.scale_gizmo.attach(self.target_node)

    def _on_subgizmo_action(self, action: Dict[str, Any]) -> None:
        """
        Callback used by MoveGizmo / RotateGizmo / ScaleGizmo to report that a
        transform action (move/rotate/scale) has been committed.
        """
        self._history.append(action)
        # New user action invalidates redo chain.
        self._redo_history.clear()

    def _sync_light_position_if_needed(self, node: Optional[NodePath]) -> None:
        """When target node wraps an RP light, keep RP light position in sync."""
        try:
            if node is None or node.isEmpty() or (not node.hasPythonTag("rp_light_object")):
                return
            light_obj = node.getPythonTag("rp_light_object")
            if not light_obj:
                return
            world_pos = node.getPos(self.world.render)
            try:
                light_obj.setPos(world_pos)
            except Exception:
                try:
                    light_obj.setPos(world_pos.x, world_pos.y, world_pos.z)
                except Exception:
                    try:
                        light_obj.pos = world_pos
                    except Exception:
                        pass
        except Exception:
            pass

    # ------------------------------------------------------------------ #
    # Input helpers (hotkeys / mouse states)
    # ------------------------------------------------------------------ #
    def _on_mouse2_down(self, extra=None) -> None:
        """Right mouse button pressed: used by camera controller for fly mode."""
        self._rmb_down = True

    def _on_mouse2_up(self, extra=None) -> None:
        """Right mouse button released."""
        self._rmb_down = False

    def _on_key_w(self, *args) -> None:
        """Switch to Move mode unless RMB is held (camera fly)."""
        if self._rmb_down:
            return
        # if HAS_IMGUI:
        #     io = imgui.get_io()
        #     if io.want_capture_mouse: return
        # if get_state().hover_imgui:return
        self.set_mode(TransformGizmoMode.MOVE)

    def _on_key_e(self, *args) -> None:
        """Switch to Rotate mode unless RMB is held (camera fly)."""
        if self._rmb_down:
            return
        # if HAS_IMGUI:
        #     io = imgui.get_io()
        #     if io.want_capture_mouse: return
        # if get_state().hover_imgui:return
        self.set_mode(TransformGizmoMode.ROTATE)

    def _on_key_r(self, *args) -> None:
        """Switch to Scale mode unless RMB is held (camera fly)."""
        if self._rmb_down:
            return
        # if HAS_IMGUI:
        #     io = imgui.get_io()
        #     if io.want_capture_mouse: return
        # if get_state().hover_imgui:return
        self.set_mode(TransformGizmoMode.SCALE)


__all__ = ["TransformGizmo", "TransformGizmoMode"]
