
from functools import partial
from typing import Callable, Dict, List, Optional

import omni.ui as ui
from omni.ui import constant as fl

from .style import ICON_PATH


class ThumbnailItem(ui.AbstractItem):
    """
    Data item for a thumbnail.
    Args:
        thumbnail (str): Url to thumbnail.
        url (str): Url to item.
        text (str): Display text for this item. Default use file name of url.
    """
    def __init__(self, thumbnail: str, url: str, text: Optional[str] = None):
        self.thumbnail = thumbnail
        self.url = url
        self.text = text if text else url.split("/")[-1]
        super().__init__()

    def execute(self):
        pass


class ThumbnailModel(ui.AbstractItemModel):
    """
    Data model for thumbnail items.
    Args:
        items (List[ThumbnailItem]): List of thumbnail items
    """
    def __init__(self, items: List[ThumbnailItem]):
        super().__init__()
        self._items = items

    def get_item_children(self, item: Optional[ThumbnailItem] = None) -> List[ThumbnailItem]:
        return self._items


class ThumbnailGrid:
    """
    Widget to show a grid (thumbnail + label).
    Args:
        model (ThumbnailModel): Data model.
        label_height (int): Height for item label. Default 30.
        thumbnail_width (int): Thumbnail width. Default 128.
        thumbnail_height (int): Thumbnail height. Default 128.
        thumbnail_padding_x (int): X padding between different thumbnail items. Default 2.
        thumbnail_padding_y (int): Y padding between different thumbnail items. Default 4.
    """
    def __init__(
            self,
            model: ThumbnailModel,
            on_selection_changed: Callable[[Optional[ThumbnailItem]], None] = None,
            label_height:int=fl._find("welcome_open_lable_size"),
            thumbnail_width: int=fl._find("welcome_open_thumbnail_size"),
            thumbnail_height: int=fl._find("welcome_open_thumbnail_size"),
            thumbnail_padding_x: int=2,
            thumbnail_padding_y: int=4,
        ):
        self._model = model
        self._on_selection_changed = on_selection_changed
        self._label_height = label_height
        self._thumbnail_width = thumbnail_width
        self._thumbnail_height = thumbnail_height
        self._thumbnail_padding_x = thumbnail_padding_x
        self._thumbnail_padding_y = thumbnail_padding_y

        self._column_width = self._thumbnail_width + 2 * self._thumbnail_padding_x
        self._row_height = self._thumbnail_height + 2 * self._thumbnail_padding_y + self._label_height

        self.__selected_thumbnail: Optional[ThumbnailItem] = None
        self.__frames: Dict[ThumbnailItem, ui.Frame] = {}

        self._build_ui()

    @property
    def selected(self) -> Optional[ThumbnailItem]:
        return self.__selected_thumbnail

    @selected.setter
    def selected(self, item: Optional[ThumbnailItem]) -> None:
        self.select_item(item)

    def _build_ui(self):
        self._container = ui.ScrollingFrame(
            style_type_name_override="ThumbnailGrid.Frame",
            vertical_scrollbar_policy=ui.ScrollBarPolicy.SCROLLBAR_AS_NEEDED,
            horizontal_scrollbar_policy=ui.ScrollBarPolicy.SCROLLBAR_ALWAYS_OFF,
        )
        self._refresh()
        self._sub_model = self._model.subscribe_item_changed_fn(lambda m, i: self._refresh())

    def _refresh(self):
        self._container.clear()
        with self._container:
            self._grid = ui.VGrid(
                column_width=self._column_width,
                row_hegiht=self._row_height,
                content_clipping=True,
            )
            with self._grid:
                children = self._model.get_item_children()
                if not children:
                    return
                for item in children:
                    with ui.VStack():
                        ui.Spacer(height=self._thumbnail_padding_y)
                        with ui.HStack():
                            ui.Spacer(width=self._thumbnail_padding_x)
                            self.__frames[item] = ui.Frame(
                                build_fn=lambda i=item: self._build_thumbnail_item(i),
                                mouse_double_clicked_fn=lambda x, y, a, f, i=item: i.execute(),
                                mouse_pressed_fn=lambda x, y, a, f, i=item: self.select_item(i),
                            )
                            ui.Spacer(width=self._thumbnail_padding_x)
                        ui.Spacer(height=self._thumbnail_padding_y)

            self._grid.set_mouse_pressed_fn(lambda x, y, btn, flag: self.select_item(None))

    def _build_thumbnail_item(self, item: ThumbnailItem) -> None:
        with ui.ZStack():
            with ui.VStack():
                self.build_thumbnail(item)
                self.build_label(item)
            ui.Rectangle(style_type_name_override="ThumbnailGrid.Selection")

    def build_thumbnail(self, item: ThumbnailItem) -> None:
        def on_image_progress(default_image: ui.Image, thumbnail_image: ui.Image, progress: float):
            """Called when the image loading progress is changed."""
            if progress != 1.0:
                # We only need to catch the moment when the image is loaded.
                return

            # Hide the default_image at the back.
            default_image.visible = False

            # Remove the callback to avoid circular references.
            thumbnail_image.set_progress_changed_fn(None)

        with ui.ZStack(width=self._thumbnail_width, height=self._thumbnail_height):
            # Default image for thumbnail not defined or invalid thumbnail url or download in progress
            # Hidden when actually thumbnail downloaded
            default_image = ui.Image(f"{ICON_PATH}/usd_stage_256.png", style_type_name_override="ThumbnailGrid.Image")
            if item.thumbnail:
                thumbnail_image = ui.Image(item.thumbnail, style_type_name_override="ThumbnailGrid.Image")
                thumbnail_image.set_progress_changed_fn(partial(on_image_progress, default_image, thumbnail_image))
            

    def build_label(self, item: ThumbnailItem) -> None:
        with ui.ZStack():
            ui.Rectangle(style_type_name_override="ThumbnailGrid.Item.Background")
            ui.Label(
                item.text,
                height=self._label_height,
                word_wrap=True,
                elided_text=True,
                skip_draw_when_clipped=True,
                alignment=ui.Alignment.CENTER,
                style_type_name_override="ThumbnailGrid.Item",
            )

    def select_item(self, item: Optional[ThumbnailItem]) -> None:
        if self.__selected_thumbnail:
            self.__frames[self.__selected_thumbnail].selected = False
        if item:
            self.__frames[item].selected = True
        self.__selected_thumbnail = item
        if self._on_selection_changed:
            self._on_selection_changed(item)

    def find_item(self, url: str) -> Optional[ThumbnailItem]:
        for item in self._model.get_item_children(None):
            if item.url == url:
                return item
        return None

    def select_url(self, url: str) -> None:
        if url:
            item = self.find_item(url)
            self.select_item(item)
        else:
            self.select_item(None)