# 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.
#
"""The documentation for ColorWidget"""

from .doc_page import DocPage
from omni import ui
from pxr import Gf
from pxr import UsdGeom
from typing import Any
import inspect
import omni.kit.commands
import omni.usd
import weakref

SPACING = 5


class SetDisplayColorCommand(omni.kit.commands.Command):
    """
    Change prim display color undoable **Command**. Unlike ChangePropertyCommand, it can undo property creation.

    Args:
        gprim (Gprim): Prim to change display color on.
        value: Value to change to.
        value: Value to undo to.
    """

    def __init__(self, gprim: UsdGeom.Gprim, color: Any, prev: Any):
        self._gprim = gprim
        self._color = color
        self._prev = prev

    def do(self):
        color_attr = self._gprim.CreateDisplayColorAttr()
        color_attr.Set([self._color])

    def undo(self):
        color_attr = self._gprim.GetDisplayColorAttr()
        if self._prev is None:
            color_attr.Clear()
        else:
            color_attr.Set([self._prev])


class FloatModel(ui.SimpleFloatModel):
    def __init__(self, parent):
        super().__init__()
        self._parent = weakref.ref(parent)

    def begin_edit(self):
        parent = self._parent()
        parent.begin_edit(None)

    def end_edit(self):
        parent = self._parent()
        parent.end_edit(None)


class USDColorItem(ui.AbstractItem):
    def __init__(self, model):
        super().__init__()
        self.model = model


class USDColorModel(ui.AbstractItemModel):
    def __init__(self):
        super().__init__()

        # Create root model
        self._root_model = ui.SimpleIntModel()
        self._root_model.add_value_changed_fn(lambda a: self._item_changed(None))

        # Create three models per component
        self._items = [USDColorItem(FloatModel(self)) for i in range(3)]
        for item in self._items:
            item.model.add_value_changed_fn(lambda a, item=item: self._on_value_changed(item))

        # Omniverse contexts
        self._usd_context = omni.usd.get_context()
        self._selection = self._usd_context.get_selection()
        self._events = self._usd_context.get_stage_event_stream()
        self._stage_event_sub = self._events.create_subscription_to_pop(
            self._on_stage_event, name="omni.example.ui colorwidget stage update"
        )

        # Privates
        self._subscription = None
        self._gprim = None
        self._prev_color = None
        self._edit_mode_counter = 0

    def _on_stage_event(self, event):
        """Called with subscription to pop"""
        if event.type == int(omni.usd.StageEventType.SELECTION_CHANGED):
            self._on_selection_changed()

    def _on_selection_changed(self):
        """Called when the user changes the selection"""
        selection = self._selection.get_selected_prim_paths()
        stage = self._usd_context.get_stage()
        self._subscription = None
        self._gprim = None

        # When TC runs tests, it's possible that stage is None
        if selection and stage:
            self._gprim = UsdGeom.Gprim.Get(stage, selection[0])
            if self._gprim:
                color_attr = self._gprim.GetDisplayColorAttr()
                usd_watcher = omni.usd.get_watcher()
                self._subscription = usd_watcher.subscribe_to_change_info_path(
                    color_attr.GetPath(), self._on_usd_changed
                )

        # Change the widget color
        self._on_usd_changed()

    def _on_value_changed(self, item):
        """Called when the submodel is chaged"""
        if not self._gprim:
            return

        if self._edit_mode_counter > 0:
            # Change USD only if we are in the edit mode.
            color_attr = self._gprim.CreateDisplayColorAttr()

            color = Gf.Vec3f(
                self._items[0].model.get_value_as_float(),
                self._items[1].model.get_value_as_float(),
                self._items[2].model.get_value_as_float(),
            )

            color_attr.Set([color])

        self._item_changed(item)

    def _on_usd_changed(self, path=None):
        """Called with UsdWatcher when something in USD is changed"""
        color = self._get_current_color() or Gf.Vec3f(0.0)

        for i in range(len(self._items)):
            self._items[i].model.set_value(color[i])

    def _get_current_color(self):
        """Returns color of the current object"""
        if self._gprim:
            color_attr = self._gprim.GetDisplayColorAttr()
            if color_attr:
                color_array = color_attr.Get()
                if color_array:
                    return color_array[0]

    def get_item_children(self, item):
        """Reimplemented from the base class"""
        return self._items

    def get_item_value_model(self, item, column_id):
        """Reimplemented from the base class"""
        if item is None:
            return self._root_model
        return item.model

    def begin_edit(self, item):
        """
        Reimplemented from the base class.
        Called when the user starts editing.
        """
        if self._edit_mode_counter == 0:
            self._prev_color = self._get_current_color()

        self._edit_mode_counter += 1

    def end_edit(self, item):
        """
        Reimplemented from the base class.
        Called when the user finishes editing.
        """
        self._edit_mode_counter -= 1

        if not self._gprim or self._edit_mode_counter > 0:
            return

        color = Gf.Vec3f(
            self._items[0].model.get_value_as_float(),
            self._items[1].model.get_value_as_float(),
            self._items[2].model.get_value_as_float(),
        )

        omni.kit.commands.execute("SetDisplayColor", gprim=self._gprim, color=color, prev=self._prev_color)


def color_field():
    with ui.HStack(spacing=SPACING):
        color_model = ui.ColorWidget(width=0, height=0).model
        for item in color_model.get_item_children():
            component = color_model.get_item_value_model(item)
            ui.FloatField(component)


def color_drag():
    with ui.HStack(spacing=SPACING):
        color_model = ui.ColorWidget(0.125, 0.25, 0.5, width=0, height=0).model
        for item in color_model.get_item_children():
            component = color_model.get_item_value_model(item)
            ui.FloatDrag(component)


def color_combo():
    with ui.HStack(spacing=SPACING):
        color_model = ui.ColorWidget(width=0, height=0).model
        ui.ComboBox(color_model)


def two_usd_color_boxes():
    with ui.HStack(spacing=SPACING):
        ui.ColorWidget(USDColorModel(), width=0)
        ui.Label("First ColorWidget", name="text")

    with ui.HStack(spacing=SPACING):
        ui.ColorWidget(USDColorModel(), width=0)
        ui.Label("Second ColorWidget has independed model", name="text")


class ColorWidgetDoc(DocPage):
    """Document for ColorWidget"""

    def create_doc(self, navigate_to=None):
        self._section_title("ColorWidget")

        self._text(
            "The ColorWidget widget is a button that displays the color from the item model and can open a picker "
            "window. The color dialog's function is to allow users to choose colors."
        )

        self._text("Fields connected to the color model")
        color_field()
        self._code(inspect.getsource(color_field).split("\n", 1)[-1])

        self._text("Drags connected to the color model")
        color_drag()
        self._code(inspect.getsource(color_drag).split("\n", 1)[-1])

        self._text("ComboBox connected to the color model.")
        color_combo()
        self._code(inspect.getsource(color_combo).split("\n", 1)[-1])

        # TODO omni.ui: restore functionality for Kit Next
        if False:
            self._text("Two different models that watch USD Stage. Undo/Redo example.")
            two_usd_color_boxes()

            with ui.ScrollingFrame(height=250, vertical_scrollbar_policy=ui.ScrollBarPolicy.SCROLLBAR_ALWAYS_ON):
                self._code(
                    inspect.getsource(FloatModel)
                    + "\n"
                    + inspect.getsource(USDColorItem)
                    + "\n"
                    + inspect.getsource(USDColorModel)
                    + "\n"
                    + inspect.getsource(two_usd_color_boxes)
                )

        ui.Spacer(height=20)
