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


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

    def __init__(self, selected_paths: List[str]):
        """
        Toggles the load/unload payload of the selected primitives undoable **Command**.

        Args:
            selected_paths: Old selected prim paths.
        """
        self._stage = omni.usd.get_context().get_stage()
        self._selected_paths = selected_paths.copy()

    def _toggle_load(self):
        for selected_path in self._selected_paths:
            selected_prim = self._stage.GetPrimAtPath(selected_path)
            if selected_prim.IsLoaded():
                selected_prim.Unload()
            else:
                selected_prim.Load()

    def do(self):
        self._toggle_load()

    def undo(self):
        self._toggle_load()


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

    def __init__(self, selected_paths: List[str], value: bool):
        """
        Set the load/unload payload of the selected primitives undoable **Command**.

        Args:
            selected_paths: Old selected prim paths.
            value: True = load, False = unload
        """
        self._stage = omni.usd.get_context().get_stage()
        self._selected_paths = selected_paths.copy()
        self._processed_path = set()
        self._value = value
        self._is_undo = False

    def _set_load(self):
        if self._is_undo:
            paths = self._processed_path
        else:
            paths = self._selected_paths
        for selected_path in paths:
            selected_prim = self._stage.GetPrimAtPath(selected_path)
            if (selected_prim.IsLoaded() and self._value) or (not selected_prim.IsLoaded() and not self._value):
                if selected_path in self._processed_path:
                    self._processed_path.remove(selected_path)
                continue
            if self._value:
                selected_prim.Load()
            else:
                selected_prim.Unload()
            self._processed_path.add(selected_path)

    def do(self):
        self._set_load()

    def undo(self):
        self._is_undo = True
        self._value = not self._value
        self._set_load()
        self._value = not self._value
        self._processed_path = set()
        self._is_undo = False


omni.kit.commands.register_all_commands_in_module(__name__)
