# Copyright (c) 2022, 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.
#

from __future__ import annotations

import asyncio
import colorsys
import weakref
from collections import defaultdict
from typing import DefaultDict, Dict, Set, Union
from weakref import ProxyType
import concurrent.futures

import carb
import carb.dictionary
import carb.events
import carb.profiler
import carb.settings
import omni.timeline
import omni.usd
from omni.kit.async_engine import run_coroutine
from omni.kit.manipulator.transform.style import COLOR_X, COLOR_Y, COLOR_Z, abgr_to_color
from omni.kit.manipulator.transform.settings_constants import Constants as TrCon
from omni.ui import color as cl
from omni.ui import scene as sc
from pxr import Sdf, Tf, Usd, UsdGeom

from .model import PrimTransformModel
from omni.kit.viewport.manipulator.transform import Viewport1WindowState
from .settings_constants import Constants
from .utils import flatten, get_local_transform_pivot_inv

LARGE_SELECTION_CAP = 20


class PreventViewportOthers(sc.GestureManager):
    def can_be_prevented(self, gesture):
        return True

    def should_prevent(self, gesture, preventer):
        if preventer.state == sc.GestureState.BEGAN or preventer.state == sc.GestureState.CHANGED:
            if issubclass(type(preventer), ClickMarkerGesture):
                if issubclass(type(gesture), ClickMarkerGesture):
                    return gesture.gesture_payload.ray_distance > preventer.gesture_payload.ray_distance
                elif isinstance(gesture, ClickMarkerGesture):
                    return False
                else:
                    return True
            else:
                return True

        return super().should_prevent(gesture, preventer)


class ClickMarkerGesture(sc.DragGesture):
    def __init__(
        self,
        prim_path: Sdf.Path,
        marker: ProxyType[ReferencePrimMarker],
    ):
        super().__init__(manager=PreventViewportOthers())
        self._prim_path = prim_path
        self._marker = marker
        self._vp1_window_state = None

    def on_began(self):
        self._viewport_on_began()

    def on_canceled(self):
        self._viewport_on_ended()

    def on_ended(self):
        self._viewport_on_ended()
        self._marker.on_pivot_marker_picked(self._prim_path)

    def _viewport_on_began(self):
        self._viewport_on_ended()
        if self._marker.legacy:
            self._vp1_window_state = Viewport1WindowState()

    def _viewport_on_ended(self):
        if self._vp1_window_state:
            self._vp1_window_state.destroy()
            self._vp1_window_state = None


class MarkerHoverGesture(sc.HoverGesture):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.items = []
        self._began_count = 0
        self._original_colors = []
        self._original_ends = []

    def on_began(self):
        self._began_count += 1
        if self._began_count == 1:
            self._original_colors = [None] * len(self.items)
            self._original_ends = [None] * len(self.items)
            for i, item in enumerate(self.items):
                self._original_colors[i] = item.color
                item.color = self._make_brighter(item.color)

                if isinstance(item, sc.Line):
                    self._original_ends[i] = item.end
                    end = [dim * 1.2 for dim in item.end]
                    item.end = end

                if isinstance(item, sc.Arc):
                    ...  # TODO change color to white, blocked by OM-56044

    def on_ended(self):
        self._began_count -= 1
        if self._began_count <= 0:
            self._began_count = 0
            for i, item in enumerate(self.items):
                item.color = self._original_colors[i]

                if isinstance(item, sc.Line):
                    item.end = self._original_ends[i]

                if isinstance(item, sc.Arc):
                    ...  # TODO change color back, blocked by OM-56044

    def _make_brighter(self, color):
        hsv = colorsys.rgb_to_hsv(color[0], color[1], color[2])
        rgb = colorsys.hsv_to_rgb(hsv[0], hsv[1], hsv[2] * 1.2)
        return cl(rgb[0], rgb[1], rgb[2], color[3])


class ReferencePrimMarker(sc.Manipulator):
    def __init__(
        self, usd_context_name: str = "", manipulator_model: ProxyType[PrimTransformModel] = None, legacy: bool = False
    ):
        super().__init__()
        self._settings = carb.settings.get_settings()
        self._dict = carb.dictionary.get_dictionary()
        self._usd_context_name = usd_context_name
        self._usd_context = omni.usd.get_context(self._usd_context_name)
        self._manipulator_model = manipulator_model
        self._legacy = legacy
        self._selection = self._usd_context.get_selection()
        self._timeline = omni.timeline.get_timeline_interface()
        self._current_time = self._timeline.get_current_time()

        # dict from prefixes -> dict of affected markers.
        self._markers: DefaultDict[Dict] = defaultdict(dict)
        self._stage_listener = None
        self._pending_changed_paths: Set[Sdf.Path] = set()
        self._process_pending_change_task_or_future: Union[asyncio.Task, concurrent.futures.Future, None] = None

        # order=1 to ensure the event handler is called after prim manipulator
        self._stage_event_sub = self._usd_context.get_stage_event_stream().create_subscription_to_pop(
            self._on_stage_event, name="ReferencePrimMarker stage event", order=1
        )

        self._placement_sub = self._settings.subscribe_to_node_change_events(
            Constants.MANIPULATOR_PLACEMENT_SETTING, self._on_placement_changed
        )
        self._placement = self._settings.get(Constants.MANIPULATOR_PLACEMENT_SETTING)

        self._op_sub = self._settings.subscribe_to_node_change_events(TrCon.TRANSFORM_OP_SETTING, self._on_op_changed)
        self._selected_op = self._settings.get(TrCon.TRANSFORM_OP_SETTING)

    def destroy(self):
        self._stage_event_sub = None
        if self._placement_sub:
            self._settings.unsubscribe_to_change_events(self._placement_sub)
            self._placement_sub = None

        if self._op_sub:
            self._settings.unsubscribe_to_change_events(self._op_sub)
            self._op_sub = None

        if self._stage_listener:
            self._stage_listener.Revoke()
            self._stage_listener = None

        if self._process_pending_change_task_or_future and not self._process_pending_change_task_or_future.done():
            self._process_pending_change_task_or_future.cancel()
            self._process_pending_change_task_or_future = None

        self._timeline_sub = None

    @property
    def usd_context_name(self) -> str:
        return self._usd_context_name

    @usd_context_name.setter
    def usd_context_name(self, value: str):
        if value != self._usd_context_name:
            new_usd_context = omni.usd.get_context(value)

            if not new_usd_context:
                carb.log_error(f"Invalid usd context name {value}")
                return

            if self._stage_listener:
                self._stage_listener.Revoke()
                self._stage_listener = None
            self._usd_context_name = value
            self._usd_context = new_usd_context

            # order=1 to ensure the event handler is called after prim manipulator
            self._stage_event_sub = self._usd_context.get_stage_event_stream().create_subscription_to_pop(
                self._on_stage_event, name="ReferencePrimMarker stage event", order=1
            )
            self.invalidate()

    @property
    def manipulator_model(self) -> ProxyType[PrimTransformModel]:
        return self._manipulator_model

    @manipulator_model.setter
    def manipulator_model(self, value):
        if value != self._manipulator_model:
            self._manipulator_model = value
            self.invalidate()

    @property
    def legacy(self) -> bool:
        return self._legacy

    @legacy.setter
    def legacy(self, value):
        self._legacy = value

    def on_pivot_marker_picked(self, path: Sdf.Path):
        if self._manipulator_model.set_pivot_prim_path(path):
            # Hide marker on the new pivot prim and show marker on old pivot prim
            old_pivot_marker = self._markers.get(self._pivot_prim_path, {}).get(self._pivot_prim_path, None)
            if old_pivot_marker:
                old_pivot_marker.visible = True

            self._pivot_prim_path = self._manipulator_model.get_pivot_prim_path()

            new_pivot_marker = self._markers.get(self._pivot_prim_path, {}).get(self._pivot_prim_path, None)
            if new_pivot_marker:
                new_pivot_marker.visible = False

    def on_build(self):
        if self._stage_listener:
            self._stage_listener.Revoke()
            self._stage_listener = None

        self._timeline_sub = None
        self._markers.clear()

        if self._placement != Constants.MANIPULATOR_PLACEMENT_PICK_REF_PRIM:
            return

        # don't build marker on "select" mode
        if self._selected_op == TrCon.TRANSFORM_OP_SELECT:
            return

        selected_xformable_paths = self._manipulator_model.xformable_prim_paths
        stage = self._usd_context.get_stage()
        self._current_time = self._timeline.get_current_time()

        selection_count = len(selected_xformable_paths)

        # skip if there's only one xformable
        if selection_count <= 1:
            return

        if selection_count > LARGE_SELECTION_CAP:
            carb.log_warn(
                f"{selection_count} is greater than the maximum selection cap {LARGE_SELECTION_CAP}, "
                f"{Constants.MANIPULATOR_PLACEMENT_PICK_REF_PRIM} mode will be disabled and fallback to "
                f"{Constants.MANIPULATOR_PLACEMENT_LAST_PRIM_PIVOT} due to performance concern."
            )
            return

        timecode = self._get_current_time_code()
        self._pivot_prim_path = self._manipulator_model.get_pivot_prim_path()
        xform_cache = UsdGeom.XformCache(timecode)

        for path in selected_xformable_paths:
            prim = stage.GetPrimAtPath(path)
            pivot_inv = get_local_transform_pivot_inv(prim, timecode)
            transform = pivot_inv.GetInverse() * xform_cache.GetLocalToWorldTransform(prim)
            transform.Orthonormalize()  # in case of a none uniform scale

            marker_transform = sc.Transform(transform=flatten(transform), visible=self._pivot_prim_path != path)
            with marker_transform:
                with sc.Transform(scale_to=sc.Space.SCREEN):
                    gesture = ClickMarkerGesture(path, marker=weakref.proxy(self))
                    hover_gesture = MarkerHoverGesture()
                    x_line = sc.Line(
                        (0, 0, 0),
                        (50, 0, 0),
                        color=abgr_to_color(COLOR_X),
                        thickness=2,
                        intersection_thickness=8,
                        gestures=[gesture, hover_gesture],
                    )
                    y_line = sc.Line(
                        (0, 0, 0),
                        (0, 50, 0),
                        color=abgr_to_color(COLOR_Y),
                        thickness=2,
                        intersection_thickness=8,
                        gestures=[gesture, hover_gesture],
                    )
                    z_line = sc.Line(
                        (0, 0, 0),
                        (0, 0, 50),
                        color=abgr_to_color(COLOR_Z),
                        thickness=2,
                        intersection_thickness=8,
                        gestures=[gesture, hover_gesture],
                    )
                    with sc.Transform(look_at=sc.Transform.LookAt.CAMERA):
                        point = sc.Arc(
                            radius=6,
                            wireframe=False,
                            tesselation=8,
                            color=cl.white,  # TODO default color grey, blocked by OM-56044
                        )
                    hover_gesture.items = [x_line, y_line, z_line, point]

            prefixes = path.GetPrefixes()
            for prefix in prefixes:
                self._markers[prefix][path] = marker_transform

        if self._markers:
            self._stage_listener = Tf.Notice.Register(Usd.Notice.ObjectsChanged, self._on_objects_changed, stage)
            self._timeline_sub = self._timeline.get_timeline_event_stream().create_subscription_to_pop(
                self._on_timeline_event
            )

    def _on_stage_event(self, event: carb.events.IEvent):
        if event.type == int(omni.usd.StageEventType.SELECTION_CHANGED):
            self._on_selection_changed()

    def _on_selection_changed(self):
        if self._placement == Constants.MANIPULATOR_PLACEMENT_PICK_REF_PRIM:
            self.invalidate()

    def _on_placement_changed(self, item, event_type):
        placement = self._dict.get(item)
        if placement != self._placement:
            self._placement = placement
            self.invalidate()

    def _on_op_changed(self, item, event_type):
        selected_op = self._dict.get(item)
        if selected_op != self._selected_op:
            if selected_op == TrCon.TRANSFORM_OP_SELECT or self._selected_op == TrCon.TRANSFORM_OP_SELECT:
                self.invalidate()
            self._selected_op = selected_op

    def _on_timeline_event(self, e: carb.events.IEvent):
        if e.type == int(omni.timeline.TimelineEventType.CURRENT_TIME_TICKED):
            current_time = e.payload["currentTime"]
            if current_time != self._current_time:
                self._current_time = current_time
                if self._placement != Constants.MANIPULATOR_PLACEMENT_PICK_REF_PRIM:
                    # TODO only invalidate transform if this prim or ancestors transforms are time varying?
                    self.invalidate()

    @carb.profiler.profile
    async def _process_pending_change(self):
        processed_transforms = set()
        timecode = self._get_current_time_code()
        xform_cache = UsdGeom.XformCache(timecode)
        stage = self._usd_context.get_stage()

        for path in self._pending_changed_paths:
            prim_path = path.GetPrimPath()
            affected_transforms = self._markers.get(prim_path, {})
            if affected_transforms:
                if not UsdGeom.Xformable.IsTransformationAffectedByAttrNamed(path.name):
                    continue

                for path, transform in affected_transforms.items():
                    if transform in processed_transforms:
                        continue

                    prim = stage.GetPrimAtPath(path)
                    if not prim:
                        continue

                    pivot_inv = get_local_transform_pivot_inv(prim, timecode)
                    xform = pivot_inv.GetInverse() * xform_cache.GetLocalToWorldTransform(prim)
                    xform.Orthonormalize()  # in case of a none uniform scale

                    transform_value = flatten(xform)
                    transform.transform = transform_value
                    processed_transforms.add(transform)

        self._pending_changed_paths.clear()

    @carb.profiler.profile
    def _on_objects_changed(self, notice, sender):
        self._pending_changed_paths.update(notice.GetChangedInfoOnlyPaths())

        if self._process_pending_change_task_or_future is None or self._process_pending_change_task_or_future.done():
            self._process_pending_change_task_or_future = run_coroutine(self._process_pending_change())

    def _get_current_time_code(self):
        return Usd.TimeCode(
            omni.usd.get_frame_time_code(self._current_time, self._usd_context.get_stage().GetTimeCodesPerSecond())
        )
