import asyncio
import os
import re
import carb
import omni.ui
import omni.client
from typing import Callable, List, Tuple
from omni.kit.window.filepicker import FilePickerDialog
from omni.kit.widget.filebrowser import FileBrowserItem
from . import FileBrowserSelectionType, FileBrowserMode


LAST_LOCATION = "/persistent/app/omni.kit.widget.layers/last_file_picker_location"


class FileBrowserUI:
    def __init__(
        self,
        title: str,
        mode: FileBrowserMode,
        selection_type: FileBrowserSelectionType,
        filter_options: List[Tuple[str, str]],
        save_extensions: List[str] = [],
        allow_multi_selection=False,
        build_options_pane_fn: Callable[[List[FileBrowserItem]], bool] = None,
        on_selection_changed: Callable[[List[FileBrowserItem]], bool] = None,
        **kwargs
    ):
        if mode == FileBrowserMode.OPEN:
            confirm_text = "Open"
        else:
            confirm_text = "Save"
        self._file_picker = FilePickerApp(
            title,
            confirm_text,
            mode,
            selection_type,
            filter_options,
            save_extensions,
            allow_multi_selection,
            build_options_pane_fn,
            on_selection_changed,
            **kwargs
        )

    def destroy(self):
        self._file_picker.hide_dialog()
        self._file_picker.destroy()
        self._file_picker = None

    def set_current_directory(self, dir: str):
        self._file_picker.set_current_directory(dir)

    def set_current_filename(self, filename: str):
        self._file_picker.set_current_filename(filename)

    def get_current_filename(self):
        return self._file_picker.get_current_filename()

    def open(self, select_fn: Callable[[List[str]], None], cancel_fn: Callable[[], None]):
        self._file_picker.set_custom_fn(select_fn, cancel_fn)
        self._file_picker.show_dialog()

    def destroy(self):
        self._file_picker.destroy()
        self._file_picker = None

    def hide(self):
        self._file_picker.hide_dialog()


class FilePickerApp:
    """
    Standalone app to demonstrate the use of the FilePicker dialog.

    Args:
        title (str): Title of the window.
        apply_button_name (str): Name of the confirm button.
        mode (FileBrowserMode): The file picker mode that whether it's to open or save.
        selection_type (FileBrowserSelectionType): The file type that confirm event will respond to.
        item_filter_options (list): Array of filter options. Element of array
        is a tuple that first element of this tuple is the regex string for filtering,
        and second element of this tuple is the descriptions, like ("*.*", "All Files").
        By default, it will list all files.
        save_extensions: The real extenion name that will be saved.
        allow_multi_selections: Allow to select multiple files.
        build_options_pane_fn (Callable[[List[FileBrowserItem]], bool]): Function to
        build options panel.
        on_selection_changed (Callable[[List[FileBrowserItem]], bool]): Function to
        monitor selection changed.
    """

    def __init__(
        self,
        title: str,
        apply_button_name: str,
        mode: FileBrowserMode,
        selection_type: FileBrowserSelectionType = FileBrowserSelectionType.ALL,
        item_filter_options: list = [("*.*", "All Files (*.*)")],
        save_extensions: list = [".usd"],
        allow_multi_selections: bool = False,
        build_options_pane_fn: Callable[[List[FileBrowserItem]], bool] = None,
        on_selection_changed: Callable[[List[FileBrowserItem]], bool] = None,
        **kwargs
    ):
        self._settings = carb.settings.get_settings()
        self._settings.set_default(LAST_LOCATION, "")
        self._title = title
        self._filepicker = None
        self._mode = mode
        self._selection_type = selection_type
        self._custom_select_fn = None
        self._custom_cancel_fn = None
        self._apply_button_name = apply_button_name
        self._filter_regexes = []
        self._filter_descriptions = []
        self._current_directory = self._settings.get_as_string(LAST_LOCATION)
        self._allow_multi_selections = allow_multi_selections
        self._build_options_pane_fn = build_options_pane_fn
        self._selection_changed_fn = on_selection_changed
        for regex, desc in item_filter_options:
            if not isinstance(regex, re.Pattern):
                regex = re.compile(regex, re.IGNORECASE)
            self._filter_regexes.append(regex)
            self._filter_descriptions.append(desc)
        self._save_extensions = save_extensions
        self._build_ui(**kwargs)

    def destroy(self):
        self._settings = None
        self._custom_select_fn = None
        self._custom_cancel_fn = None
        self._build_options_pane_fn = None
        self._build_options_pane_fn = None
        self._selection_changed_fn = None
        self._on_selection_changed = None
        self.hide_dialog()
        self._filepicker.destroy()
        self._filepicker = None

    def set_custom_fn(self, select_fn: Callable[[List[str]], None], cancel_fn: Callable[[], None]):
        self._custom_select_fn = select_fn
        self._custom_cancel_fn = cancel_fn

    def show_dialog(self):
        self._filepicker.show(self._current_directory)

    def hide_dialog(self):
        self._filepicker.hide()

    def set_current_directory(self, dir: str):
        if not dir:
            return

        self._current_directory = dir
        if not self._current_directory.endswith("/"):
            self._current_directory += "/"

        self._settings.set_string(LAST_LOCATION, self._current_directory)

    def set_current_filename(self, filename: str):
        self._filepicker.set_filename(filename)

    def get_current_filename(self):
        return self._filepicker.get_filename()

    def _build_ui(self, **kwargs):
        # Create the dialog
        self._filepicker = FilePickerDialog(
            self._title,
            allow_multi_selection=self._allow_multi_selections,
            apply_button_label=self._apply_button_name,
            click_apply_handler=self._on_click_open,
            click_cancel_handler=self._on_click_cancel,
            item_filter_options=self._filter_descriptions,
            item_filter_fn=lambda item: self._on_filter_item(item),
            error_handler=lambda m: self._on_error(m),
            options_pane_build_fn=self._build_options_pane_fn,
            selection_changed_fn=self._on_selection_changed,
            **kwargs
        )

        if self._selection_type == FileBrowserSelectionType.DIRECTORY_ONLY:
            self._filepicker.set_filebar_label_name("Folder name")
        elif self._selection_type == FileBrowserSelectionType.FILE_ONLY:
            self._filepicker.set_filebar_label_name("File name")
        else:
            self._filepicker.set_filebar_label_name("File or Folder name")

        # Start off hidden
        self.hide_dialog()

    def _on_selection_changed(self, items: List[FileBrowserItem]):
        if self._selection_changed_fn:
            self._selection_changed_fn(items)

        if len(items) != 1:
            return False

        item = items[0]
        if item.is_folder and self._selection_type == FileBrowserSelectionType.FILE_ONLY:
            return False

        if not item.is_folder and self._selection_type == FileBrowserSelectionType.DIRECTORY_ONLY:
            return False

        self.set_current_filename(item.name)

    def _on_filter_item(self, item: FileBrowserItem) -> bool:
        if not item or item.is_folder:
            return True

        if self._selection_type == FileBrowserSelectionType.DIRECTORY_ONLY:
            return False

        if self._filepicker.current_filter_option >= len(self._filter_regexes):
            return False

        regex = self._filter_regexes[self._filepicker.current_filter_option]
        if regex.match(item.path):
            return True
        else:
            return False

    def _on_error(self, msg: str):
        pass

    def _on_click_open(self, filename: str, dirname: str):
        """
        The meat of the App is done in this callback when the user clicks 'Accept'. This is
        a potentially costly operation so we implement it as an async operation.  The inputs
        are the filename and directory name. Together they form the fullpath to the selected
        file.
        """
        if not dirname:
            return

        if self._mode == FileBrowserMode.SAVE and not filename:
            return

        selection_paths = self._filepicker.get_current_selections()
        if dirname:
            if filename:
                current_directory_path = self._filepicker.get_current_directory()

                # Stores current directory
                self._settings.set_string(LAST_LOCATION, current_directory_path)
                if current_directory_path.endswith("/"):
                    current_directory_path = current_directory_path[:-1]

                current_directory_name = os.path.basename(current_directory_path)
                # FIXME: It's possible that the file name in the input dialog
                # is the same as the current folder name, which is selected.
                # But FileDialog class will not include folders in selection_paths,
                # so here it's exclude that it's folder selection.
                if not selection_paths and current_directory_name == filename:
                    fullpath = f"{dirname}"
                else:
                    if not dirname.endswith("/"):
                        dirname = dirname + "/"
                    fullpath = omni.client.combine_urls(dirname, filename)
                    fullpath = fullpath.replace("\\", "/")
            else:
                fullpath = dirname
        else:
            fullpath = filename

        if self._selection_type == FileBrowserSelectionType.FILE_ONLY and not filename:
            return

        if not selection_paths:
            selection_paths = [fullpath]

        async def timeout_stat(path, timeout):
            try:
                result, entry = await asyncio.wait_for(omni.client.stat_async(path), timeout=timeout)
            except (Exception, asyncio.TimeoutError) as e:
                result = omni.client.Result.ERROR_NOT_FOUND
                entry = None

            return result, entry

        if not self._allow_multi_selections:
            result, entry = asyncio.get_event_loop().run_until_complete(timeout_stat(fullpath, timeout=1))
            if result == omni.client.Result.OK and entry.flags & omni.client.ItemFlags.CAN_HAVE_CHILDREN:
                is_folder = True
            else:
                is_folder = False

            if self._mode != FileBrowserMode.SAVE:
                if (is_folder and self._selection_type == FileBrowserSelectionType.FILE_ONLY) or (
                    not is_folder and self._selection_type == FileBrowserSelectionType.DIRECTORY_ONLY
                ):
                    return

        self.hide_dialog()

        if self._mode == FileBrowserMode.SAVE:
            _, ext = os.path.splitext(fullpath)
            if self._save_extensions and ext not in self._save_extensions:
                if self._filepicker.current_filter_option < len(self._save_extensions):
                    fullpath += self._save_extensions[self._filepicker.current_filter_option]
        if self._custom_select_fn:
            if self._allow_multi_selections and len(selection_paths) > 1:
                self._custom_select_fn(selection_paths)
            else:
                self._custom_select_fn([fullpath])

    def _on_click_cancel(self, filename: str, dirname: str):
        """
        This function is called when the user clicks 'Cancel'.
        """
        self.hide_dialog()

        if self._custom_cancel_fn:
            self._custom_cancel_fn()
