import omni.kit.test
from ..layer_undo import *
from pxr import Sdf, Usd, UsdGeom, Kind, Gf


class TestUsdUndo(omni.kit.test.AsyncTestCase):
    async def test_layer_undo(self):
        stage = Usd.Stage.CreateInMemory()
        usd_undo = UsdEditTargetUndo(stage.GetEditTarget())

        mesh = UsdGeom.Mesh.Define(stage, "/root")
        mesh.CreateNormalsAttr().Set([Gf.Vec3f(1, 0, 0)])
        mesh.SetNormalsInterpolation("constant")

        model = Usd.ModelAPI(mesh)
        model.SetKind(Kind.Tokens.group)

        org_stage_str = stage.ExportToString()
        print("Original stage:")
        print(org_stage_str)

        usd_undo.reserve("/root.normals")
        mesh.GetNormalsAttr().Set([Gf.Vec3f(0, 1, 0)])
        mesh.SetNormalsInterpolation("faceVarying")

        usd_undo.reserve("/root", "kind")
        model.SetKind(Kind.Tokens.component)

        usd_undo.reserve("/root/newPrim")
        stage.DefinePrim("/root/newPrim")

        modified_stage_str = stage.ExportToString()
        print("Modified stage:")
        print(modified_stage_str)

        assert org_stage_str != modified_stage_str

        usd_undo.undo()
        undone_stage_str = stage.ExportToString()
        print("Undone stage:")
        print(undone_stage_str)

        assert org_stage_str == undone_stage_str
