from pathlib import Path
import omni.usd
import omni.client
import omni.kit.app
import omni.kit.usd.layers as layers
from omni.ui.tests.test_base import OmniUiTest
import omni.kit.collaboration.presence_layer.utils as pl_utils

from pxr import Usd, Sdf, UsdGeom, Gf
from typing import List

from omni.kit.usd.layers.tests.mock_utils import MockLiveSyncingApi, join_new_simulated_user, quit_simulated_user

CURRENT_PATH = Path(omni.kit.app.get_app().get_extension_manager().get_extension_path_by_module(__name__) + "/data")

class TestCollaborationSelectionOutline(OmniUiTest):
    user_counter = 0

    # Before running each test
    async def setUp(self):
        self.previous_retry_values = omni.client.set_retries(0, 0, 0)
        self.app = omni.kit.app.get_app()
        self.usd_context = omni.usd.get_context()
        self.layers = layers.get_layers(self.usd_context)
        self.live_syncing = self.layers.get_live_syncing()
        await omni.usd.get_context().new_stage_async()

        self._golden_img_dir = CURRENT_PATH.absolute().resolve().joinpath("tests")
        self.simulated_user_names = []
        self.simulated_user_ids   = []
        self.test_session_name = "test"
        self.stage_url = "omniverse://__faked_omniverse_server__/test/live_session.usd"

    async def tearDown(self):
        if omni.usd.get_context().get_stage():
            await omni.usd.get_context().close_stage_async()
        self.stage = None
        omni.client.set_retries(*self.previous_retry_values)
        self.live_syncing.stop_all_live_sessions()

    async def __create_fake_stage(self):
        format = Sdf.FileFormat.FindByExtension(".usd")
        # Sdf.Layer.New will not save layer so it won't fail.
        # This can be used to test layer identifier with omniverse sheme without
        # touching real server.
        layer = Sdf.Layer.New(format, self.stage_url)
        stage = Usd.Stage.Open(layer.identifier)
        session = self.live_syncing.create_live_session(self.test_session_name, layer_identifier=layer.identifier)
        self.assertTrue(session)

        await self.usd_context.attach_stage_async(stage)
        self.live_syncing.join_live_session(session)

        return stage, layer

    def __manipulate_stage(self, stage):
        sphere = stage.DefinePrim('/sphere01', 'Sphere')
        cube = stage.DefinePrim('/cube01', 'Cube')
        UsdGeom.XformCommonAPI(sphere).SetTranslate((100, 0, 0))
        UsdGeom.XformCommonAPI(sphere).SetScale(Gf.Vec3f(40.0, 40.0, 40.0))
        UsdGeom.XformCommonAPI(cube).SetTranslate((0, 100, 0))
        UsdGeom.XformCommonAPI(cube).SetScale(Gf.Vec3f(40.0, 40.0, 40.0))

    async def __create_simulated_users(self, count=2):
        for i in range(count):
            user_name = chr(ord('A') + TestCollaborationSelectionOutline.user_counter) # f"test{i}"
            user_id = user_name
            self.simulated_user_names.append(user_name)
            self.simulated_user_ids.append(user_id)

            join_new_simulated_user(user_name, user_id)
            await self.wait_n_updates(2)
            TestCollaborationSelectionOutline.user_counter += 1

    def __get_shared_stage(self, current_session: layers.LiveSession):
        shared_data_stage_url = current_session.url + "/shared_data/users.live"
        layer = Sdf.Layer.FindOrOpen(shared_data_stage_url)
        if not layer:
            layer = Sdf.Layer.CreateNew(shared_data_stage_url)

        return Usd.Stage.Open(layer)

    async def __select_prims(self, shared_stage: Usd.Stage, user_id: str, selections: List[str]):
        selection_property_path = pl_utils.get_selection_property_path(user_id)
        property_spec = pl_utils.get_or_create_property_spec(
            shared_stage.GetRootLayer(), selection_property_path,
            Sdf.ValueTypeNames.StringArray
        )
        property_spec.default = selections
        await self.wait_n_updates()

    async def __prepare_stage(self, num_users=2):
        await self.wait_n_updates(n_frames=10)

        # create stage and add users
        stage, layer = await self.__create_fake_stage()
        await self.wait_n_updates(n_frames=10)

        self.__manipulate_stage(stage)
        await self.__create_simulated_users(num_users)
        await self.wait_n_updates(n_frames=10)

        current_session = self.live_syncing.get_current_live_session()
        shared_stage = self.__get_shared_stage(current_session)
        assert(stage is not shared_stage)
        return shared_stage

    async def __check_and_leave(self):
        await self.finalize_test(golden_img_dir=self._golden_img_dir)
        self.live_syncing.stop_all_live_sessions()
        await self.wait_n_updates(n_frames=10)

    @MockLiveSyncingApi(user_name="test", user_id="test")
    async def test_1_selection(self):
        shared_stage = await self.__prepare_stage()

        selections = ["/cube01"]
        await self.__select_prims(shared_stage, self.simulated_user_ids[0], selections)
        await self.wait_n_updates(n_frames=10)

        await self.__check_and_leave()

    @MockLiveSyncingApi(user_name="test", user_id="test")
    async def test_2_select_two_prims(self):
        shared_stage = await self.__prepare_stage()

        selections = ["/sphere01", "/cube01"]
        await self.__select_prims(shared_stage, self.simulated_user_ids[0], selections)
        await self.wait_n_updates(n_frames=10)

        await self.__check_and_leave()

    @MockLiveSyncingApi(user_name="test", user_id="test")
    async def test_3_select_identical_prim(self):
        shared_stage = await self.__prepare_stage()

        selections = ["/cube01"]
        await self.__select_prims(shared_stage, self.simulated_user_ids[0], selections)
        await self.wait_n_updates(n_frames=10)
        await self.__select_prims(shared_stage, self.simulated_user_ids[1], selections)
        await self.wait_n_updates(n_frames=10)

        await self.__check_and_leave()

    @MockLiveSyncingApi(user_name="test", user_id="test")
    async def test_4_selection_overflow(self):
        num_users = 4
        shared_stage = await self.__prepare_stage(num_users=num_users)

        selections = ["/cube01"]
        for i in range(num_users):
            await self.__select_prims(shared_stage, self.simulated_user_ids[i], selections)

        await self.__check_and_leave()

    @MockLiveSyncingApi(user_name="test", user_id="test")
    async def test_5_select_and_unselect(self):
        shared_stage = await self.__prepare_stage()

        selections = ["/cube01", "/sphere01"]
        await self.__select_prims(shared_stage, self.simulated_user_ids[0], selections)
        await self.wait_n_updates(n_frames=10)

        selections = ["/sphere01"]
        await self.__select_prims(shared_stage, self.simulated_user_ids[0], selections)
        await self.wait_n_updates(n_frames=10)

        await self.__check_and_leave()
