import omni.kit.commands
import omni.usd
from typing import List
from pxr import Sdf


class ParentPrimsCommand(omni.kit.commands.Command):

    def __init__(
        self,
        parent_path: str,
        child_paths: List[str],
        on_move_fn: callable = None,
        keep_world_transform: bool = True
    ):
        """
        Move prims into children of "parent" primitives undoable **Command**.

        Args:
            parent_path: prim path to become parent of child_paths
            child_paths: prim paths to become children of parent_prim
            keep_world_transform: If it needs to keep the world transform after parenting.
        """
        self._parent_path = parent_path
        self._child_paths = child_paths.copy()
        self._on_move_fn = on_move_fn
        self._keep_world_transform = keep_world_transform

    def do(self):
        with omni.kit.undo.group():
            for path in self._child_paths:
                path_to = self._parent_path + "/" + Sdf.Path(path).name
                omni.kit.commands.execute(
                    "MovePrim",
                    path_from=path,
                    path_to=path_to,
                    on_move_fn=self._on_move_fn,
                    destructive=False,
                    keep_world_transform=self._keep_world_transform
                )

    def undo(self):
        pass


class UnparentPrimsCommand(omni.kit.commands.Command):

    def __init__(
        self,
        paths: List[str],
        on_move_fn: callable = None,
        keep_world_transform: bool = True
    ):
        """
        Move prims into "/" primitives undoable **Command**.

        Args:
            paths: prim path to become parent of child_paths
            keep_world_transform: If it needs to keep the world transform after parenting.
        """
        self._paths = paths.copy()
        self._on_move_fn = on_move_fn
        self._keep_world_transform = keep_world_transform

    def do(self):
        with omni.kit.undo.group():
            for path in self._paths:
                path_to = "/" + Sdf.Path(path).name
                omni.kit.commands.execute(
                    "MovePrim",
                    path_from=path,
                    path_to=path_to,
                    on_move_fn=self._on_move_fn,
                    destructive=False,
                    keep_world_transform=self._keep_world_transform
                )

    def undo(self):
        pass


omni.kit.commands.register_all_commands_in_module(__name__)
