import carb
import omni.kit.test
import omni.kit.ui_test as ui_test
import omni.usd
import omni.ui as ui
import omni.client
import unittest
import omni.kit.app
import omni.kit.usd.layers as layers
import omni.kit.clipboard
import omni.kit.collaboration.presence_layer as pl
import omni.kit.collaboration.presence_layer.utils as pl_utils
import omni.kit.collaboration.presence_layer.extension as pl_extension

from omni.kit.viewport.utility import get_active_viewport_window
from omni.kit.collaboration.viewport.camera.manipulator.peer_user_camera import PeerUserCamera
from omni.ui.tests.test_base import OmniUiTest
from pathlib import Path
from pxr import Usd, Sdf, UsdGeom
from omni.kit.ui_test import Vec2
from omni.kit.usd.layers.tests.mock_utils import MockLiveSyncingApi, join_new_simulated_user, quit_simulated_user

CURRENT_PATH = Path(__file__).parent.joinpath("../../../../../../data")


class ToggleExpansionState():
    def __init__(self):
        self.__settings = carb.settings.get_settings()
        self.__key = "/persistent/exts/omni.kit.viewport.menubar.camera/expand"
        self.__restore_value = self.__settings.get(self.__key)
    def set(self, value: bool):
        self.__settings.set(self.__key, value)
    def __del__(self):
        self.__settings.set(self.__key, self.__restore_value)


class TestCollaborationViewportCamera(OmniUiTest):

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

        self._golden_img_dir = CURRENT_PATH.absolute().resolve().joinpath("tests")
        self.layers = layers.get_layers(self.usd_context)
        self.live_syncing = layers.get_live_syncing()
        self.local_builtin_cameras = [
            "/OmniverseKit_Persp", "/OmniverseKit_Front", "/OmniverseKit_Top", "/OmniverseKit_Right"
        ]

        self.test_session_name = "test"
        self.simulated_user_names = []
        self.simulated_user_ids   = []
        self.stage_url = "omniverse://__fake_omniverse_server__/test/camera.usd"

    async def tearDown(self):
        pass

    async def wait(self, frames=10):
        for i in range(frames):
            await self.app.next_update_async()

    async def __create_simulated_users(self, count=2):
        for i in range(count):
            user_name = 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(2)

    async def __create_fake_stage_and_join_session(self, bound_non_default_camera=False):
        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)
        camera = stage.DefinePrim("/Camera", "Camera")
        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)

        if bound_non_default_camera:
            viewport_window = get_active_viewport_window()
            viewport_api = viewport_window.viewport_api
            viewport_api.camera_path = str(camera.GetPath())
            await self.wait()

            self.assertEqual(viewport_api.camera_path, str(camera.GetPath()))

        self.live_syncing.join_live_session(session)
        self.assertTrue(self.live_syncing.get_current_live_session())
        await self.wait()

        return stage, layer

    async def __bound_camera(self, shared_stage: Usd.Stage, user_id: str, camera_path: Sdf.Path):
        if not camera_path:
            camera_path = Sdf.Path.emptyPath

        camera_path = Sdf.Path(camera_path)
        bound_camera_property_path = pl_utils.get_bound_camera_property_path(user_id)
        property_spec = pl_utils.get_or_create_property_spec(
            shared_stage.GetRootLayer(), bound_camera_property_path,
            Sdf.ValueTypeNames.String
        )

        builtin_camera_name = pl_utils.is_local_builtin_camera(camera_path)
        if not builtin_camera_name:
            property_spec.default = str(camera_path)
        else:
            property_spec.default = builtin_camera_name

        if builtin_camera_name:
            persp_camera = pl_utils.get_user_shared_root_path(user_id).AppendChild(builtin_camera_name)
            shared_stage.DefinePrim(persp_camera, "Camera")

        await self.wait()

    @MockLiveSyncingApi("test", "test")
    async def test_camera_switch(self):
        await self.wait()
        await self.__create_fake_stage_and_join_session(True)
        await self.wait()

        viewport_window = get_active_viewport_window()
        viewport_api = viewport_window.viewport_api
        self.assertEqual(viewport_api.camera_path, self.local_builtin_cameras[0])

    @MockLiveSyncingApi("test", "test")
    async def test_follow_user_menu(self):
        expand_state = ToggleExpansionState()
        expand_state.set(False)

        await self.wait()
        await self.__create_fake_stage_and_join_session()
        await self.wait()
        await self.__create_simulated_users()
        await self.wait()

        presence_layer = pl.get_presence_layer_interface(self.usd_context)
        shared_stage = presence_layer.get_shared_data_stage()
        self.assertTrue(shared_stage)

        # make "Cameras" menu and "Follow User" menu show by simulating mouse moving and clicking
        await ui_test.emulate_mouse_move(Vec2(40, 40))
        await ui_test.emulate_mouse_click()
        await self.wait()
        await ui_test.emulate_mouse_move(Vec2(40, 240))
        await ui_test.emulate_mouse_click()
        await self.wait()


        await self.__bound_camera(shared_stage, self.simulated_user_ids[1], Sdf.Path(self.local_builtin_cameras[0]))
        await self.wait()

        presence_layer.enter_follow_mode(self.simulated_user_ids[1])
        await self.wait()

        await self.finalize_test(golden_img_dir=self._golden_img_dir)

        camera_menu = ui_test.find("Viewport//Frame/**/Menu[*].text=='Perspective'")
        self.assertTrue(camera_menu)

        menu_items = camera_menu.find_all("**/")
        menu_item_lables = []
        for widget in menu_items:
            if isinstance(widget.widget, (ui.Menu, ui.MenuItem)):
                menu_item_lables.append(widget.widget.text)

        #make sure there are two menu items under follow user menu item
        follow_users = menu_item_lables[-2:]
        self.assertEqual(set(follow_users), set(self.simulated_user_names))

        # Log off one user from the live session
        quit_simulated_user(self.simulated_user_ids[1])
        await self.wait()

        # make "Cameras" menu and "Follow User" menu show by simulating mouse moving and clicking
        await ui_test.emulate_mouse_move(Vec2(0, 0))
        await ui_test.emulate_mouse_click()
        await ui_test.emulate_mouse_move(Vec2(40, 40))
        await ui_test.emulate_mouse_click()
        await self.wait()
        await ui_test.emulate_mouse_move(Vec2(40, 240))
        await ui_test.emulate_mouse_click()
        await self.wait()

        menu_items = camera_menu.find_all("**/")
        menu_item_lables = []
        for widget in menu_items:
            if isinstance(widget.widget, ui.Menu) or isinstance(widget.widget, ui.MenuItem):
                menu_item_lables.append(widget.widget.text)

        #make sure there are only one menu item under follow user menu item
        follow_users = menu_item_lables[-1:]
        self.assertEqual(follow_users, [self.simulated_user_ids[0]])

        self.live_syncing.stop_all_live_sessions()
        await self.wait()

    async def test_camera_manipulator(self):
        expand_state = ToggleExpansionState()
        expand_state.set(False)

        stage = self.usd_context.get_stage()
        prim = stage.DefinePrim("/root", "Xform")

        user = layers.LiveSessionUser("test", "test", "Kit")
        camera_widget = PeerUserCamera(user, prim.GetPath(), self.usd_context)
        try:
            await self.wait()
            await self.finalize_test(golden_img_dir=self._golden_img_dir, golden_img_name="camera_label.png")

            # Moves the prim out of the viewport frame.
            UsdGeom.XformCommonAPI(prim).SetTranslate((-10000, 0, 10000))
            camera_widget.on_world_translation_changed()
            await self.wait()
            await self.finalize_test(golden_img_dir=self._golden_img_dir, golden_img_name="offscreen_pointer.png")

            UsdGeom.XformCommonAPI(prim).SetTranslate((10000, 0, 10000))
            camera_widget.on_world_translation_changed()
            await self.wait()
            await self.finalize_test(golden_img_dir=self._golden_img_dir, golden_img_name="offscreen_pointer_2.png")
        finally:
            camera_widget.destroy()
            camera_widget = None
            await self.wait()
