import carb
import omni.kit.commands
from typing import List, Optional, Any
from pxr import Usd, Sdf, UsdGeom


class PrimVarCommand(omni.kit.commands.Command):
    """
    Set primvar undoable **Command**.

    Args:
        prim_path (list): List of paths of prims.
        prim_name (str): Primvar name.
        prim_type (): Primvar variable type (EG. Sdf.ValueTypeNames.Bool)
        value (any): New primvar value. If primvar doesn't exist, it will be created
    """

    def __init__(
        self,
        prim_path: List[str],
        prim_name: str,
        prim_type: str,
        value: Any,
        usd_context_name: Optional[str] = "",
    ):
        self._prim_path = prim_path
        self._prim_name = prim_name
        self._prim_type = prim_type
        self._value = value
        self._usd_context = omni.usd.get_context(usd_context_name)
        self._undo_values = {}

    def do(self):
        stage = self._usd_context.get_stage()
        for path in self._prim_path:
            if path:
                primvars_api = UsdGeom.PrimvarsAPI(stage.GetPrimAtPath(path))
                value = primvars_api.GetPrimvar(self._prim_name)
                if value:
                    if value.GetTypeName() != self._prim_type:
                         carb.log_error(f"PrimVarCommand: cannot set value as {path}.{self._prim_name} is type {value.GetTypeName()} and expected type is {self._prim_type}")
                    else:
                        self._undo_values[str(path)] = value.Get()
                        value.Set(self._value)
                else:
                    self._undo_values[str(path)] = None
                    primvars_api.CreatePrimvar(self._prim_name, self._prim_type).Set(self._value)

    def undo(self):
        stage = self._usd_context.get_stage()
        for path in self._undo_values.keys():

            primvars_api = UsdGeom.PrimvarsAPI(stage.GetPrimAtPath(path))
            value = primvars_api.GetPrimvar(self._prim_name)
            orig_value = self._undo_values[path]
            if orig_value:
                value.Set(orig_value)
            else:
                primvars_api.RemovePrimvar(self._prim_name)

        self._undo_values = {}

class TogglePrimVarCommand(omni.kit.commands.Command):
    """
    Toggle primvar undoable **Command**.

    Args:
        prim_path (list): List of paths of prims.
        prim_name (str): Primvar name.
    """

    def __init__(
        self,
        prim_path: List[str],
        prim_name: str,
        usd_context_name: Optional[str] = "",
    ):
        self._prim_path = prim_path
        self._prim_name = prim_name
        self._usd_context = omni.usd.get_context(usd_context_name)
        self._undo_values = {}

    def do(self):
        stage = self._usd_context.get_stage()
        for path in self._prim_path:
            if path:
                primvars_api = UsdGeom.PrimvarsAPI(stage.GetPrimAtPath(path))
                value = primvars_api.GetPrimvar(self._prim_name)
                if value:
                    if value.GetTypeName() != Sdf.ValueTypeNames.Bool:
                        carb.log_error(f"TogglePrimVarCommand: cannot set value as {value.GetTypeName()} isn't a {self._prim_type}")
                    else:
                        self._undo_values[str(path)] = value.Get()
                        value.Set(not value.Get())
                else:
                    self._undo_values[path] = None
                    primvars_api.CreatePrimvar(self._prim_name, Sdf.ValueTypeNames.Bool).Set(True)

    def undo(self):
        stage = self._usd_context.get_stage()
        for path in self._undo_values.keys():

            primvars_api = UsdGeom.PrimvarsAPI(stage.GetPrimAtPath(path))
            value = primvars_api.GetPrimvar(self._prim_name)
            orig_value = self._undo_values[path]
            if orig_value:
                value.Set(orig_value)
            else:
                primvars_api.RemovePrimvar(self._prim_name)

        self._undo_values = {}


class ToggleInstanceableCommand(omni.kit.commands.Command):
    """
    Toggle instanceable undoable **Command**.

    Args:
        prim_path (list): List of paths of prims.
    """

    def __init__(
        self,
        prim_path: List[str],
        usd_context_name: Optional[str] = "",
    ):
        self._prim_path = prim_path
        self._usd_context = omni.usd.get_context(usd_context_name)
        self._undo_values = {}

    def do(self):
        stage = self._usd_context.get_stage()
        for path in self._prim_path:
            if path:
                prim = stage.GetPrimAtPath(path)
                value = prim.IsInstanceable()
                self._undo_values[str(path)] = value
                prim.SetInstanceable(not value)

    def undo(self):
        stage = self._usd_context.get_stage()
        for path in self._undo_values.keys():
            prim = stage.GetPrimAtPath(path)
            value = self._undo_values[path]
            prim.SetInstanceable(value)

        self._undo_values = {}
