import carb
import sys
import os
import omni.ext
import omni.ui
import omni.kit.ui
import omni.kit.app
import omni.stats
import carb.settings

WINDOW_NAME = "Statistics"


class Extension(omni.ext.IExt):
    """Statistics view extension"""

    class ComboBoxItem(omni.ui.AbstractItem):
        def __init__(self, text):
            super().__init__()
            self.model = omni.ui.SimpleStringModel(text)

    class ComboBoxModel(omni.ui.AbstractItemModel):
        def __init__(self):
            super().__init__()
            self._current_index = omni.ui.SimpleIntModel()
            self._current_index.add_value_changed_fn(self._changed_model)
            self._items = []

        def _changed_model(self, model):
            # update stats at the end of the frame instead
            self._item_changed(None)

        def get_item_children(self, item):
            return self._items

        def get_item_value_model(self, item, column_id):
            if item is None:
                return self._current_index
            return item.model

        def _update_scopes(self, stat_iface):
            scopes = stat_iface.get_scopes()
            scope_count = len(scopes)
            item_count = len(self._items)
            if item_count != scope_count:
                selected_scope = ""
                if item_count != 0:
                    selected_scope = self._current_index
                    self._items.clear()
                for i in range(scope_count):
                    scope = scopes[i]
                    self._items.append(Extension.ComboBoxItem(scope["name"]))
                    # verify if a scope is already selected
                    if selected_scope == scope["name"]:
                        self._current_index.as_int = i
                # A default selection if it is not already picked
                if not selected_scope:
                    self._current_index.as_int = 0
            # update stats
            if scope_count != 0:
                selected_scope_node = scopes[self._current_index.as_int]
                self._item_changed(None)
                return selected_scope_node
            return None

    def __init__(self):
        self._window = None
        self._app = None
        self._stats_mode = None
        self._scope_description = None
        self._stats = None
        self._stat_grid = None
        self._scope_combo_model = None
        self._stats_names = None
        self._stats_values = None
        self._stats_desc = None
        self._show_names = False  # Don't show names by default, just descriptions
        super().__init__()

    def get_name(self):
        return WINDOW_NAME

    def _menu_callback(self, menu, value):
        self._window.visible = value
        # update menu state
        if self._menu:
            omni.kit.ui.get_editor_menu().set_value(f"Window/{WINDOW_NAME}", value)

    def on_startup(self):
        self._app = omni.kit.app.get_app()
        self._stats = omni.stats.get_stats_interface()
        menu_path = f"Window/{WINDOW_NAME}"
        self._window = omni.ui.Window(
            WINDOW_NAME,
            width=400,
            height=600,
            padding_x=10,
            visible=False,
            dockPreference=omni.ui.DockPreference.RIGHT_TOP,
        )
        self._window.deferred_dock_in("Details", omni.ui.DockPolicy.CURRENT_WINDOW_IS_ACTIVE)

        with self._window.frame:
            with omni.ui.VStack(height=0, spacing=8):
                with omni.ui.HStack():
                    omni.ui.Label("Scope to display:")
                    self._scope_combo_model = Extension.ComboBoxModel()
                    self._scope_combo = omni.ui.ComboBox(self._scope_combo_model)

                with omni.ui.HStack():
                    omni.ui.Label("Scope description:")
                    self._scope_description = omni.ui.Label("", word_wrap=True)

                omni.ui.Line()
                # green color for header
                style_header = {"color": 0xFF00B976}
                with omni.ui.HStack(style=style_header):
                    if self._show_names:
                        omni.ui.Label("Name", alignment=omni.ui.Alignment.LEFT)
                    omni.ui.Label("Description", alignment=omni.ui.Alignment.LEFT)
                    omni.ui.Label("Amount", alignment=omni.ui.Alignment.RIGHT)
                omni.ui.Line()

                # For performance, draw with two labels (per-frame update)
                with omni.ui.HStack(style=style_header):
                    if self._show_names:
                        self._stats_names = omni.ui.Label("", alignment=omni.ui.Alignment.LEFT)
                    self._stats_desc = omni.ui.Label("", alignment=omni.ui.Alignment.LEFT)
                    self._stats_values = omni.ui.Label("", alignment=omni.ui.Alignment.RIGHT)

        try:
            editor_menu = omni.kit.ui.get_editor_menu()
            self._menu = editor_menu.add_item(menu_path, self._menu_callback, toggle=True)
            editor_menu.set_priority(menu_path, 15)
        except:
            pass

        self._sub_event = self._app.get_update_event_stream().create_subscription_to_pop(
            self._on_update, name="omni.stats frame statistics"
        )
        self._window.set_visibility_changed_fn(self._visiblity_changed_fn)

    def on_shutdown(self):
        self._menu = None
        self._window = None
        self._sub_event = None

    def _visiblity_changed_fn(self, visible):
        # update menu state
        if self._menu:
            omni.kit.ui.get_editor_menu().set_value(f"Window/{WINDOW_NAME}", visible)
            
        carb.settings.get_settings().set("/profiler/enableDeviceUtilizationQuery", visible)

    def _update_stats(self, scope_node):
        stat_nodes = self._stats.get_stats(scope_node["scopeId"])
        stats_names = ""
        stats_values = ""
        stats_descs = ""
        # Sort nodes in descending order based on the alphabet
        stat_nodes = sorted(stat_nodes, key=lambda node: node["description"].lower(), reverse=False)
        for node in stat_nodes:
            if self._show_names:
                stats_names += node["name"] + "\n"
            stats_descs += node["description"] + "\n"
            if node["type"] == 0:
                stats_values += "{:,}".format(node["value"]) + "\n"
            elif node["type"] == 1:
                stats_values += "{:.3f}".format(node["value"]) + "\n"
        if self._show_names:
            self._stats_names.text = stats_names
        self._stats_values.text = stats_values
        self._stats_desc.text = stats_descs

    def _on_update(self, evt):
        if not self._window.visible:
            return
        scope_node = self._scope_combo_model._update_scopes(self._stats)
        if scope_node is not None:
            self._scope_description.text = scope_node["description"]
            self._update_stats(scope_node)
