from collections import defaultdict
from inspect import isclass
from functools import partial
import omni.kit.commands
import omni.kit.undo
from omni import ui
from .history import HistoryModel, HistoryDelegate, MainItem
from .search import SearchModel, SearchDelegate, CommandItem, Columns


class Window:
    def __init__(self, title, menu_path):
        self._menu_path = menu_path
        self._win_history = ui.Window(title, width=700, height=500)
        self._win_history.set_visibility_changed_fn(self._on_visibility_changed)
        self._win_search = None
        self._search_frame = None
        self._doc_frame = None
        self._history_delegate = HistoryDelegate()
        self._search_delegate = SearchDelegate()

        self._build_ui()

        omni.kit.commands.subscribe_on_change(self._update_ui)
        self._update_ui()

    def on_shutdown(self):
        self._win_history = None
        self._win_search = None
        self._search_frame = None
        self._doc_frame = None
        self._history_delegate = None
        self._search_delegate = None

        omni.kit.commands.unsubscribe_on_change(self._update_ui)

    def show(self):
        self._win_history.visible = True
        self._win_history.focus()

    def hide(self):
        self._win_history.visible = False
        if self._win_search is not None:
            self._win_search.visible = False

    def _on_visibility_changed(self, visible):
        if not visible:
            omni.kit.ui.get_editor_menu().set_value(self._menu_path, False)
            if self._win_search is not None:
                self._win_search.visible = False

    def _build_ui(self):
        with self._win_history.frame:
            with ui.VStack():
                with ui.HStack(height=20):
                    ui.Button("Clear history", width=60, clicked_fn=self._clear_history)
                    ui.Button("Search commands", width=60, clicked_fn=self._show_registered)
                with ui.ScrollingFrame(
                    horizontal_scrollbar_policy=ui.ScrollBarPolicy.SCROLLBAR_AS_NEEDED,
                    vertical_scrollbar_policy=ui.ScrollBarPolicy.SCROLLBAR_AS_NEEDED,
                    style_type_name_override="TreeView",
                ):
                    self._history_model = HistoryModel()
                    self._history_tree_view = ui.TreeView(
                        self._history_model,
                        delegate=self._history_delegate,
                        root_visible=False,
                        style={"TreeView.Item::error": {"color": 0xFF0000BB}},
                        column_widths=[ui.Fraction(1)],
                    )
                    self._history_tree_view.set_mouse_double_clicked_fn(self._on_double_click)
                with ui.HStack(height=20):
                    ui.Label("Generate script to clipboard from:", width=0)
                    ui.Button("Top-level commands", width=60, clicked_fn=partial(self._copy_to_clipboard, False))
                    ui.Button("Selected commands", width=60, clicked_fn=partial(self._copy_to_clipboard, True))

    def _on_double_click(self, x, y, b, m):
        if len(self._history_tree_view.selection) <= 0:
            return

        item = self._history_tree_view.selection[0]
        if isinstance(item, MainItem):
            self._history_tree_view.set_expanded(item, not self._history_tree_view.is_expanded(item), False)

    def _copy_to_clipboard(self, only_selected):
        omni.kit.clipboard.copy(self._generate_command_script(only_selected))

    def _clear_history(self):
        omni.kit.undo.clear_history()
        self._history_model._commands_changed()

    def _show_registered(self):
        if self._win_search is None:
            self._win_search = ui.Window(
                "Search Commands", width=700, height=500, dockPreference=ui.DockPreference.MAIN
            )
            with self._win_search.frame:
                with ui.VStack():
                    with ui.HStack(height=20):
                        ui.Label("Find: ", width=0)
                        self._search_field = ui.StringField()
                        self._search_field.model.add_value_changed_fn(lambda _: self._refresh_list())
                    with ui.HStack(height=300):
                        self._search_frame = ui.ScrollingFrame(
                            horizontal_scrollbar_policy=ui.ScrollBarPolicy.SCROLLBAR_AS_NEEDED,
                            vertical_scrollbar_policy=ui.ScrollBarPolicy.SCROLLBAR_AS_NEEDED,
                            style_type_name_override="TreeView",
                        )
                        with self._search_frame:
                            self._search_tree_model = SearchModel()
                            self._search_tree_view = ui.TreeView(
                                self._search_tree_model,
                                delegate=self._search_delegate,
                                selection_changed_fn=self._on_selection_changed,
                                root_visible=False,
                                header_visible=True,
                                columns_resizable=True,
                                column_widths=[x.width for x in Columns.ORDER],
                            )
                    self._doc_frame = ui.ScrollingFrame(
                        horizontal_scrollbar_policy=ui.ScrollBarPolicy.SCROLLBAR_AS_NEEDED,
                        vertical_scrollbar_policy=ui.ScrollBarPolicy.SCROLLBAR_AS_NEEDED,
                    )
            self._refresh_list()

        self._win_search.visible = True
        self._win_search.focus()

    def _on_selection_changed(self, selected_items):
        with self._doc_frame:
            with ui.VStack(height=40):
                for item in selected_items:
                    if isinstance(item, CommandItem):
                        ui.Label(f"{item.command_model.as_string}{item.class_sig}", height=30)
                        ui.Label(f"Documentation:", height=30)
                        if item.doc is None or item.doc == omni.kit.commands.Command.__doc__:
                            ui.Label("\n\t<None>")
                        else:
                            ui.Label(item.doc.strip())

    def _update_ui(self):
        self._history_model._commands_changed()
        self._refresh_list()  # TODO: check if this is called too much

    def _refresh_list(self):
        if self._search_frame is None:
            return

        search_str = self._search_field.model.get_value_as_string().lower().strip()
        self._search_tree_model._clear_commands()

        commands = omni.kit.commands.get_commands()
        for cmd_list in commands.values():
            for command in cmd_list.values():
                if search_str == "" or command.__name__.lower().find(search_str) != -1:
                    if self._search_tree_model:
                        class_sig = omni.kit.commands.get_command_class_signature(command.__name__)
                        self._search_tree_model._add_command(CommandItem(command, class_sig))

        self._search_tree_view.clear_selection()
        self._search_tree_model._commands_changed()

    def _generate_command_script(self, only_selected):
        history = omni.kit.undo.get_history().values()

        imports = "import omni.kit.commands\n"
        code_str = ""

        arg_imports = defaultdict(set)

        def parse_module_name(name):
            if name == "builtins":
                return

            nonlocal imports
            lst = name.split(".")
            if len(lst) == 1:
                imports += "import " + name + "\n"
            else:
                arg_imports[lst[0]].add(".".join(lst[1:]))

        def arg_import(val):
            if isclass(val):
                parse_module_name(val.__module__)
            else:
                parse_module_name(val.__class__.__module__)
            return val

        def gen_cmd_str(cmd):
            if len(cmd.kwargs.items()) > 0:
                args = ",\n\t".join(["{}={!r}".format(k, arg_import(v)) for k, v in cmd.kwargs.items()])
                return f"\nomni.kit.commands.execute('{cmd.name}',\n\t{args})\n"
            else:
                return f"\nomni.kit.commands.execute('{cmd.name}')\n"

        # generate a script executing top-level commands or all selected commands
        if only_selected:
            for item in self._history_tree_view.selection:
                if isinstance(item, MainItem):
                    code_str += gen_cmd_str(item._data)
        else:
            for cmd in history:
                if cmd.level == 0:
                    code_str += gen_cmd_str(cmd)

        for k, v in arg_imports.items():
            imports += f"from {k} import " + ", ".join(v) + "\n"

        return imports + code_str
