import carb
import omni.kit.test

import omni.kit.undo
import omni.kit.commands
import omni.usd
from pxr import Sdf, Usd


def get_stage_default_prim_path(stage):
    if stage.HasDefaultPrim():
        return stage.GetDefaultPrim().GetPath()
    else:
        return Sdf.Path.absoluteRootPath


class TestCommandUsd(omni.kit.test.AsyncTestCase):

    async def test_toggle_payload_selected(self):
        carb.log_info("Test TogglePayLoadLoadSelectedPrimsCommand")
        await omni.usd.get_context().new_stage_async()
        usd_context = omni.usd.get_context()
        selection = usd_context.get_selection()
        stage = usd_context.get_stage()
        default_prim_path = get_stage_default_prim_path(stage)

        payload1 = Usd.Stage.CreateInMemory("payload1.usd")
        payload1.DefinePrim("/payload1/scope1", "Xform")
        payload1.DefinePrim("/payload1/scope1/xform", "Cube")

        payload2 = Usd.Stage.CreateInMemory("payload2.usd")
        payload2.DefinePrim("/payload2/scope2", "Xform")
        payload2.DefinePrim("/payload2/scope2/xform", "Cube")

        payload3 = Usd.Stage.CreateInMemory("payload3.usd")
        payload3.DefinePrim("/payload3/scope3", "Xform")
        payload3.DefinePrim("/payload3/scope3/xform", "Cube")

        payload4 = Usd.Stage.CreateInMemory("payload4.usd")
        payload4.DefinePrim("/payload4/scope4", "Xform")
        payload4.DefinePrim("/payload4/scope4/xform", "Cube")

        ps1 = stage.DefinePrim(default_prim_path.AppendChild("payload1"), "Xform")
        ps1.GetPayloads().AddPayload(
            Sdf.Payload(payload1.GetRootLayer().identifier, "/payload1"))
        ps2 = stage.DefinePrim(default_prim_path.AppendChild("payload2"), "Xform")
        ps2.GetPayloads().AddPayload(
            Sdf.Payload(payload2.GetRootLayer().identifier, "/payload2"))
        ps3 = stage.DefinePrim(default_prim_path.AppendChild("payload3"), "Xform")
        ps3.GetPayloads().AddPayload(
            Sdf.Payload(payload3.GetRootLayer().identifier, "/payload3"))
        ps4 = stage.DefinePrim(ps3.GetPath().AppendChild("payload4"), "Xform")
        ps4.GetPayloads().AddPayload(
            Sdf.Payload(payload4.GetRootLayer().identifier, "/payload4"))

        # unload everything
        stage.Unload()

        self.assertTrue(not ps1.IsLoaded())
        self.assertTrue(not ps2.IsLoaded())
        self.assertTrue(not ps3.IsLoaded())
        self.assertTrue(not ps4.IsLoaded())

        # if nothing selected, payload state should not change.
        selection.clear_selected_prim_paths()
        paths = selection.get_selected_prim_paths()
        omni.kit.commands.execute("TogglePayLoadLoadSelectedPrims", selected_paths=paths)
        self.assertTrue(not ps1.IsLoaded())
        self.assertTrue(not ps2.IsLoaded())
        self.assertTrue(not ps3.IsLoaded())
        self.assertTrue(not ps4.IsLoaded())

        # load payload1
        selection.set_selected_prim_paths(
            [
                ps1.GetPath().pathString
            ],
            False,
        )
        paths = selection.get_selected_prim_paths()
        omni.kit.commands.execute("TogglePayLoadLoadSelectedPrims", selected_paths=paths)
        self.assertTrue(ps1.IsLoaded())

        # unload payload1
        omni.kit.commands.execute("TogglePayLoadLoadSelectedPrims", selected_paths=paths)
        self.assertTrue(not ps1.IsLoaded())

        # load payload1, 2 and 3. 4 will load
        selection.set_selected_prim_paths(
            [
                ps1.GetPath().pathString,
                ps2.GetPath().pathString,
                ps3.GetPath().pathString,
            ],
            False,
        )
        paths = selection.get_selected_prim_paths()
        omni.kit.commands.execute("TogglePayLoadLoadSelectedPrims", selected_paths=paths)
        self.assertTrue(ps1.IsLoaded())
        self.assertTrue(ps2.IsLoaded())
        self.assertTrue(ps3.IsLoaded())
        self.assertTrue(ps4.IsLoaded())

        # unload 4
        selection.set_selected_prim_paths(
            [
                ps4.GetPath().pathString
            ],
            False,
        )
        paths = selection.get_selected_prim_paths()
        omni.kit.commands.execute("TogglePayLoadLoadSelectedPrims", selected_paths=paths)
        self.assertTrue(not ps4.IsLoaded())

        # undo
        omni.kit.undo.undo()
        self.assertTrue(ps4.IsLoaded())

        # redo
        omni.kit.undo.redo()
        self.assertTrue(not ps4.IsLoaded())

    async def test_set_payload_selected(self):
        carb.log_info("Test SetPayLoadLoadSelectedPrimsCommand")
        await omni.usd.get_context().new_stage_async()
        usd_context = omni.usd.get_context()
        selection = usd_context.get_selection()
        stage = usd_context.get_stage()
        default_prim_path = get_stage_default_prim_path(stage)

        payload1 = Usd.Stage.CreateInMemory("payload1.usd")
        payload1.DefinePrim("/payload1/scope1", "Xform")
        payload1.DefinePrim("/payload1/scope1/xform", "Cube")

        payload2 = Usd.Stage.CreateInMemory("payload2.usd")
        payload2.DefinePrim("/payload2/scope2", "Xform")
        payload2.DefinePrim("/payload2/scope2/xform", "Cube")

        payload3 = Usd.Stage.CreateInMemory("payload3.usd")
        payload3.DefinePrim("/payload3/scope3", "Xform")
        payload3.DefinePrim("/payload3/scope3/xform", "Cube")

        payload4 = Usd.Stage.CreateInMemory("payload4.usd")
        payload4.DefinePrim("/payload4/scope4", "Xform")
        payload4.DefinePrim("/payload4/scope4/xform", "Cube")

        ps1 = stage.DefinePrim(default_prim_path.AppendChild("payload1"), "Xform")
        ps1.GetPayloads().AddPayload(
            Sdf.Payload(payload1.GetRootLayer().identifier, "/payload1"))
        ps2 = stage.DefinePrim(default_prim_path.AppendChild("payload2"), "Xform")
        ps2.GetPayloads().AddPayload(
            Sdf.Payload(payload2.GetRootLayer().identifier, "/payload2"))
        ps3 = stage.DefinePrim(default_prim_path.AppendChild("payload3"), "Xform")
        ps3.GetPayloads().AddPayload(
            Sdf.Payload(payload3.GetRootLayer().identifier, "/payload3"))
        ps4 = stage.DefinePrim(ps3.GetPath().AppendChild("payload4"), "Xform")
        ps4.GetPayloads().AddPayload(
            Sdf.Payload(payload4.GetRootLayer().identifier, "/payload4"))

        # unload everything
        stage.Unload()

        self.assertTrue(not ps1.IsLoaded())
        self.assertTrue(not ps2.IsLoaded())
        self.assertTrue(not ps3.IsLoaded())
        self.assertTrue(not ps4.IsLoaded())

        # if nothing selected, payload state should not change.
        selection.clear_selected_prim_paths()
        paths = selection.get_selected_prim_paths()
        omni.kit.commands.execute("SetPayLoadLoadSelectedPrims", selected_paths=paths, value=True)
        self.assertTrue(not ps1.IsLoaded())
        self.assertTrue(not ps2.IsLoaded())
        self.assertTrue(not ps3.IsLoaded())
        self.assertTrue(not ps4.IsLoaded())

        # load payload1
        selection.set_selected_prim_paths(
            [
                ps1.GetPath().pathString
            ],
            False,
        )
        paths = selection.get_selected_prim_paths()
        omni.kit.commands.execute("SetPayLoadLoadSelectedPrims", selected_paths=paths, value=True)
        self.assertTrue(ps1.IsLoaded())
        omni.kit.commands.execute("SetPayLoadLoadSelectedPrims", selected_paths=paths, value=True)
        self.assertTrue(ps1.IsLoaded())

        # unload payload1
        omni.kit.commands.execute("SetPayLoadLoadSelectedPrims", selected_paths=paths, value=False)
        self.assertTrue(not ps1.IsLoaded())
        omni.kit.commands.execute("SetPayLoadLoadSelectedPrims", selected_paths=paths, value=False)
        self.assertTrue(not ps1.IsLoaded())

        # load payload1, 2 and 3. 4 will load
        selection.set_selected_prim_paths(
            [
                ps1.GetPath().pathString,
                ps2.GetPath().pathString,
                ps3.GetPath().pathString,
            ],
            False,
        )
        paths = selection.get_selected_prim_paths()
        omni.kit.commands.execute("SetPayLoadLoadSelectedPrims", selected_paths=paths, value=True)
        self.assertTrue(ps1.IsLoaded())
        self.assertTrue(ps2.IsLoaded())
        self.assertTrue(ps3.IsLoaded())
        self.assertTrue(ps4.IsLoaded())

        selection.set_selected_prim_paths(
            [
                ps4.GetPath().pathString
            ],
            False,
        )
        paths = selection.get_selected_prim_paths()
        # reload 4
        omni.kit.commands.execute("SetPayLoadLoadSelectedPrims", selected_paths=paths, value=True)
        self.assertTrue(ps4.IsLoaded())
        # unload 4
        omni.kit.commands.execute("SetPayLoadLoadSelectedPrims", selected_paths=paths, value=False)
        self.assertTrue(not ps4.IsLoaded())

        # undo
        omni.kit.undo.undo()
        self.assertTrue(ps4.IsLoaded())

        # redo
        omni.kit.undo.redo()
        self.assertTrue(not ps4.IsLoaded())

        omni.kit.commands.execute("SetPayLoadLoadSelectedPrims", selected_paths=paths, value=False)
        self.assertTrue(not ps4.IsLoaded())
        omni.kit.commands.execute("SetPayLoadLoadSelectedPrims", selected_paths=paths, value=False)
        self.assertTrue(not ps4.IsLoaded())
        omni.kit.undo.undo()
        self.assertTrue(not ps4.IsLoaded())
        omni.kit.undo.redo()
        self.assertTrue(not ps4.IsLoaded())

        omni.kit.commands.execute("SetPayLoadLoadSelectedPrims", selected_paths=paths, value=True)  # -1
        self.assertTrue(ps4.IsLoaded())
        omni.kit.commands.execute("SetPayLoadLoadSelectedPrims", selected_paths=paths, value=True)  # 0
        self.assertTrue(ps4.IsLoaded())
        omni.kit.undo.undo()
        self.assertTrue(ps4.IsLoaded())
        omni.kit.undo.redo()
        self.assertTrue(ps4.IsLoaded())

        omni.kit.commands.execute("SetPayLoadLoadSelectedPrims", selected_paths=paths, value=False)  # 1
        self.assertTrue(not ps4.IsLoaded())  # 1
        omni.kit.undo.undo()
        self.assertTrue(ps4.IsLoaded())  # 0
        omni.kit.undo.redo()
        self.assertTrue(not ps4.IsLoaded())  # 1
        omni.kit.undo.undo()
        self.assertTrue(ps4.IsLoaded())  # 0

        # triple undo
        omni.kit.commands.execute("SetPayLoadLoadSelectedPrims", selected_paths=paths, value=True)  # 2
        omni.kit.commands.execute("SetPayLoadLoadSelectedPrims", selected_paths=paths, value=True)  # 3
        omni.kit.commands.execute("SetPayLoadLoadSelectedPrims", selected_paths=paths, value=False)  # 4
        omni.kit.commands.execute("SetPayLoadLoadSelectedPrims", selected_paths=paths, value=True)  # 5
        self.assertTrue(ps4.IsLoaded())  # 5
        omni.kit.undo.undo()
        self.assertTrue(not ps4.IsLoaded())  # 4
        omni.kit.undo.undo()
        self.assertTrue(ps4.IsLoaded())  # 3
        omni.kit.undo.undo()
        self.assertTrue(ps4.IsLoaded())  # 2

        # more undo
        omni.kit.undo.undo()
        self.assertTrue(ps4.IsLoaded())  # 0
        omni.kit.undo.undo()
        self.assertTrue(ps4.IsLoaded())  # -1
