# Copyright (c) 2021, 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 .event import Event
from .event import EventSubscription
from .stage_column_delegate_registry import StageColumnDelegateRegistry
from typing import Callable
from typing import List
from typing import Optional
from typing import Tuple
from .stage_icons import StageIcons as Icons

import omni.ui as ui
import weakref


class ColumnMenuDelegate(ui.AbstractItemDelegate):
    """
    Delegate is the representation layer. TreeView calls the methods
    of the delegate to create custom widgets for each item.
    """

    MENU_STYLES = {
        "Menu.CheckBox": {"background_color": 0x0, "margin": 0},
        "Menu.CheckBox::drag": {
            "image_url": Icons().get("drag"),
            "color": 0xFF505050,
            "alignment": ui.Alignment.CENTER,
        },
        "Menu.CheckBox.Image": {"image_url": Icons().get("check_off"), "color": 0xFF8A8777},
        "Menu.CheckBox.Image:checked": {"image_url": Icons().get("check_on"), "color": 0x34C7FFFF},
    }

    def __init__(self):
        super().__init__()

    def destroy(self):
        pass

    def build_header(self, column_id):
        pass

    def build_branch(self, model, item, column_id, level, expanded):
        """Create a branch widget that opens or closes subtree"""
        pass

    def build_widget(self, model, item, column_id, level, expanded):
        """Create a widget per column per item"""
        with ui.HStack(height=0, style=ColumnMenuDelegate.MENU_STYLES):
            ui.ToolButton(
                item.checked_model,
                style_type_name_override="Menu.CheckBox",
                width=18,
                height=18,
                image_width=18,
                image_height=18,
            )
            ui.Label(item.name_model.as_string)
            ui.Image(style_type_name_override="Menu.CheckBox", name="drag", width=10)


class ColumnMenuItem(ui.AbstractItem):
    """Single item of the model"""

    def __init__(self, model, text, checked):
        super().__init__()
        self.name_model = ui.SimpleStringModel(text)
        self.checked_model = ui.SimpleBoolModel(checked)
        self.checked_model.add_value_changed_fn(self._on_checked_changed)
        self.__weak_model = weakref.ref(model)

    def _on_checked_changed(self, model):
        model = self.__weak_model()
        if model:
            model._item_changed(self)
            model._on_column_checked_changed([i[0] for i in model.get_columns()])

    def __repr__(self):
        return f'"{self.name_model.as_string}"'


class ColumnMenuModel(ui.AbstractItemModel):
    """
    Represents the model for available columns.
    """

    def __init__(self, enabled: Optional[List[str]] = None, accepted: Optional[List[str]] = None):
        super().__init__()

        self._column_delegate_sub = StageColumnDelegateRegistry().subscribe_delegate_changed(
            self._on_column_delegate_changed
        )

        self._accepted = accepted

        if enabled is None:
            self._children = []
        else:
            self._children = [
                ColumnMenuItem(self, name, True) for name in enabled if not self._accepted or name in self._accepted
            ]
        self._on_column_delegate_changed(enabled is None)

        self._on_column_checked_changed = Event()

    def destroy(self):
        self._column_delegate_sub = None
        self._children = None

    def get_item_children(self, item):
        """Returns all the children when the widget asks it."""
        if item is not None:
            # Since we are doing a flat list, we return the children of root only.
            # If it's not root we return.
            return []

        return self._children

    def get_item_value_model_count(self, item):
        """The number of columns"""
        return 1

    def get_item_value_model(self, item, column_id):
        """
        Return value model.
        It's the object that tracks the specific value.
        In our case we use ui.SimpleStringModel.
        """
        return item.name_model

    def get_drag_mime_data(self, item):
        """Returns Multipurpose Internet Mail Extensions (MIME) data for be able to drop this item somewhere"""
        # As we don't do Drag and Drop to the operating system, we return the string.
        return item.name_model.as_string

    def drop_accepted(self, target_item, source, drop_location=-1):
        """Reimplemented from AbstractItemModel. Called to highlight target when drag and drop."""
        # If target_item is None, it's the drop to root. Since it's
        # list model, we support reorganizetion of root only and we
        # don't want to create tree.
        return not target_item and drop_location >= 0

    def drop(self, target_item, source, drop_location=-1):
        """Reimplemented from AbstractItemModel. Called when dropping something to the item."""
        try:
            source_id = self._children.index(source)
        except ValueError:
            try:
                current_names = [c.name_model.as_string for c in self._children]
                source_id = current_names.index(source)
                source = self._children[source_id]

            except ValueError:
                # Not in the list. This is the source from another model.
                return

        if source_id == drop_location:
            # Nothing to do
            return

        self._children.pop(source_id)

        if drop_location > len(self._children):
            # Drop it to the end
            self._children.append(source)
        else:
            if source_id < drop_location:
                # Becase when we removed source, the array became shorter
                drop_location = drop_location - 1

            self._children.insert(drop_location, source)

        self._item_changed(None)

    def get_columns(self) -> List[Tuple[str, bool]]:
        """
        Return all the columns in the format

           `[("Visibility", True), ("Type", False)]`

        The first item is the column name, and the second item is a flag that
        is True if the column is enabled.
        """
        current_names = [c.name_model.as_string for c in self._children]
        current_checked = [c.checked_model.as_bool for c in self._children]
        return list(zip(current_names, current_checked))

    def subscribe_delegate_changed(self, fn: Callable[[List[str]], None]) -> EventSubscription:
        """
        Return the object that will automatically unsubscribe when destroyed.
        """
        return EventSubscription(self._on_column_checked_changed, fn)

    def _on_column_delegate_changed(self, enabled=False):
        """Called by StageColumnDelegateRegistry"""
        current_names = [c.name_model.as_string for c in self._children]
        column_delegate_names = StageColumnDelegateRegistry().get_column_delegate_names()

        # Exclude "Name" as it's enabled always.
        if "Name" in column_delegate_names:
            column_delegate_names.remove("Name")

        # Filter columns. We only need the columns that is in self._accepted list.
        if self._accepted:
            column_delegate_names = [name for name in column_delegate_names if name in self._accepted]

        # Form the new list keeping the order of the old columns
        new_names = []
        for name in current_names:
            if name in column_delegate_names:
                new_names.append(name)

        # Put the new columns to the end of the list
        unused_names = [name for name in column_delegate_names if name not in new_names]
        new_names += unused_names

        if current_names == new_names:
            # Nothing changed
            return

        # Create item for new columns and reuse the items for old columns
        new_children = []
        for name in new_names:
            if name in current_names:
                i = current_names.index(name)
                new_children.append(self._children[i])
            else:
                new_children.append(ColumnMenuItem(self, name, enabled))

        # Replace children
        self._children = new_children

        self._item_changed(None)
