# 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.
#
import omni.ui as ui
from omni.kit.widget.stage import StageWidget
from pxr import Usd
from .singleton import Singleton


@Singleton
class ViewerWindows:
    """This object keeps all the Stage Viewer windows"""

    def __init__(self):
        self.__windows = {}

    def open_window(self, filepath):
        """Open StageViewer window with the stage file opened in it"""
        if filepath in self.__windows:
            self.__windows[filepath].visible = True
        else:
            stage = Usd.Stage.Open(filepath, Usd.Stage.LoadNone)
            if not stage:
                return
            window = StageViewer(stage)
            # When window is closed, remove it from the list
            window.set_visibility_changed_fn(lambda _, f=filepath: self.close(f))
            self.__windows[filepath] = window

    def close(self, filepath):
        """Close and remove specific window"""
        self.__windows[filepath].destroy()
        self.__windows[filepath] = None
        del self.__windows[filepath]

    def close_all(self):
        """Close and remove all windows"""
        self.__windows = {}


class StageViewer(ui.Window):
    """The Stage window that displays a custom stage"""

    def __init__(self, stage: Usd.Stage, **kwargs):
        if stage:
            name = stage.GetRootLayer().identifier
        else:
            name = "None"

        kwargs["flags"] = ui.WINDOW_FLAGS_NO_SCROLLBAR
        kwargs["width"] = 600
        kwargs["height"] = 800
        super().__init__(f"Stage {name}", **kwargs)

        with self.frame:
            # We only show the Type column
            self.__stage_widget = StageWidget(stage, columns_accepted=["Type"], lazy_payloads=True)

    def destroy(self):
        """
        Called by extension before destroying this object. It doesn't happen automatically.
        Without this hot reloading doesn't work.
        """
        self.__stage_widget.destroy()
        self.__stage_widget = None
        super().destroy()
