import carb
import omni.kit.test
import omni.usd
import omni.client
import unittest
import omni.kit.app

from omni.kit.usd.layers import get_layers, get_layer_event_payload, LayerEventType, LayerUtils
from .test_base import enable_server_tests
from pxr import Usd, Sdf


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

    # Before running each test
    async def setUp(self):
        self.app = omni.kit.app.get_app()
        self.usd_context = omni.usd.get_context()
        self.layers = get_layers(self.usd_context)
        await omni.usd.get_context().new_stage_async()
        
        if enable_server_tests():
            self.test_folder = "omniverse://localhost/Projects/omni.kit.usd.layers/test_live_syncing/"
            await omni.client.delete_async(self.test_folder)
            self.stage_url = self.test_folder + "test.usd"
            self.sublayer = self.test_folder + "test2.usd"
            self.sublayer2 = self.test_folder + "test3.usd"
            stage = Usd.Stage.CreateNew(self.stage_url)
            sublayer = Sdf.Layer.CreateNew(self.sublayer)
            stage.GetRootLayer().subLayerPaths.append(self.sublayer)
            await self.usd_context.attach_stage_async(stage)
            self.stage = self.usd_context.get_stage()
        else:
            self.test_folder = ""
            self.stage_url = ""
            self.stage = None

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

        if enable_server_tests():
            await omni.client.delete_async(self.test_folder)

    async def wait(self, frames=10):
        for i in range(frames):
            await self.app.next_update_async()
    
    async def test_other_interfaces(self):
        layers = get_layers()
        live_syncing = layers.get_live_syncing()
        live_syncing.mute_live_session_merge_notice("omniverse://fake-test-server/a.usd")
        live_syncing.mute_live_session_merge_notice("omniverse://fake-test-server/b.usd")
        self.assertTrue(live_syncing.is_live_session_merge_notice_muted("omniverse://fake-test-server/a.usd"))
        self.assertTrue(live_syncing.is_live_session_merge_notice_muted("omniverse://fake-test-server/b.usd"))
        live_syncing.unmute_live_session_merge_notice("omniverse://fake-test-server/a.usd")
        self.assertFalse(live_syncing.is_live_session_merge_notice_muted("omniverse://fake-test-server/a.usd"))
        self.assertTrue(live_syncing.is_live_session_merge_notice_muted("omniverse://fake-test-server/b.usd"))

    @unittest.skipIf(not enable_server_tests(), "")
    async def test_join_session_with_url(self):
        layers = get_layers()
        live_syncing = layers.get_live_syncing()
        session_url = omni.client.combine_urls(self.stage_url, ".live/test.live/customized_session.live")
        self.assertFalse(live_syncing.join_live_session_by_url(self.stage_url, session_url, False))
        self.assertTrue(live_syncing.join_live_session_by_url(self.stage_url, session_url, True))
        current_live_session = live_syncing.get_current_live_session()
        self.assertTrue(current_live_session)
        self.assertTrue(live_syncing.is_stage_in_live_session())
        self.assertTrue(live_syncing.is_in_live_session())
        self.assertTrue(live_syncing.is_layer_in_live_session(self.stage_url))
        self.assertEqual(current_live_session.name, "customized_session")
        self.assertEqual(current_live_session.url, session_url)
        self.assertTrue(live_syncing.get_live_session_by_url(session_url))
        live_syncing.stop_all_live_sessions()

    @unittest.skipIf(not enable_server_tests(), "")
    async def test_session_management(self):
        for layer_identifier in [self.sublayer, self.stage_url]:
            await self.__test_session_management_internal(layer_identifier, False)
    
    @unittest.skipIf(not enable_server_tests(), "")
    async def test_session_management_merge_to_new_layer(self):
            await self.__test_session_management_internal(self.sublayer, True)
    
    async def __test_session_management_internal(self, sublayer_identifier, merge_to_new_layer=False):
        layers = get_layers()
        live_syncing = layers.get_live_syncing()
        is_root_layer = sublayer_identifier == self.usd_context.get_stage_url()

        event_stream = layers.get_event_stream()

        payload = None
        merge_started = None
        merge_ended = None

        def _on_layer_event(event: carb.events.IEvent):
            nonlocal payload
            nonlocal merge_started
            nonlocal merge_ended
            temp = get_layer_event_payload(event)
            if temp.event_type == LayerEventType.LIVE_SESSION_LIST_CHANGED or temp.event_type == LayerEventType.LIVE_SESSION_STATE_CHANGED:
                payload = temp
            elif temp.event_type == LayerEventType.LIVE_SESSION_MERGE_STARTED:
                merge_started = temp
            elif temp.event_type == LayerEventType.LIVE_SESSION_MERGE_ENDED:
                merge_ended = temp
        
        subscription = event_stream.create_subscription_to_pop(_on_layer_event, name="layer events")
        
        self.assertEqual(len(live_syncing.get_all_live_sessions(sublayer_identifier)), 0)

        live_session = live_syncing.create_live_session("test", layer_identifier=sublayer_identifier)
        self.assertTrue(live_session)
        self.assertEqual(live_session.name, "test")
        self.assertTrue(live_session.merge_permission)
        self.assertTrue(live_session.channel_url)
        self.assertTrue(live_session.url)
        self.assertTrue(live_session.root)
        self.assertTrue(live_session.owner)

        sessions = live_syncing.get_all_live_sessions(sublayer_identifier)
        self.assertEqual(len(sessions), 1)
        session = sessions[0]
        self.assertEqual(live_session.name, session.name)
        self.assertEqual(live_session.owner, session.owner)
        self.assertEqual(live_session.url, session.url)
        self.assertEqual(live_session.root, session.root)
        self.assertEqual(live_session.channel_url, session.channel_url)
        self.assertEqual(live_session.merge_permission, session.merge_permission)

        await self.wait()
        self.assertTrue(payload)
        self.assertEqual(payload.event_type, LayerEventType.LIVE_SESSION_LIST_CHANGED)

        live_session = live_syncing.create_live_session("test", layer_identifier=sublayer_identifier)
        # Create session with the same name will fail.
        self.assertFalse(live_session)
        
        # Create session with empty name will fail.
        live_session = live_syncing.create_live_session("", layer_identifier=sublayer_identifier)
        self.assertFalse(live_session)

        live_session = live_syncing.create_live_session("test2", layer_identifier=sublayer_identifier)
        sessions = live_syncing.get_all_live_sessions(sublayer_identifier)
        self.assertEqual(len(sessions), 2)
        live_syncing.stop_live_session(sublayer_identifier)
        await self.wait()

        payload = None
        self.assertTrue(live_syncing.join_live_session(live_session))
        current_session = live_syncing.get_current_live_session(sublayer_identifier)
        self.assertTrue(current_session)
        self.assertTrue(current_session.name, live_session.name)
        result, server_info = await omni.client.get_server_info_async(sublayer_identifier)
        self.assertTrue(result == omni.client.Result.OK)
        self.assertEqual(server_info.username, current_session.logged_user_name)
        self.assertEqual(server_info.connection_id, current_session.logged_user_id)

        self.assertTrue(live_syncing.is_layer_in_live_session(sublayer_identifier))
        if is_root_layer:
            self.assertTrue(self.usd_context.is_stage_live())
            self.assertEqual(self.usd_context.get_stage_live_mode(), omni.usd.StageLiveModeType.ALWAYS_ON)
        await self.wait()
        self.assertTrue(payload)
        self.assertEqual(payload.event_type, LayerEventType.LIVE_SESSION_STATE_CHANGED)

        all_sublayers = set(LayerUtils.get_all_sublayers(self.stage, is_root_layer, True, False))
        all_sublayers.discard(self.stage_url)
        all_sublayers.discard(self.sublayer)
        all_sublayers.discard(self.sublayer2)
        for sublayer in all_sublayers:
            self.assertTrue(live_syncing.is_live_session_layer(sublayer))

        all_sublayers = set(live_syncing.get_current_live_session_layers(sublayer_identifier))
        all_sublayers.discard(self.sublayer2)
        self.assertEqual(len(all_sublayers), 1)
        self.assertTrue(live_syncing.is_live_session_layer(all_sublayers.pop()))

        # Makes some edits
        root_live_layer = Sdf.Find(current_session.root)
        self.assertTrue(root_live_layer)

        layer_content = """\
            #usda 1.0
            (
                subLayers = [
                    @omniverse://test-ov-fake-server/fake/path/sublayers/level2/sublayer2.usd@
                ]
            )

            over "root"
            {
                over "test_prim" (
                    prepend payload = @../invalid/payload2.usd@
                    prepend references = @../invalid/reference2.usd@
                )
                {
                    asset[] AssetArray = [@1.usd@, @2.usd@, @OmniPBR.mdl@]
                    bool[] BoolArray = [0, 1]
                    float[] FloatArray = [1, 2, 3]
                    int[] IntArray = [1, 2, 3]
                    string[] StringArray = ["string1", "string2", "string3"]
                    int test = 3
                    asset test2 = @../../invalid/path.usd@
                    asset test3 = @OmniPBR.mdl@
                    token[] TokenArray = ["token1", "token2", "token3"]
                }
            }
        """

        if is_root_layer:
            expected_string = """
#usda 1.0
(
    customLayerData = {
        dictionary omni_layer = {
            dictionary locked = {
                bool "./test2.usd" = 0
            }
        }
    }
    metersPerUnit = 0.01
    subLayers = [
        @omniverse://localhost/Projects/omni.kit.usd.layers/test_live_syncing/test2.usd@
    ]
)

over "root"
{
    over "test_prim" (
        prepend payload = @./.live/test.live/invalid/payload2.usd@
        prepend references = @./.live/test.live/invalid/reference2.usd@
    )
    {
        asset[] AssetArray = [@1.usd@, @2.usd@, @OmniPBR.mdl@]
        bool[] BoolArray = [0, 1]
        float[] FloatArray = [1, 2, 3]
        int[] IntArray = [1, 2, 3]
        string[] StringArray = ["string1", "string2", "string3"]
        int test = 3
        asset test2 = @./.live/invalid/path.usd@
        asset test3 = @OmniPBR.mdl@
        token[] TokenArray = ["token1", "token2", "token3"]
    }
}
            """
        else:
            expected_string = """
#usda 1.0

over "root"
{
    over "test_prim" (
        prepend payload = @./.live/test2.live/invalid/payload2.usd@
        prepend references = @./.live/test2.live/invalid/reference2.usd@
    )
    {
        asset[] AssetArray = [@1.usd@, @2.usd@, @OmniPBR.mdl@]
        bool[] BoolArray = [0, 1]
        float[] FloatArray = [1, 2, 3]
        int[] IntArray = [1, 2, 3]
        string[] StringArray = ["string1", "string2", "string3"]
        int test = 3
        asset test2 = @./.live/invalid/path.usd@
        asset test3 = @OmniPBR.mdl@
        token[] TokenArray = ["token1", "token2", "token3"]
    }
}
            """

        root_live_layer.ImportFromString(layer_content)
        payload = None
        merge_started = None
        merge_ended = None
        if merge_to_new_layer:
            new_sublayer = Sdf.Layer.CreateNew(self.sublayer2)
            live_syncing.merge_live_session_changes_to_specific_layer(sublayer_identifier, self.sublayer2, True, True)
            all_sublayers = set(LayerUtils.get_all_sublayers(self.stage, False, True, False))
            self.assertIn(self.sublayer2, all_sublayers)
            layer_identifier = sublayer_identifier
        elif is_root_layer:
            live_syncing.merge_changes_to_base_layers(True)
            layer_identifier = self.stage.GetRootLayer().identifier
        else:
            live_syncing.merge_live_session_changes(sublayer_identifier, True)
            layer_identifier = sublayer_identifier
            
        self.assertTrue(merge_started)
        self.assertTrue(merge_ended)
        self.assertTrue(merge_started.layer_identifier == layer_identifier)
        self.assertTrue(merge_ended.layer_identifier == layer_identifier)
        self.assertTrue(merge_ended.success)
 
        current_session = live_syncing.get_current_live_session(sublayer_identifier)
        self.assertFalse(current_session)
        if is_root_layer:
            self.assertFalse(live_syncing.is_in_live_session())
        else:
            self.assertFalse(live_syncing.is_layer_in_live_session(sublayer_identifier))
        await self.wait()
        self.assertTrue(payload)
        self.assertEqual(payload.event_type, LayerEventType.LIVE_SESSION_STATE_CHANGED)
        if merge_to_new_layer:
            layer = Sdf.Find(self.sublayer2)
        elif is_root_layer:
            layer = self.stage.GetRootLayer()
        else:
            layer = Sdf.Find(self.sublayer)

        string = layer.ExportToString()
        self.assertEqual(string.strip(), expected_string.strip())

        if merge_to_new_layer:
            LayerUtils.remove_sublayer(self.stage.GetRootLayer(), 0)


