import carb
import omni.kit.test

import omni.kit.undo
import omni.kit.commands
import omni.timeline
import omni.usd
import tempfile
import omni.client

from omni.kit.widget.layers import LayerUtils
from pxr import Sdf, UsdGeom
from .base import TestLayerNonUIBase


class TestCppBindings(TestLayerNonUIBase):
    async def setUp(self):
        await super().setUp()

        self.previous_retry_values = omni.client.set_retries(0, 0, 0)

    async def tearDown(self):
        omni.client.set_retries(*self.previous_retry_values)

        await super().tearDown()

    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_layers_state_apis(self):
        layers = self.usd_context.get_layers()
        stage = self.usd_context.get_stage()
        self.assertFalse(layers.is_layer_muteness_global())
        layers.set_layer_muteness_scope(True)
        self.assertTrue(layers.is_layer_muteness_global())
        layers.set_layer_muteness_scope(False)
        self.assertFalse(layers.is_layer_muteness_global())

        layer2 = Sdf.Layer.CreateAnonymous()
        LayerUtils.insert_sublayer(stage.GetRootLayer(), 0, layer2.identifier)

        self.assertFalse(layers.is_layer_locally_muted(layer2.identifier))
        self.assertFalse(layers.is_layer_globally_muted(layer2.identifier))
        stage = self.usd_context.get_stage()
        stage.MuteLayer(layer2.identifier)
        await omni.kit.app.get_app().next_update_async()
        layers.set_layer_muteness_scope(False)
        self.assertTrue(layers.is_layer_locally_muted(layer2.identifier))
        self.assertFalse(layers.is_layer_globally_muted(layer2.identifier))
        self.assertFalse(omni.usd.is_layer_globally_muted(self.usd_context, layer2.identifier))
        layers.set_layer_muteness_scope(True)
        self.assertTrue(layers.is_layer_muteness_global())
        self.assertTrue(layers.is_layer_locally_muted(layer2.identifier))
        self.assertFalse(layers.is_layer_globally_muted(layer2.identifier))

        self.assertEqual(layers.get_layer_edit_mode(), omni.usd.LayerEditMode.NORMAL)
        layers.set_layer_edit_mode(omni.usd.LayerEditMode.AUTO_AUTHORING)
        self.assertEqual(layers.get_layer_edit_mode(), omni.usd.LayerEditMode.AUTO_AUTHORING)
        layers.set_layer_edit_mode(omni.usd.LayerEditMode.NORMAL)
        self.assertEqual(layers.get_layer_edit_mode(), omni.usd.LayerEditMode.NORMAL)

        self.assertFalse(layers.is_layer_locked(layer2.identifier))
        layers.set_layer_lock_state(layer2.identifier, True)
        self.assertTrue(layers.is_layer_locked(layer2.identifier))
        layers.set_layer_lock_state(layer2.identifier, False)
        self.assertFalse(layers.is_layer_locked(layer2.identifier))
        self.assertFalse(layers.is_layer_locked_by_other(layer2.identifier))
        self.assertEqual(layers.get_layer_lock_user_name(layer2.identifier), "")
        self.assertTrue(layers.is_layer_writable(layer2.identifier))
        self.assertFalse(layers.is_layer_savable(layer2.identifier))

        layers.set_layer_edit_mode(omni.usd.LayerEditMode.AUTO_AUTHORING)
        layers.set_default_edit_layer_identifier(layer2.identifier)
        self.assertEqual(layers.get_default_edit_layer_identifier(), layer2.identifier)
        self.assertFalse(layers.is_auto_authoring_layer(layer2.identifier))

        edit_target = stage.GetEditTarget()
        self.assertTrue(layers.is_auto_authoring_layer(edit_target.GetLayer().identifier))
        with omni.usd.active_authoring_layer_context(self.usd_context):
            edit_target = stage.GetEditTarget()
            self.assertFalse(layers.is_auto_authoring_layer(edit_target.GetLayer().identifier))
            self.assertEqual(edit_target.GetLayer().identifier, layer2.identifier)

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

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

        sublayers = layers.get_used_sublayers()
        self.check_sublayers(sublayers, [root_layer.identifier])
        layer0 = Sdf.Layer.CreateAnonymous()
        LayerUtils.insert_sublayer(root_layer, 0, layer0.identifier)
        sublayers = layers.get_used_sublayers()
        self.check_sublayers(sublayers, [root_layer.identifier, layer0.identifier])
        layer1 = Sdf.Layer.CreateAnonymous()
        LayerUtils.insert_sublayer(root_layer, 1, layer1.identifier)
        sublayers = layers.get_used_sublayers()
        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_used_sublayers()
        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_used_sublayers()
        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_used_sublayers()
        self.check_sublayers(sublayers, [root_layer.identifier, layer1.identifier])

    async def test_get_dirty_sublayers(self):
        usd_context = omni.usd.get_context()
        layers = usd_context.get_layers()
        stage = 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_used_sublayers()
            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_dirty_sublayers()
            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_dirty_sublayers()
            self.check_sublayers(dirty_sublayers, [layer1.identifier, layer2.identifier, layer4.identifier])
