from typing import Dict, Optional

import carb
import omni.ui as ui
from omni.kit.welcome.window import show_window as show_welcome_window
from omni.ui import constant as fl

from .stage_tab import StageTab
from .style import OPEN_PAGE_STYLE


class OpenWidget:
    """
    Widget for open page.
    """
    def __init__(self):
        self.__tabs: Dict[str, callable] = {
            "STAGE": self.__build_stage,
            "SAMPLES": self.__build_sample,
            "FILE": self.__build_file,
        }
        self.__tab_frames: Dict[str, ui.Frame] = {}
        self.__stage_tab: Optional[StageTab] = None
        self._build_ui()

    def destroy(self):
        if self.__stage_tab:
            self.__stage_tab.destroy()
            self.__stage_tab = None

    def _build_ui(self):
        self._collection = ui.RadioCollection()
        with ui.ZStack(style=OPEN_PAGE_STYLE):
            # Background
            ui.Rectangle(style_type_name_override="Content")
            with ui.VStack():
                # Tab buttons
                with ui.VStack(height=fl._find("welcome_page_title_height")):
                    ui.Spacer()
                    with ui.HStack(height=0):
                        ui.Spacer()
                        with ui.HStack(spacing=30):
                            for tab in self.__tabs.keys():
                                ui.RadioButton(text=tab, width=0, spacing=8, image_width=24, image_height=24, name=tab, radio_collection=self._collection, clicked_fn=lambda t=tab: self._show_tab(t))
                        ui.Spacer()
                    ui.Spacer()
                # Tab frames
                with ui.ZStack():
                    for tab in self.__tabs.keys():
                        self.__tab_frames[tab] = ui.Frame(build_fn=lambda t=tab: self.__tabs[t](), visible=tab == "STAGE")

    def _show_tab(self, name: str) -> None:
        for tab in self.__tabs.keys():
            if tab == name:
                self.__tab_frames[name].visible = True
            else:
                self.__tab_frames[tab].visible = False

    def __build_stage(self):
        self.__stage_tab = StageTab()

    def __build_sample(self):
        try:
            import omni.kit.app
            manager = omni.kit.app.get_app().get_extension_manager()
            manager.set_extension_enabled_immediate("omni.kit.browser.sample", True)
            from omni.kit.browser.sample import get_instance as get_sample_instance
            from omni.kit.browser.sample.browser_widget import SampleTreeFolderBrowserWidget
            from omni.kit.browser.core import BrowserSearchBar

            # Remove options button
            class SampleWidget(SampleTreeFolderBrowserWidget):
                def _build_search_bar(self):
                    self._search_bar = BrowserSearchBar(options_menu=None)

            # Hack to close welcome window after sample opened
            def __open_sample(item):
                from .open_action import run_open_action
                run_open_action("omni.kit.window.file", "open_stage", item.url)

            model = get_sample_instance().get_model()
            model.execute = __open_sample

            # Sample widget
            with ui.ZStack():
                ui.Rectangle(style_type_name_override="Stage.Separator")
                with ui.VStack():
                    ui.Spacer(height=4)
                    with ui.HStack():
                        ui.Spacer(width=2)
                        self.__sample_widget = SampleWidget(model)
                    # Open widgets
                    with ui.ZStack(height=26):
                        ui.Rectangle(style_type_name_override="Stage.Separator")
                        with ui.VStack():
                            ui.Spacer(height=2)
                            with ui.HStack(spacing=2):
                                ui.Rectangle(style_type_name_override="Stage.Blank")
                                # Button size/style to be same as omni.kit.window.filepicker.FilePickerWidget
                                ui.Button("Open File", width=117, style_type_name_override="Stage.Button", mouse_pressed_fn=lambda x, y, a, f:self.__open_sample())
                                ui.Button("Close", width=117, style_type_name_override="Stage.Button", mouse_pressed_fn=lambda x, y, a, f:show_welcome_window(visible=False))
        except Exception as e:
            carb.log_error(f"Failed to import sample browser widget: {e}")

    def __build_file(self):
        try:
            from omni.kit.window.filepicker import FilePickerWidget

            def __open_file(file_name, dir_name):
                from .open_action import run_open_action
                run_open_action("omni.kit.window.file", "open_stage", dir_name + file_name)

            with ui.ZStack():
                ui.Rectangle(style_type_name_override="Stage.Separator")
                with ui.VStack():
                    ui.Spacer(height=2)
                    with ui.HStack():
                        ui.Spacer(width=2)
                        FilePickerWidget(
                            "Welcome",
                            enable_checkpoints=True,
                            apply_button_label="Open File",
                            cancel_button_label="Close",
                            click_apply_handler=__open_file,
                            click_cancel_handler=lambda f, d: show_welcome_window(visible=False),
                            modal=True,
                        )
        except Exception as e:
            carb.log_error(f"Failed to import file picker widget: {e}")

    def __open_sample(self):
        selections = self.__sample_widget.detail_selection
        if selections:
            self.__sample_widget._browser_model.execute(selections[0])
            show_welcome_window(visible=False)