import asyncio
import carb
import os
import tempfile
import omni.kit.test
import omni.usd
import omni.client
import unittest

from stat import S_IREAD, S_IWRITE
from pxr import Sdf, Usd, UsdGeom, Gf
from .test_base import enable_server_tests
from omni.kit.usd.layers import get_layers, LayerUtils, get_layer_event_payload, LayerEventType


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

    # Before running each test
    async def setUp(self):
        self.usd_context = omni.usd.get_context()
        await omni.usd.get_context().new_stage_async()

    async def tearDown(self):
        if omni.usd.get_context().get_stage():
            await omni.usd.get_context().close_stage_async()

    async def test_layers_state_muteness_interfaces(self):
        layers = get_layers()
        layers_state = layers.get_layers_state()

        stage = omni.usd.get_context().get_stage()

        layer0 = Sdf.Layer.CreateAnonymous()
        layer1 = Sdf.Layer.CreateAnonymous()
        layer2 = Sdf.Layer.CreateAnonymous()
        stage.GetRootLayer().subLayerPaths.append(layer0.identifier)
        stage.GetRootLayer().subLayerPaths.append(layer1.identifier)
        stage.GetRootLayer().subLayerPaths.append(layer2.identifier)

        self.assertFalse(layers_state.is_muteness_global())

        stage.MuteLayer(layer0.identifier)
        await omni.kit.app.get_app().next_update_async()
        self.assertTrue(layers_state.is_layer_locally_muted(layer0.identifier))
        self.assertFalse(layers_state.is_layer_globally_muted(layer0.identifier))
        self.assertFalse(layers_state.is_layer_locally_muted(layer1.identifier))
        self.assertFalse(layers_state.is_layer_locally_muted(layer2.identifier))

        stage.UnmuteLayer(layer0.identifier)
        await omni.kit.app.get_app().next_update_async()
        self.assertFalse(layers_state.is_layer_locally_muted(layer0.identifier))

        # Mute layers before muteness scope switches.
        stage.MuteLayer(layer0.identifier)
        await omni.kit.app.get_app().next_update_async()
        self.assertTrue(layers_state.is_layer_locally_muted(layer0.identifier))
        layers_state.set_muteness_scope(True)
        self.assertTrue(layers_state.is_muteness_global())
        self.assertTrue(layers_state.is_layer_locally_muted(layer0.identifier))
        self.assertFalse(layers_state.is_layer_globally_muted(layer0.identifier))
        # Layer is not muted since it's muted in local mode.
        self.assertFalse(stage.IsLayerMuted(layer0.identifier))

        stage.MuteLayer(layer0.identifier)
        await omni.kit.app.get_app().next_update_async()
        # Local muteness flag is still true.
        self.assertTrue(layers_state.is_layer_locally_muted(layer0.identifier))
        self.assertTrue(layers_state.is_layer_globally_muted(layer0.identifier))
        self.assertTrue(stage.IsLayerMuted(layer0.identifier))
        stage.UnmuteLayer(layer0.identifier)
        await omni.kit.app.get_app().next_update_async()
        self.assertFalse(layers_state.is_layer_globally_muted(layer0.identifier))
        self.assertFalse(stage.IsLayerMuted(layer0.identifier))

        # Swiches back to local scope again.
        layers_state.set_muteness_scope(False)
        self.assertTrue(layers_state.is_layer_locally_muted(layer0.identifier))
        self.assertFalse(layers_state.is_layer_globally_muted(layer0.identifier))
        self.assertTrue(stage.IsLayerMuted(layer0.identifier))

        # Pass None for identifier will not cause any issues
        self.assertFalse(layers_state.is_layer_locally_muted(None))

    async def test_live_update_muteness(self):
        # The test is not to do real live update, but tries to simulate live update
        # by changing the custom data of muteness directly.
        layers = get_layers()
        layers_state = layers.get_layers_state()
        layers_state.set_muteness_scope(False)

        stage = omni.usd.get_context().get_stage()
        layer0 = Sdf.Layer.CreateAnonymous()
        stage.GetRootLayer().subLayerPaths.append(layer0.identifier)
        LayerUtils.set_layer_global_muteness(stage.GetRootLayer(), layer0.identifier, True)
        await omni.kit.app.get_app().next_update_async()
        self.assertTrue(layers_state.is_layer_globally_muted(layer0.identifier))
        self.assertFalse(stage.IsLayerMuted(layer0.identifier))

        layers_state.set_muteness_scope(True)
        await omni.kit.app.get_app().next_update_async()
        self.assertTrue(layers_state.is_layer_globally_muted(layer0.identifier))
        self.assertTrue(stage.IsLayerMuted(layer0.identifier))

    async def test_permissions_query_interface(self):
        layers = get_layers()
        layers_state = layers.get_layers_state()
        with tempfile.TemporaryDirectory() as tempdir:
            file = os.path.join(tempdir, "test.usd")
            layer = Sdf.Layer.CreateNew(file)
            layer.Save()
            layer = None
            os.chmod(file, S_IREAD)

            await omni.usd.get_context().open_stage_async(file)
            is_writable = layers_state.is_layer_writable(file)
            is_read_only_on_disk = layers_state.is_layer_readonly_on_disk(file)
            is_savable = layers_state.is_layer_savable(file)
            await omni.usd.get_context().new_stage_async()

            os.chmod(file, S_IWRITE)

        self.assertFalse(is_writable)
        self.assertTrue(is_read_only_on_disk)
        self.assertFalse(is_savable)
        self.assertFalse(layers_state.is_layer_locked(file))

        stage = omni.usd.get_context().get_stage()
        root_layer = stage.GetRootLayer()
        LayerUtils.set_layer_lock_status(root_layer, root_layer.identifier, True)
        # Root layer cannot be locked
        self.assertTrue(layers_state.is_layer_writable(root_layer.identifier))
        # Anonymous layer cannot be saved.
        self.assertFalse(layers_state.is_layer_savable(root_layer.identifier))
        self.assertFalse(layers_state.is_layer_locked(file))

        with tempfile.TemporaryDirectory() as tempdir:
            file = os.path.join(tempdir, "test.usd")
            layer = Sdf.Layer.CreateNew(file)
            root_layer.subLayerPaths.append(file)
            LayerUtils.set_layer_lock_status(root_layer, file, True)

            # Locked layer cannot be writable or savable.
            self.assertFalse(layers_state.is_layer_writable(file))
            # Even it's locked, it's still not readonly on disk.
            self.assertFalse(layers_state.is_layer_readonly_on_disk(file))
            self.assertFalse(layers_state.is_layer_savable(file))
            self.assertTrue(layers_state.is_layer_locked(file))

            # Close stage to release temp file
            layer = None
            stage = None
            root_layer = None
            await omni.usd.get_context().close_stage_async()

    async def test_set_get_layer_names(self):
        layers = get_layers()
        layers_state = layers.get_layers_state()
        stage = omni.usd.get_context().get_stage()
        root_layer = stage.GetRootLayer()
        session_layer = stage.GetSessionLayer()

        self.assertEqual(layers_state.get_layer_name(root_layer.identifier), "Root Layer")
        self.assertEqual(layers_state.get_layer_name(session_layer.identifier), "Session Layer")

        sublayer = Sdf.Layer.CreateAnonymous()
        root_layer.subLayerPaths.append(sublayer.identifier)
        layers_state.set_layer_name(sublayer.identifier, "test")
        self.assertEqual(layers_state.get_layer_name(sublayer.identifier), "test")

        self.assertEqual("abc.usd", layers_state.get_layer_name("c:/a/b/abc.usd"))
        self.assertEqual("abc.usda", layers_state.get_layer_name("c:/a/b/abc.usda"))
        self.assertEqual("abc.usda", layers_state.get_layer_name("omniverse://ov-invalid-fake-server/a/b/abc.usda"))
        layer = Sdf.Layer.CreateAnonymous()
        self.assertEqual(layer.identifier, layers_state.get_layer_name(layer.identifier))
        self.assertEqual("", layers_state.get_layer_name(""))
        self.assertEqual(
            "a b c.usda", layers_state.get_layer_name("omniverse://ov-invalid-fake-server/a/b/a%20b%20c.usda")
        )

    @unittest.skipIf(not enable_server_tests(), "")
    async def test_layer_refresh(self):
        layers = get_layers()
        layers_state = layers.get_layers_state()
        stage = omni.usd.get_context().get_stage()

        file = "omniverse://localhost/Projects/tests/omni.kit.usd.layers/test_layer_refresh/test.usd"
        layer = Sdf.Layer.FindOrOpen(file)
        if not layer:
            layer = Sdf.Layer.CreateNew(file)
            layer.Save()

        stage.GetRootLayer().subLayerPaths.append(layer.identifier)
        self.assertFalse(layers_state.is_layer_outdated(layer.identifier))

        custom_data = layer.customLayerData
        custom_data['test'] = 1
        layer.customLayerData = custom_data
        layer.Save()
        for i in range(100):
            await omni.kit.app.get_app().next_update_async()

        self.assertFalse(layers_state.is_layer_outdated(layer.identifier))
        layer.Reload()
        self.assertFalse(layers_state.is_layer_outdated(layer.identifier))

        custom_data = layer.customLayerData
        custom_data['test'] = 1
        layer.customLayerData = custom_data
        # Saves with customized interface will update latest state
        layer.Save()
        for i in range(100):
            await omni.kit.app.get_app().next_update_async()
        self.assertFalse(layers_state.is_layer_outdated(layer.identifier))

    async def test_layer_edit_target_changed(self):
        # Test layer edit target change if it's muted/locked/removed.
        layers = get_layers()
        layers_state = layers.get_layers_state()
        stage = omni.usd.get_context().get_stage()
        sublayer0 = Sdf.Layer.CreateAnonymous()
        root_layer = stage.GetRootLayer()
        LayerUtils.insert_sublayer(root_layer, 0, sublayer0.identifier, False)

        # Mute/lock/remove current edit target will change its edit target to root layer.
        LayerUtils.set_edit_target(stage, sublayer0.identifier)
        self.assertEqual(sublayer0.identifier, LayerUtils.get_edit_target(stage))
        stage.MuteLayer(sublayer0.identifier)
        await omni.kit.app.get_app().next_update_async()
        self.assertEqual(root_layer.identifier, LayerUtils.get_edit_target(stage))

        stage.UnmuteLayer(sublayer0.identifier)
        LayerUtils.set_edit_target(stage, sublayer0.identifier)
        self.assertEqual(sublayer0.identifier, LayerUtils.get_edit_target(stage))
        layers_state.set_layer_lock_state(sublayer0.identifier, True)
        self.assertEqual(root_layer.identifier, LayerUtils.get_edit_target(stage))

        layers_state.set_layer_lock_state(sublayer0.identifier, False)
        LayerUtils.set_edit_target(stage, sublayer0.identifier)
        self.assertEqual(sublayer0.identifier, LayerUtils.get_edit_target(stage))
        LayerUtils.remove_sublayer(root_layer, 0)
        self.assertEqual(root_layer.identifier, LayerUtils.get_edit_target(stage))

    async def test_layer_muteness_events(self):
        layers = get_layers()
        layers_state = layers.get_layers_state()
        event_stream = layers.get_event_stream()

        payload = None

        def _on_layer_event(event: carb.events.IEvent):
            nonlocal payload
            payload = get_layer_event_payload(event)

        subscription = event_stream.create_subscription_to_pop(_on_layer_event, name="layer events")
        layers_state = layers.get_layers_state()
        layers_state.set_muteness_scope(True)

        self.assertIsNotNone(payload)
        self.assertEqual(payload.event_type, LayerEventType.MUTENESS_SCOPE_CHANGED)

        stage = omni.usd.get_context().get_stage()
        sublayer0 = Sdf.Layer.CreateAnonymous()
        sublayer1 = Sdf.Layer.CreateAnonymous()
        root_layer = stage.GetRootLayer()
        LayerUtils.insert_sublayer(root_layer, 0, sublayer0.identifier, False)
        LayerUtils.insert_sublayer(root_layer, 1, sublayer1.identifier, False)
        await omni.kit.app.get_app().next_update_async()

        stage.MuteAndUnmuteLayers([sublayer0.identifier, sublayer1.identifier], [])
        await omni.kit.app.get_app().next_update_async()
        await omni.kit.app.get_app().next_update_async()

        self.assertIsNotNone(payload)
        self.assertEqual(payload.event_type, LayerEventType.MUTENESS_STATE_CHANGED)
        self.assertEqual(set(payload.identifiers_or_spec_paths), set([sublayer0.identifier, sublayer1.identifier]))

        stage.UnmuteLayer(sublayer1.identifier)
        await omni.kit.app.get_app().next_update_async()
        await omni.kit.app.get_app().next_update_async()
        self.assertIsNotNone(payload)
        self.assertEqual(payload.event_type, LayerEventType.MUTENESS_STATE_CHANGED)
        self.assertEqual(set(payload.identifiers_or_spec_paths), set([sublayer1.identifier]))

        layers_state.set_muteness_scope(False)
        # Switch scope will load local muteness.
        self.assertIsNotNone(payload)
        self.assertEqual(payload.event_type, LayerEventType.MUTENESS_STATE_CHANGED)
        # Since only sublayer0's state is changed.
        self.assertEqual(set(payload.identifiers_or_spec_paths), set([sublayer0.identifier]))

    async def test_layer_lock_events(self):
        layers = get_layers()
        layers_state = layers.get_layers_state()
        event_stream = layers.get_event_stream()

        payload = None

        def _on_layer_event(event: carb.events.IEvent):
            nonlocal payload
            payload = get_layer_event_payload(event)

        stage = omni.usd.get_context().get_stage()
        sublayer0 = Sdf.Layer.CreateAnonymous()
        sublayer1 = Sdf.Layer.CreateAnonymous()
        root_layer = stage.GetRootLayer()
        LayerUtils.insert_sublayer(root_layer, 0, sublayer0.identifier, False)
        LayerUtils.insert_sublayer(root_layer, 1, sublayer1.identifier, False)

        subscription = event_stream.create_subscription_to_pop(_on_layer_event, name="layer events")

        # Lock
        layers_state.set_layer_lock_state(sublayer0.identifier, True)
        await omni.kit.app.get_app().next_update_async()
        await omni.kit.app.get_app().next_update_async()
        self.assertIsNotNone(payload)
        self.assertEqual(payload.event_type, LayerEventType.LOCK_STATE_CHANGED)
        self.assertEqual(set(payload.identifiers_or_spec_paths), set([sublayer0.identifier]))

        # Unlock
        payload = None
        await omni.kit.app.get_app().next_update_async()
        await omni.kit.app.get_app().next_update_async()
        layers_state.set_layer_lock_state(sublayer0.identifier, False)
        await omni.kit.app.get_app().next_update_async()
        await omni.kit.app.get_app().next_update_async()
        self.assertIsNotNone(payload)
        self.assertEqual(payload.event_type, LayerEventType.LOCK_STATE_CHANGED)
        self.assertEqual(set(payload.identifiers_or_spec_paths), set([sublayer0.identifier]))

        payload = None
        with Sdf.ChangeBlock():
            LayerUtils.set_layer_lock_status(root_layer, sublayer0.identifier, True)
            LayerUtils.set_layer_lock_status(root_layer, sublayer1.identifier, True)

        await omni.kit.app.get_app().next_update_async()
        await omni.kit.app.get_app().next_update_async()
        self.assertIsNotNone(payload)
        self.assertEqual(payload.event_type, LayerEventType.LOCK_STATE_CHANGED)
        self.assertEqual(set(payload.identifiers_or_spec_paths), set([sublayer0.identifier, sublayer1.identifier]))

    async def test_layer_dirty_state_events(self):
        layers = get_layers()
        event_stream = layers.get_event_stream()

        payload = None

        def _on_layer_event(event: carb.events.IEvent):
            nonlocal payload
            temp = get_layer_event_payload(event)
            if temp.event_type == LayerEventType.DIRTY_STATE_CHANGED:
                payload = temp

        with tempfile.TemporaryDirectory() as tempdir:
            file = os.path.join(tempdir, "test.usd")
            file1 = os.path.join(tempdir, "test1.usd")
            layer = Sdf.Layer.CreateNew(file)
            layer.Save()

            layer1 = Sdf.Layer.CreateNew(file1)
            layer1.Save()

            stage = omni.usd.get_context().get_stage()
            root_layer = stage.GetRootLayer()
            LayerUtils.insert_sublayer(root_layer, 0, layer.identifier, False)
            LayerUtils.insert_sublayer(root_layer, 1, layer1.identifier, False)

            subscription = event_stream.create_subscription_to_pop(_on_layer_event, name="layer events")

            # Change anonymous layer will not emit dirty event
            with Usd.EditContext(stage, stage.GetRootLayer()):
                stage.DefinePrim("/root/prim")
                await omni.kit.app.get_app().next_update_async()
                self.assertIsNone(payload)

            with Usd.EditContext(stage, layer):
                stage.DefinePrim("/root/prim2")
                await omni.kit.app.get_app().next_update_async()
                self.assertIsNotNone(payload)
                self.assertEqual(payload.event_type, LayerEventType.DIRTY_STATE_CHANGED)
                self.assertEqual(set(payload.identifiers_or_spec_paths), set([layer.identifier]))
                self.assertTrue(layer.dirty)

                # Change a dirty layer will not emit dirty event again
                payload = None
                stage.DefinePrim("/root/prim3")
                self.assertIsNone(payload)

            with Usd.EditContext(stage, layer1):
                stage.DefinePrim("/root/prim4")
                await omni.kit.app.get_app().next_update_async()

            # Save dirty layers will emit dirty change event
            payload = None
            stage.Save()
            await omni.kit.app.get_app().next_update_async()
            self.assertIsNotNone(payload)
            self.assertEqual(payload.event_type, LayerEventType.DIRTY_STATE_CHANGED)
            self.assertEqual(set(payload.identifiers_or_spec_paths), set([layer.identifier, layer1.identifier]))
            self.assertFalse(layer.dirty)
            self.assertFalse(layer1.dirty)

            layer = None
            layer1 = None
            root_layer = None
            await omni.usd.get_context().new_stage_async()

    async def test_layer_outdate_state_events(self):
        pass

    async def test_layer_prim_specs_events(self):
        layers = get_layers()
        event_stream = layers.get_event_stream()

        payload = None

        def _on_layer_event(event: carb.events.IEvent):
            nonlocal payload
            temp = get_layer_event_payload(event)
            if temp.event_type == LayerEventType.PRIM_SPECS_CHANGED:
                payload = temp

        subscription = event_stream.create_subscription_to_pop(_on_layer_event, name="layer events")

        stage = omni.usd.get_context().get_stage()
        all_prims = []
        for i in range(100):
            prim = stage.DefinePrim(f"/test_{i}")
            all_prims.append(str(prim.GetPath()))

        await omni.kit.app.get_app().next_update_async()
        self.assertIsNotNone(payload)
        self.assertEqual(payload.event_type, LayerEventType.PRIM_SPECS_CHANGED)
        root_layer_identifier = stage.GetRootLayer().identifier
        self.assertIn(root_layer_identifier, payload.layer_spec_paths)
        self.assertEqual(set(payload.layer_spec_paths[root_layer_identifier]), set(all_prims))

        stage.GetRootLayer().Clear()
        await omni.kit.app.get_app().next_update_async()
        self.assertIsNotNone(payload)
        self.assertEqual(payload.event_type, LayerEventType.PRIM_SPECS_CHANGED)
        root_layer_identifier = stage.GetRootLayer().identifier
        self.assertIn(root_layer_identifier, payload.layer_spec_paths)
        self.assertEqual(set(payload.layer_spec_paths[root_layer_identifier]), set([str(Sdf.Path.absoluteRootPath)]))

    async def test_sublayers_changed_events(self):
        layers = get_layers()
        event_stream = layers.get_event_stream()

        payload = None

        def _on_layer_event(event: carb.events.IEvent):
            nonlocal payload
            temp = get_layer_event_payload(event)
            if temp.event_type == LayerEventType.SUBLAYERS_CHANGED:
                payload = temp

        subscription = event_stream.create_subscription_to_pop(_on_layer_event, name="layer events")

        stage = omni.usd.get_context().get_stage()
        root_layer = stage.GetRootLayer()
        sublayer0 = Sdf.Layer.CreateAnonymous()
        sublayer1 = Sdf.Layer.CreateAnonymous()

        root_layer.subLayerPaths.append(sublayer0.identifier)
        await omni.kit.app.get_app().next_update_async()
        self.assertIsNotNone(payload)
        self.assertEqual(payload.event_type, LayerEventType.SUBLAYERS_CHANGED)

        payload = None
        LayerUtils.remove_sublayer(root_layer, 0)
        await omni.kit.app.get_app().next_update_async()
        self.assertIsNotNone(payload)
        self.assertEqual(payload.event_type, LayerEventType.SUBLAYERS_CHANGED)

        root_layer.subLayerPaths.append(sublayer0.identifier)
        root_layer.subLayerPaths.append(sublayer1.identifier)

        # Layer move
        payload = None
        LayerUtils.move_layer(root_layer.identifier, 0, sublayer1.identifier, 0, True)
        await omni.kit.app.get_app().next_update_async()
        self.assertIsNotNone(payload)
        self.assertEqual(payload.event_type, LayerEventType.SUBLAYERS_CHANGED)

        # Not all edits will emit sublayer changes events.
        payload = None
        stage.DefinePrim("/root/test")
        self.assertIsNone(payload)

    async def test_layer_info_changed_events(self):
        layers = get_layers()
        event_stream = layers.get_event_stream()

        payload = None

        def _on_layer_event(event: carb.events.IEvent):
            nonlocal payload
            temp = get_layer_event_payload(event)
            if temp.event_type == LayerEventType.INFO_CHANGED:
                payload = temp

        subscription = event_stream.create_subscription_to_pop(_on_layer_event, name="layer events")

        stage = omni.usd.get_context().get_stage()
        root_layer = stage.GetRootLayer()
        sublayer0 = Sdf.Layer.CreateAnonymous()
        root_layer.subLayerPaths.append(sublayer0.identifier)
        root_layer.subLayerOffsets[0] = Sdf.LayerOffset(2.0, 3.0)

        up_axis = UsdGeom.GetStageUpAxis(stage)
        with Sdf.ChangeBlock():
            if up_axis == UsdGeom.Tokens.z:
                UsdGeom.SetStageUpAxis(stage, UsdGeom.Tokens.y)
            else:
                UsdGeom.SetStageUpAxis(stage, UsdGeom.Tokens.z)

            root_layer.comment = "test"
            root_layer.documentation = "test2"

            UsdGeom.SetStageMetersPerUnit(stage, 100.0)

            stage.SetStartTimeCode(100.99)
            stage.SetEndTimeCode(1000.99)
            root_layer.framesPerSecond = 100
            stage.SetTimeCodesPerSecond(1000)

        await omni.kit.app.get_app().next_update_async()
        self.assertIsNotNone(payload)
        self.assertEqual(payload.event_type, LayerEventType.INFO_CHANGED)
        self.assertIn(root_layer.identifier, payload.layer_info_data)
        expected_changes = (
            "documentation",
            "upAxis",
            "comment",
            "metersPerUnit",
            "startTimeCode",
            "framesPerSecond",
            "endTimeCode",
            "timeCodesPerSecond",
        )
        self.assertEqual(set(payload.layer_info_data[root_layer.identifier]), set(expected_changes))

    async def test_edit_target_changed(self):
        layers = get_layers()
        event_stream = layers.get_event_stream()

        payload = None

        def _on_layer_event(event: carb.events.IEvent):
            nonlocal payload
            temp = get_layer_event_payload(event)
            if temp.event_type == LayerEventType.EDIT_TARGET_CHANGED:
                payload = temp

        subscription = event_stream.create_subscription_to_pop(_on_layer_event, name="layer events")

        stage = omni.usd.get_context().get_stage()
        root_layer = stage.GetRootLayer()
        sublayer0 = Sdf.Layer.CreateAnonymous()
        root_layer.subLayerPaths.append(sublayer0.identifier)

        edit_target = Usd.EditTarget(sublayer0)
        stage.SetEditTarget(edit_target)
        await omni.kit.app.get_app().next_update_async()
        self.assertIsNotNone(payload)
        self.assertEqual(payload.event_type, LayerEventType.EDIT_TARGET_CHANGED)

        # Set the same edit target will not trigger any events
        payload = None
        stage.SetEditTarget(edit_target)
        await omni.kit.app.get_app().next_update_async()
        self.assertIsNone(payload)

    def check_sublayers(self, sublayer_paths, expected_layer_identifiers):
        sublayer_paths = sorted(sublayer_paths)
        expected_layer_identifiers = sorted(expected_layer_identifiers)
        self.assertTrue(
            sublayer_paths == expected_layer_identifiers,
            f"Sublayers array does not match, got: {sublayer_paths}, expected: {expected_layer_identifiers}",
        )

    async def test_get_used_sublayers(self):
        layers = get_layers()
        stage = self.usd_context.get_stage()
        root_layer = stage.GetRootLayer()

        sublayers = layers.get_layers_state().get_local_layer_identifiers()
        self.check_sublayers(sublayers, [root_layer.identifier])
        layer0 = Sdf.Layer.CreateAnonymous()
        LayerUtils.insert_sublayer(root_layer, 0, layer0.identifier)
        sublayers = layers.get_layers_state().get_local_layer_identifiers()
        self.check_sublayers(sublayers, [root_layer.identifier, layer0.identifier])
        layer1 = Sdf.Layer.CreateAnonymous()
        LayerUtils.insert_sublayer(root_layer, 1, layer1.identifier)
        sublayers = layers.get_layers_state().get_local_layer_identifiers()
        self.check_sublayers(sublayers, [root_layer.identifier, layer0.identifier, layer1.identifier])

        layer2 = Sdf.Layer.CreateAnonymous()
        layer3 = Sdf.Layer.CreateAnonymous()
        layer4 = Sdf.Layer.CreateAnonymous()
        layer5 = Sdf.Layer.CreateAnonymous()
        LayerUtils.insert_sublayer(layer2, 0, layer3.identifier)
        LayerUtils.insert_sublayer(layer2, 1, layer4.identifier)
        LayerUtils.insert_sublayer(layer4, 0, layer5.identifier)

        LayerUtils.insert_sublayer(root_layer, 2, layer2.identifier)
        sublayers = layers.get_layers_state().get_local_layer_identifiers()
        self.check_sublayers(
            sublayers,
            [
                root_layer.identifier,
                layer0.identifier,
                layer1.identifier,
                layer2.identifier,
                layer3.identifier,
                layer4.identifier,
                layer5.identifier,
            ],
        )

        # Removes layer0
        LayerUtils.remove_sublayer(root_layer, 0)
        sublayers = layers.get_layers_state().get_local_layer_identifiers()
        self.check_sublayers(
            sublayers,
            [
                root_layer.identifier,
                layer1.identifier,
                layer2.identifier,
                layer3.identifier,
                layer4.identifier,
                layer5.identifier,
            ],
        )

        # Removes layer2 will remove layer2, layer3, layer4, layer5
        LayerUtils.remove_sublayer(root_layer, 1)
        sublayers = layers.get_layers_state().get_local_layer_identifiers()
        self.check_sublayers(sublayers, [root_layer.identifier, layer1.identifier])

    async def test_get_dirty_sublayers(self):
        layers = get_layers()
        stage = self.usd_context.get_stage()
        root_layer = stage.GetRootLayer()

        with tempfile.TemporaryDirectory() as tempdir:
            format = Sdf.FileFormat.FindByExtension(".usd")
            layer0 = Sdf.Layer.New(format, f"{tempdir}/1.usd")
            layer1 = Sdf.Layer.New(format, f"{tempdir}/2.usd")
            layer2 = Sdf.Layer.New(format, f"{tempdir}/3.usd")
            layer3 = Sdf.Layer.New(format, f"{tempdir}/4.usd")
            layer4 = Sdf.Layer.New(format, f"{tempdir}/5.usd")
            layer5 = Sdf.Layer.New(format, f"{tempdir}/6.usd")
            LayerUtils.insert_sublayer(root_layer, 0, layer0.identifier, False)
            LayerUtils.insert_sublayer(root_layer, 0, layer1.identifier, False)
            LayerUtils.insert_sublayer(layer2, 0, layer3.identifier, False)
            LayerUtils.insert_sublayer(layer2, 0, layer4.identifier, False)
            LayerUtils.insert_sublayer(layer4, 0, layer5.identifier, False)
            LayerUtils.insert_sublayer(root_layer, 0, layer2.identifier, False)
            sublayers = layers.get_layers_state().get_local_layer_identifiers()
            self.check_sublayers(
                sublayers,
                [
                    root_layer.identifier,
                    layer0.identifier,
                    layer1.identifier,
                    layer2.identifier,
                    layer3.identifier,
                    layer4.identifier,
                    layer5.identifier,
                ],
            )

            # Checkes dirtiness of layers since layer2 and layer4 have been touched.
            # They should be dirty at this moment
            dirty_sublayers = layers.get_layers_state().get_dirty_layer_identifiers()
            self.check_sublayers(dirty_sublayers, [layer2.identifier, layer4.identifier])

            # Touches layer1
            LayerUtils.set_edit_target(stage, layer1.identifier)
            UsdGeom.Mesh.Define(stage, "/root/test")
            dirty_sublayers = layers.get_layers_state().get_dirty_layer_identifiers()
            self.check_sublayers(dirty_sublayers, [layer1.identifier, layer2.identifier, layer4.identifier])

