from typing import List
from collections import defaultdict, OrderedDict
from omni.kit.window.property.templates import SimplePropertyWidget, LABEL_WIDTH, LABEL_HEIGHT, HORIZONTAL_SPACING
from omni.kit.property.usd.usd_property_widget import (
    UsdPropertiesWidget,
    UsdPropertyUiEntry,
    UsdPropertiesWidgetBuilder,
)
from omni.kit.property.usd.usd_attribute_model import UsdAttributeModel, TfTokenAttributeModel

import asyncio
import carb
import omni.ui as ui
import omni.usd

from pxr import Usd, Tf, Vt, Sdf, UsdShade


class SubIdentifierUtils:
    class AllowedAnnoItem(ui.AbstractItem):
        def __init__(self, item, token=None):
            from omni.kit.material.library import MaterialLibraryExtension

            super().__init__()
            if isinstance(item, MaterialLibraryExtension.SubIDEntry):
                if "subid_token" in item.annotations:
                    self.token = item.annotations["subid_token"]
                else:
                    self.token = item.name
                self.model = ui.SimpleStringModel(item.name)
                if "subid_display_name" in item.annotations:
                    self.model = ui.SimpleStringModel(item.annotations['subid_display_name'])
                elif "display_name" in item.annotations:
                    self.model = ui.SimpleStringModel(item.annotations['display_name'])
                else:
                    self.model = ui.SimpleStringModel(item.name)
            else:
                if token:
                    self.token = token
                    self.model = ui.SimpleStringModel(item)

                else:
                    self.token = item
                    self.model = ui.SimpleStringModel(item)


    class SubIdentifierModel(TfTokenAttributeModel):
        def __init__(
            self,
            stage: Usd.Stage,
            attribute_paths: List[Sdf.Path],
            self_refresh: bool,
            metadata: dict,
            options: list,
            attr: Usd.Attribute,
        ):
            self._widget = None
            self._combobox_options = [attr.Get()]
            super().__init__(stage=stage, attribute_paths=attribute_paths, self_refresh=self_refresh, metadata=metadata)
            self._has_index = False

            async def load_subids():
                await omni.kit.app.get_app().next_update_async()
                await omni.kit.material.library.get_subidentifier_from_material(prim=attr.GetPrim(), on_complete_fn=self._have_list, use_functions=True)

            asyncio.ensure_future(load_subids())

        def _get_allowed_tokens(self, attr):
            return self._combobox_options

        def _update_allowed_token(self):
            super()._update_allowed_token(SubIdentifierUtils.AllowedAnnoItem)

        def _update_value(self, force=False):
            from omni.kit.property.usd.usd_model_base import UsdBase

            was_updating_value = self._updating_value
            self._updating_value = True
            if UsdBase._update_value(self, force):
                # TODO don't have to do this every time. Just needed when "allowedTokens" actually changed
                self._update_allowed_token()

                def find_allowed_token(value):
                    # try to match the full token, i.e. simple name with function parameters
                    for i in range(0, len(self._allowed_tokens)):
                        if self._allowed_tokens[i].token == value:
                            return i
                    # if the above failed, drop the parameter list of the query
                    # try to find a match based on the simple name alone
                    # note, because of overloads there can be more than one match
                    query = value.split("(", 1)[0]
                    match_count = 0
                    first_match_index = -1
                    for i in range(0, len(self._allowed_tokens)):
                        if self._allowed_tokens[i].token.split("(", 1)[0] == query:
                            if match_count == 0:
                                first_match_index = i
                            match_count += 1
                    # if there is one match based on simple name, we can safely return this one
                    if match_count == 1:
                        return first_match_index
                    # the match is not unique, we need to return a `<not found>`
                    else:
                        return -1

                index = find_allowed_token(self._value)
                if self._value and index == -1:
                    carb.log_warn(f"failed to find '{self._value}' in function name list")
                    index = len(self._allowed_tokens)
                    self._allowed_tokens.append(SubIdentifierUtils.AllowedAnnoItem(token=self._value, item=f"<sub-identifier not found: '{self._value}'>"))

                if index != -1 and self._current_index.as_int != index:
                    self._current_index.set_value(index)
                    self._item_changed(None)
            self._updating_value = was_updating_value


        def _have_list(self, mtl_list: list):
            if mtl_list:
                self._combobox_options = mtl_list
                self._set_dirty()
                if self._has_index is False:
                    self._update_value()
                    self._has_index = True

    def build_subidentifier(
        stage,
        attr_name,
        type_name,
        metadata,
        attr: Usd.Attribute,
        prim_paths: List[Sdf.Path],
        additional_label_kwargs=None,
        additional_widget_kwargs=None,
    ):
        with ui.HStack(spacing=HORIZONTAL_SPACING):
            model = None
            UsdPropertiesWidgetBuilder._create_label(attr_name, metadata, additional_label_kwargs)
            tokens = metadata.get("allowedTokens")
            if not tokens:
                tokens = [attr.Get()]
            model = SubIdentifierUtils.SubIdentifierModel(
                stage,
                attribute_paths=[path.AppendProperty(attr_name) for path in prim_paths],
                self_refresh=False,
                metadata=metadata,
                options=tokens,
                attr=attr,
            )

            widget_kwargs = {"name": "choices", "no_default": True}
            if additional_widget_kwargs:
                widget_kwargs.update(additional_widget_kwargs)
            with ui.ZStack():
                value_widget = ui.ComboBox(model, **widget_kwargs)
                value_widget.identifier = 'sdf_asset_info:mdl:sourceAsset:subIdentifier'
                mixed_overlay = UsdPropertiesWidgetBuilder._create_mixed_text_overlay()
            UsdPropertiesWidgetBuilder._create_control_state(
                model=model, value_widget=value_widget, mixed_overlay=mixed_overlay, **widget_kwargs
            )
            return model


    def annotations_build_fn(
        stage,
        attr_name,
        metadata,
        property_type,
        prim_paths: List[Sdf.Path],
        additional_label_kwargs,
        additional_widget_kwarg,
        material_annotations
    ):
        additional_label_kwargs = {"width": ui.Fraction(1), "read_only": True}

        description = material_annotations["description"]
        description = description.replace("\n\r", "\r").replace("\n", "\r")
        description_list = description.split("\r")
        if len(description_list) > 1:
            with ui.VStack(spacing=2):
                for description in description_list:
                    if not description:
                        continue

                    parts = description.split("\t", maxsplit=1)
                    if len(parts) == 1:
                        parts = description.split(" ", maxsplit=1)

                    with ui.HStack(spacing=HORIZONTAL_SPACING):
                        if len(parts) == 1:
                            parts.append("")
                        UsdPropertiesWidgetBuilder._create_label(parts[0])
                        if parts[1].startswith("http"):
                            additional_label_kwargs["tooltip"] = f"Click to open \"{parts[1]}\""
                        else:
                            additional_label_kwargs["tooltip"] = parts[1]
                        UsdPropertiesWidgetBuilder._create_text_label(parts[1], additional_label_kwargs=additional_label_kwargs)
        else:
            with ui.HStack(spacing=HORIZONTAL_SPACING):
                UsdPropertiesWidgetBuilder._create_label("Description")
                additional_label_kwargs["tooltip"] = description
                UsdPropertiesWidgetBuilder._create_text_label(description, additional_label_kwargs=additional_label_kwargs)

    def get_annotations(
        annotation_key,
        prim,
        material_annotation_list,
    ):
        if prim.IsA(UsdShade.Material):
            shader = omni.usd.get_shader_from_material(prim, False)
            attr = shader.GetPrim().GetAttribute("info:mdl:sourceAsset:subIdentifier") if shader else None
        else:
            attr = prim.GetAttribute("info:mdl:sourceAsset:subIdentifier") if prim else None
            shader = UsdShade.Shader(prim)

        if attr:
            asset = shader.GetSourceAsset("mdl") if shader else None
            mdl_file = asset.resolvedPath if asset else None
            if mdl_file and mdl_file in material_annotation_list:
                material_annotations = material_annotation_list[mdl_file]
                if attr.Get() in material_annotations:
                    material_annotations = material_annotations[attr.Get()]
                    if annotation_key in material_annotations:
                        return material_annotations

        return None
