# Copyright (c) 2018-2020, NVIDIA CORPORATION.  All rights reserved.
#
# NVIDIA CORPORATION and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION is strictly prohibited.
#
__all__ = ["PeerUserListManipulatorModel"]

import carb
import omni.usd
import omni.kit.usd.layers as layers

from omni.ui import scene as sc
from pxr import Gf, UsdGeom, Usd, Sdf, Tf, Trace


class PeerUserListManipulatorModel(sc.AbstractManipulatorModel):
    """The model tracks the position of peer user's camera in a live session."""

    class PositionItem(sc.AbstractManipulatorItem):
        def __init__(self):
            super().__init__()
            self.value = [0, 0, 0]
    
    class UserListItem(sc.AbstractManipulatorItem):
        def __init__(self):
            super().__init__()
            self.value = {}

    def __init__(self, prim_path, usd_context: omni.usd.UsdContext):
        super().__init__()

        self.position = PeerUserListManipulatorModel.PositionItem()
        self.user_list = PeerUserListManipulatorModel.UserListItem()

        self.__prim_path = Sdf.Path(prim_path)

        self.__usd_context = usd_context

        self.__update_position = True
    
    def destroy(self):
        self.position = None
        self.__usd_context = None

    def add_user(self, user_info: layers.LiveSessionUser):
        if user_info.user_id in self.user_list.value:
            return
        
        self.user_list.value[user_info.user_id] = user_info
        
        # This makes the manipulator updated
        self._item_changed(self.user_list)

    def remove_user(self, user_id):
        item = self.user_list.value.pop(user_id, None)
        if not item:
            # This makes the manipulator updated
            self._item_changed(self.user_list)
    
    def empty_user_list(self):
        return len(self.user_list.value) == 0

    def get_item(self, identifier):
        if identifier == "position":
            return self.position
        elif identifier == "user_list":
            return self.user_list
        
        return None

    def get_as_floats(self, item):
        if item == self.position:
            # Requesting position
            if self.__update_position:
                self.position.value = self._get_position()
                self.__update_position = False

            return self.position.value

        return []

    def set_floats(self, item, value):
        if item != self.position:
            return

        if not value or item.value == value:
            return

        # Set directly to the item
        item.value = value

        self.__update_position = False

        # This makes the manipulator updated
        self._item_changed(item)
    
    def on_transform_changed(self, position=None):
        """When position is provided, it will not re-calculate position but use provided value instead."""

        # Position is changed
        if position:
            self.set_floats(self.position, position)
        else:
            self.__update_position = True
            self._item_changed(self.position)

    def _get_position(self):
        """Returns center position at the top of the bbox for the current selected object."""
        stage = self.__usd_context.get_stage()
        prim = stage.GetPrimAtPath(self.__prim_path)
        if not prim:
            return []

        # Get bounding box from prim
        aabb_min, aabb_max = self.__usd_context.compute_path_world_bounding_box(
            str(self.__prim_path)
        )

        gf_range = Gf.Range3d(
            Gf.Vec3d(aabb_min[0], aabb_min[1], aabb_min[2]),
            Gf.Vec3d(aabb_max[0], aabb_max[1], aabb_max[2])
        )
        if gf_range.IsEmpty():
            # May be empty (UsdGeom.Xform for example), so build a scene-scaled constant box
            world_units = UsdGeom.GetStageMetersPerUnit(stage)
            if Gf.IsClose(world_units, 0.0, 1e-6):
                world_units = 0.01

            xformable_prim = UsdGeom.Xformable(prim)
            world_xform = xformable_prim.ComputeLocalToWorldTransform(Usd.TimeCode.Default())

            ref_position = world_xform.ExtractTranslation()
            extent = Gf.Vec3d(0.1 / world_units) # 10cm by default
            gf_range.SetMin(ref_position - extent)
            gf_range.SetMax(ref_position + extent)
        
        bboxMin = gf_range.GetMin()
        bboxMax = gf_range.GetMax()
        position = [(bboxMin[0] + bboxMax[0]) * 0.5, bboxMax[1], (bboxMin[2] + bboxMax[2]) * 0.5]

        return position
