# 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.
#
import asyncio
import weakref

import omni.kit.app
import omni.ui as ui

from .checkpoints_model import CheckpointsModel
from .checkpoints_delegate import CheckpointsDelegate
from .style import get_style


class CheckpointWidget:
    def __init__(self, url: str = "", width=700, show_none_entry=False, **kwargs):
        """
        Build a CheckpointWidget. Called while constructing UI layout.

        Args:
            url (str): The url of the file to list checkpoints from.
            width (str): Width of the checkpoint widget.
            show_none_entry (bool): Whether to show a dummy "head" entry to represent not using checkpoint.
        """

        self._checkpoint_model = CheckpointsModel(show_none_entry)
        self._checkpoint_model.set_on_list_checkpoint_fn(self._on_list_checkpoint)
        self._checkpoint_delegate = CheckpointsDelegate(kwargs.get("open_file", None))
        self._tree_view = None
        self._labels = {}
        self._on_selection_changed_fn = []
        self._on_list_checkpoint_fn = kwargs.get("on_list_checkpoint_fn", None)
        self._build_layout(width)
        self.set_url(url)
        self.set_multi_select(False)
        self._on_list_checkpoint(supports_checkpointing=True,
                                 has_checkpoints=True,
                                 current_checkpoint=None,
                                 multi_select=False)

    def __del__(self):
        self.destroy()

    def destroy(self):
        self._tree_view = None
        self._labels = None
        self._checkpoint_stack = None
        self._root_frame = None
        if self._checkpoint_model:
            self._checkpoint_model.destroy()
        self._checkpoint_model = None
        self._checkpoint_delegate = None
        self._on_selection_changed_fn.clear()

    def set_multi_select(self, state:bool):
        self._checkpoint_model.set_multi_select(state)

    def set_url(self, url: str):
        self._checkpoint_model.set_url(url)

    def set_search(self, keywords):
        self._checkpoint_model.set_search(keywords)
        
    def add_on_selection_changed_fn(self, fn):
        self._on_selection_changed_fn.append(fn)

    def set_on_list_checkpoint_fn(self, fn):
        self._on_list_checkpoint_fn = fn

    def empty(self):
        return self._checkpoint_model.empty()

    def _build_layout(self, width):
        self._root_frame = ui.Frame(visible=True, width=width)
        with self._root_frame:
            with ui.ZStack():
                self._labels["cp_not_suppored"] = ui.Label("Checkpoints Not Supported.", style={"font_size": 18})
                self._labels["multi_select"] = ui.Label("Checkpoints Cannot Be Displayed For Multiple Items.", style={"font_size": 18})

                with ui.VStack():
                    self._checkpoint_stack = ui.VStack(visible=False, style=get_style())
                    with self._checkpoint_stack:
                        with ui.ScrollingFrame(
                            height=ui.Fraction(1),
                            horizontal_scrollbar_policy=ui.ScrollBarPolicy.SCROLLBAR_ALWAYS_OFF,
                            vertical_scrollbar_policy=ui.ScrollBarPolicy.SCROLLBAR_AS_NEEDED,
                            style_type_name_override="TreeView",
                        ):
                            self._tree_view = ui.TreeView(
                                self._checkpoint_model,
                                root_visible=False,
                                header_visible=True,
                                delegate=self._checkpoint_delegate,
                                column_widths=[80, ui.Fraction(1), 20, 120, 100, 150],
                            )

                            # only allow selecting one checkpoint
                            def on_selection_changed(weak_self, selection):
                                weak_self = weak_self()
                                if not weak_self:
                                    return

                                if len(selection) > 1:
                                    selection = [selection[-1]]
                                    if weak_self._tree_view.selection != selection:
                                        weak_self._tree_view.selection = selection

                                for fn in weak_self._on_selection_changed_fn:
                                    fn(selection[0] if selection else None)

                            self._tree_view.set_selection_changed_fn(
                                lambda selection, weak_self=weakref.ref(self): on_selection_changed(weak_self, selection)
                            )

                    self._labels["no_checkpoint"] = ui.Label("No Checkpoint")

    def _on_list_checkpoint(self, supports_checkpointing, has_checkpoints, current_checkpoint, multi_select):
        if multi_select:
            self._root_frame.visible = True
            self._checkpoint_stack.visible = False
            self._labels["no_checkpoint"].visible = False
            self._labels["cp_not_suppored"].visible = False
            self._labels["multi_select"].visible = True
        elif not supports_checkpointing:
            self._checkpoint_stack.visible = False
            self._root_frame.visible = True
            self._labels["no_checkpoint"].visible = False
            self._labels["cp_not_suppored"].visible = True
            self._labels["multi_select"].visible = False
        else:
            self._checkpoint_stack.visible = has_checkpoints
            self._root_frame.visible = has_checkpoints
            self._labels["cp_not_suppored"].visible = False
            self._labels["multi_select"].visible = False
            self._labels["no_checkpoint"].visible = not has_checkpoints

        # Must delay the selection for one frame so that treeview is synced with model for scroll to selection to work.
        async def delayed_selection(weak_tree_view, callback):
            tree_view = weak_tree_view()
            if not tree_view:
                return

            await omni.kit.app.get_app().next_update_async()

            if current_checkpoint:
                tree_view.selection = [current_checkpoint]
            else:
                tree_view.selection = []

            if callback:
                callback(tree_view.selection)

        asyncio.ensure_future(delayed_selection(weakref.ref(self._tree_view), self._on_list_checkpoint_fn))
