from pxr import Sdf, Usd


class UsdLayerUndo:
    class Key:
        def __init__(self, path, info):
            self.path = path
            self.info = info

    def __init__(self, layer: Sdf.Layer):
        self._layer = layer
        self.reset()

    def reserve(self, path: Sdf.Path, info=None):
        # Check if it's included by any added paths
        path = Sdf.Path(path)
        for added_key in self._paths:
            if path.HasPrefix(added_key.path) and added_key.info is None:
                return

        # Check if it includes any added paths
        if info is None:
            for added_key in self._paths:
                if added_key.path.HasPrefix(path):
                    self._paths.pop(added_key)

        # If it doesn't exist, it's new spec. No need to reserve anything.
        key = self.Key(path, info)
        if not self._layer.GetObjectAtPath(path):
            self._paths[key] = True
            return
        else:
            self._paths[key] = False

        # Reserve existing data
        if info is None:
            Sdf.CreatePrimInLayer(self._reserve_layer, path.GetParentPath())
            Sdf.CopySpec(self._layer, path, self._reserve_layer, path)
        else:
            # We must know attribute type to create it in layer. So we don't handle it here.
            if path.IsPropertyPath():
                raise Exception("Property info is not supported.")

            spec = self._layer.GetObjectAtPath(path)
            Sdf.CreatePrimInLayer(self._reserve_layer, path)
            reserve_spec = self._reserve_layer.GetObjectAtPath(path)

            reserve_spec.SetInfo(info, spec.GetInfo(info))

    def undo(self):
        batch_edit = Sdf.BatchNamespaceEdit()
        for key, new_spec in self._paths.items():
            spec = self._layer.GetObjectAtPath(key.path)
            reserve_spec = self._reserve_layer.GetObjectAtPath(key.path)

            if new_spec:
                # Remove new added spec
                if spec:
                    batch_edit.Add(Sdf.NamespaceEdit.Remove(key.path))
            elif key.info is None:
                # Restore spec
                Sdf.CopySpec(self._reserve_layer, key.path, self._layer, key.path)
            else:
                # Restore spec info
                spec.SetInfo(key.info, reserve_spec.GetInfo(key.info))

        self._layer.Apply(batch_edit)

    def reset(self):
        self._reserve_layer = Sdf.Layer.CreateAnonymous()
        self._paths = {}


class UsdEditTargetUndo:
    def __init__(self, edit_target: Usd.EditTarget):
        self._edit_target = edit_target
        self._layer_undo = UsdLayerUndo(self._edit_target.GetLayer())

    def reserve(self, path: Sdf.Path, info=None):
        self._layer_undo.reserve(path, info)

    def undo(self):
        self._layer_undo.undo()

    def reset(self):
        self._layer_undo.reset()
