import weakref
import carb
import omni.kit.commands

from pxr import Sdf, Usd, UsdGeom, UsdPhysics
from .types import LayerMetaType
from omni.kit.widget.layers import LayerUtils


class ModifyStageAxisCommand(omni.kit.commands.Command):
    """Modify stage up axis undoable **Command**."""

    def __init__(self, stage, axis):
        if stage:
            self._stage = weakref.ref(stage)
        else:
            self._stage = None
        self._new_axis = axis
        self._old_axis = None

    def do(self):
        if self._stage and self._stage():
            self._old_axis = UsdGeom.GetStageUpAxis(self._stage())
            UsdGeom.SetStageUpAxis(self._stage(), self._new_axis)

    def undo(self):
        if self._stage and self._stage() and self._old_axis:
            UsdGeom.SetStageUpAxis(self._stage(), self._old_axis)


class ModifyLayerMetadataCommand(omni.kit.commands.Command):
    """Modify layer metadata undoable **Command**."""

    def __init__(self, layer_identifier, parent_layer_identifier, meta_index, value):
        """Constructor.

        Keyword Arguments:
        layer_identifier (str): Layer identifier to operate.
        parent_identifier (str): Parent identifier that layer_identifier resides in as sublayer.
        It's None if it's root layer.
        meta_index (omni.kit.property.layer.types.LayerMetaType): Metadata type.
        """

        super().__init__()
        self._layer_identifer = layer_identifier
        self._parent_layer_identifier = parent_layer_identifier
        self._meta_index = meta_index
        self._new_value = value
        self._old_value = None

    def _set_value(self, meta_type, value):
        layer = Sdf.Find(self._layer_identifer)
        if not layer:
            return

        if self._parent_layer_identifier:
            parent_layer = Sdf.Find(self._parent_layer_identifier)
        else:
            parent_layer = None

        try:
            if meta_type == LayerMetaType.COMMENT:
                self._old_value = layer.comment
                layer.comment = str(value)
            elif meta_type == LayerMetaType.DOC:
                self._old_value = layer.documentation
                layer.documentation = str(value)
            elif meta_type == LayerMetaType.START_TIME:
                self._old_value = layer.startTimeCode
                layer.startTimeCode = float(value)
            elif meta_type == LayerMetaType.END_TIME:
                self._old_value = layer.endTimeCode
                layer.endTimeCode = float(value)
            elif meta_type == LayerMetaType.TIMECODES_PER_SECOND:
                self._old_value = layer.timeCodesPerSecond
                layer.timeCodesPerSecond = float(value)
            elif meta_type == LayerMetaType.FPS_PER_SECOND:
                self._old_value = layer.framesPerSecond
                layer.framesPerSecond = float(value)
            elif meta_type == LayerMetaType.UNITS:
                stage = Usd.Stage.Open(layer, None, None, Usd.Stage.LoadNone)
                meters = UsdGeom.GetStageMetersPerUnit(stage)
                self._old_value = meters
                UsdGeom.SetStageMetersPerUnit(stage, float(value))
            elif meta_type == LayerMetaType.KG_PER_UNIT:
                stage = Usd.Stage.Open(layer, None, None, Usd.Stage.LoadNone)
                kilograms = UsdPhysics.GetStageKilogramsPerUnit(stage)
                self._old_value = kilograms
                UsdPhysics.SetStageKilogramsPerUnit(stage, float(value))
            elif parent_layer and meta_type == LayerMetaType.LAYER_OFFSET:
                layer_index = LayerUtils.get_sublayer_position_in_parent(
                    self._parent_layer_identifier, layer.identifier
                )
                offset = parent_layer.subLayerOffsets[layer_index]
                self._old_value = offset.offset
                parent_layer.subLayerOffsets[layer_index] = Sdf.LayerOffset(float(value), offset.scale)
            elif parent_layer and meta_type == LayerMetaType.LAYER_SCALE:
                layer_index = LayerUtils.get_sublayer_position_in_parent(
                    self._parent_layer_identifier, layer.identifier
                )
                offset = parent_layer.subLayerOffsets[layer_index]
                self._old_value = offset.scale
                parent_layer.subLayerOffsets[layer_index] = Sdf.LayerOffset(offset.offset, float(value))
        except Exception as e:
            pass

    def do(self):
        self._set_value(self._meta_index, self._new_value)

    def undo(self):
        self._set_value(self._meta_index, self._old_value)
