# 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.client
import asyncio

from typing import Union, Callable, List, Optional
from carb import log_error, log_info
from collections import OrderedDict, namedtuple
from omni.kit.widget.filebrowser import FileBrowserItem
from omni.kit.window.filepicker import FilePickerModel, asset_types
from . import FILE_TYPE_USD, FILE_TYPE_IMAGE, FILE_TYPE_SOUND, FILE_TYPE_TEXT, FILE_TYPE_VOLUME

FileOpenAction = namedtuple("FileOpenAction", "name open_fn matching_type")


class ContentBrowserModel(FilePickerModel):
    """The model class for :obj:`ContentBrowserWidget`."""

    def __init__(self, item_filter_fn=None, timeout=3.0):
        super().__init__(timeout=timeout)
        self._file_open_dict = OrderedDict()
        self._on_selection_changed_subs = set()
        self._item_filter_fn = item_filter_fn

    def add_file_open_handler(self, name: str, open_fn: Callable, file_type: Union[int, Callable]) -> str:
        """
        Registers callback/handler to open a file of matching type.

        Args:
            name (str): Unique name of handler.
            open_fn (Callable): This function is executed when a matching file is selected for open, i.e. double clicked,
                right mouse menu open, or url submitted to browser bar.  Function signature: void open_fn(url: str).
            file_type (Union[int, func]): Can either be an enumerated int that is one of: [FILE_TYPE_USD, 
                FILE_TYPE_IMAGE, FILE_TYPE_SOUND, FILE_TYPE_TEXT, FILE_TYPE_VOLUME] or a more general boolean function
                that returns True if this function should be activated on the given file.  Function
                signature: bool file_type(url: str).

        Returns:
            str: Name if successful, None otherwise.

        """
        if name:
            self._file_open_dict[name] = FileOpenAction(name, open_fn, file_type)
            return name
        else:
            return None

    def delete_file_open_handler(self, name: str):
        """
        Unregisters the named file open handler.

        Args:
            name (str): Name of the handler.

        """
        if name and name in self._file_open_dict:
            self._file_open_dict.pop(name)

    def get_file_open_handler(self, url: str) -> Callable:
        """
        Returns the matching file open handler for the given item. 

        Args:
            url str: The url of the item to open.

        """
        if not url:
            return None
        broken_url = omni.client.break_url(url)

        # Note: It's important that we use an ordered dict so that we can search by order
        # of insertion.
        for _, tup in self._file_open_dict.items():
            _, open_fn, file_type = tup
            if isinstance(file_type, Callable):
                if file_type(broken_url.path):
                    return open_fn
            else:
                asset_type = {
                    FILE_TYPE_USD: asset_types.ASSET_TYPE_USD,
                    FILE_TYPE_IMAGE: asset_types.ASSET_TYPE_IMAGE,
                    FILE_TYPE_SOUND: asset_types.ASSET_TYPE_SOUND,
                    FILE_TYPE_TEXT: asset_types.ASSET_TYPE_SCRIPT,
                    FILE_TYPE_VOLUME: asset_types.ASSET_TYPE_VOLUME,
                }.get(file_type, None)

                if asset_type and asset_types.is_asset_type(broken_url.path, asset_type):
                    return open_fn

        return None

    async def open_stage_async(self, url: str, open_loadset=omni.usd.UsdContextInitialLoadSet.LOAD_ALL):
        """
        Async function for opening a USD file.

        Args:
            url (str): Url to file.

        """
        try:
            import omni.kit.window.file

            omni.kit.window.file.open_stage(url, open_loadset)
        except Exception as e:
            log_error(str(e))
        else:
            log_info(f"Success! Opened '{url}'.\n")

    async def open_with_new_edit_layer(self, url: str, open_loadset=omni.usd.UsdContextInitialLoadSet.LOAD_ALL):
        """
        Async function for opening a USD file, then creating a new layer.

        Args:
            url (str): Url to file.

        """
        try:
            import omni.kit.window.file

            omni.kit.window.file.open_with_new_edit_layer(url, open_loadset)
        except Exception as e:
            log_error(str(e))
        else:
            log_info(f"Success! Opened '{url}' with new edit layer.\n")

    def download_items_with_callback(self, src_urls: List[str], dst_urls: List[str], callback: Optional[Callable] = None):
        """
        Downloads specified items. Upon success, executes the given callback.

        Args:
            src_urls (List[str]): Paths of items to download.
            dst_urls (List[str]): Destination paths.
            callback (Callable): Callback to execute upon success.  Function signature is void callback([str]).

        Raises:
            :obj:`Exception`

        """
        try:
            self.copy_items_with_callback(src_urls, dst_urls, callback=callback)
        except Exception:
            raise

    def subscribe_selection_changed(self, fn: Callable):
        """
        Subscribes to file selection changes.

        Args:
            fn (Callable): callback function when file selection changed.

        """
        self._on_selection_changed_subs.add(fn)

    def unsubscribe_selection_changed(self, fn: Callable):
        """
        Unsubscribes this callback from selection changes.

        Args:
            fn (Callable): callback function when file selection changed.

        """
        if fn in self._on_selection_changed_subs:
            self._on_selection_changed_subs.remove(fn)

    def notify_selection_subs(self, pane: int, selected: [FileBrowserItem]):
        for fn in self._on_selection_changed_subs or []:

            async def async_cb(pane: int, selected: [FileBrowserItem]):
                # Making sure callback was not removed between async schedule/execute
                if fn in self._on_selection_changed_subs:
                    fn(pane, selected)

            asyncio.ensure_future(async_cb(pane, selected))

    def destroy(self):
        super().destroy()
        if self._file_open_dict:
            self._file_open_dict.clear()
            self._file_open_dict = None
        if self._on_selection_changed_subs:
            self._on_selection_changed_subs.clear()
            self._on_selection_changed_subs = None
