import asyncio
import importlib.util
import pathlib
import weakref
from contextlib import suppress
from typing import List

import omni.graph.core as og
import omni.graph.tools.ogn as ogn
import omni.kit.window.property as p
import omni.ui as ui
import omni.usd
from omni.kit.property.usd.custom_layout_helper import CustomLayoutGroup, CustomLayoutProperty
from omni.kit.property.usd.usd_attribute_widget import UsdAttributeUiEntry, UsdPropertiesWidget, UsdPropertyUiEntry
from omni.kit.window.property.templates import HORIZONTAL_SPACING, LABEL_WIDTH
from pxr import Sdf, Usd

from .omnigraph_attribute_builder import OmniGraphPropertiesWidgetBuilder  # noqa: PLE0402
from .omnigraph_attribute_builder import find_first_node_with_attrib  # noqa: PLE0402

_compute_node_widget_instance = None


class ComputeNodeWidget(UsdPropertiesWidget):
    """Creates and registers a widget in the Property Window to display Compute Node contents"""

    # self.bundles : dict[bundle_name : str, tuple[omni::graph::core::Py_Bundle, list[attrib_name : str]]]
    BUNDLE_INDEX = 0
    ATTRIB_LIST_INDEX = 1

    def __init__(self, title="OmniGraph Node", collapsed=False):
        super().__init__(title=title, collapsed=collapsed)
        self.context_menu = ui.Menu("ContextMenu")
        self.bundle_update_sub = None
        self.bundles = {}
        self.graph_context: og.GraphContext = None
        self.graph: og.Graph = None
        self.node: og.Node = None
        self.selection = omni.usd.get_context().get_selection()
        self.stage: Usd.Stage = None
        self.stage_update_sub = None
        self.template = None
        self.template_paths = []
        self.time = 0
        self.window = None
        self.window_width_change = False
        self.window_width_change_fn_set = False
        self.node_type_name = None
        self.node_event_sub = None

        self._payload = None

        w = p.get_window()
        w.register_widget("prim", "compute_node", self)
        global _compute_node_widget_instance
        _compute_node_widget_instance = weakref.ref(self)

        OmniGraphPropertiesWidgetBuilder.startup()

    @staticmethod
    def get_instance():
        return _compute_node_widget_instance()

    def on_shutdown(self):
        w = p.get_window()
        w.unregister_widget("prim", "compute_node")

    def reset(self):
        super().reset()

        window = ui.Workspace.get_window("Property")
        if self.window != window:
            self.window_width_change_fn_set = False
        self.window = window
        self.window_width_change = False

    def __on_node_event(self, event):
        # Called when one of the node event happens
        self.request_rebuild()

    def on_new_payload(self, payload: List[Sdf.Path]) -> bool:
        """
        See PropertyWidget.on_new_payload
        """
        if self.bundle_update_sub:
            self.bundle_update_sub = None
        if self.stage_update_sub:
            self.stage_update_sub = None
        if len(payload) == 0:
            return False
        self.stage = omni.usd.get_context().get_stage()
        self.node = None
        for path in reversed(payload):
            node = og.get_node_by_path(path.pathString)
            if (
                node
                and node.is_valid()
                and self.stage.GetPrimAtPath(node.get_prim_path()).GetTypeName() in ("ComputeNode", "OmniGraphNode")
            ):
                self.node = node
                self.graph = self.node.get_graph()
                self.node_event_sub = node.get_event_stream().create_subscription_to_pop(
                    self.__on_node_event, name=f"{path.pathString} Event"
                )
                break
        if self.node is None:
            return False

        self.node_type_name = self.node.get_type_name()

        self._payload = payload
        self.graph_context = self.graph.get_default_graph_context() if self.graph.is_valid() else None

        self.bundles = {}
        return True

    def get_additional_kwargs(self, ui_attr: UsdAttributeUiEntry):
        additional_label_kwargs = {"alignment": ui.Alignment.RIGHT}
        additional_widget_kwargs = {"no_mixed": True}
        return additional_label_kwargs, additional_widget_kwargs

    def get_widget_prim(self):
        return self.stage.GetPrimAtPath(self._payload[-1])

    def add_template_path(self, file_path):
        """Makes a path from the file_path parameter and adds it to templates paths.
        So if there is a same kind of structure:
            /templates
                template_<ModuleName>.<NameOfComputeNodeType>.py
            <my_extension>.py

        then template path can be added by this line in the <my_extension>.py:
            omni.graph.ui.ComputeNodeWidget.get_instance().add_template_path(__file__)

        example of template filename:
            template_omni.particle.system.core.Emitter.py

        """
        template_path = pathlib.Path(file_path).parent.joinpath(pathlib.Path("templates"))

        if template_path not in self.template_paths:
            self.template_paths.append(template_path)

    def get_template_path(self, template_name):
        for path in self.template_paths:
            path = path / template_name
            if path.is_file():
                return path
        return None

    def load_template(self, props):
        if not self.node_type_name:
            return None
        if self.node_type_name.startswith("omni.graph.ui."):
            template_name = f"template_{self.node_type_name[14:]}.py"
        else:
            template_name = f"template_{self.node_type_name}.py"
        path = self.get_template_path(template_name)
        if path is None:
            return None
        spec = importlib.util.spec_from_file_location(template_name, path)
        module = importlib.util.module_from_spec(spec)
        spec.loader.exec_module(module)
        self.template = module.CustomLayout(self)
        if not self.template.enable:
            return None
        return self.template.apply(props)

    async def __delayed_build_layout(self):
        await omni.kit.app.get_app().next_update_async()
        self.window.frame.rebuild()

    def rebuild_window(self):
        asyncio.ensure_future(self.__delayed_build_layout())

    def _customize_props_layout(self, props):
        for prop in props:
            if self.node.get_attribute_exists(prop.prop_name):
                node_attr = self.node.get_attribute(prop.prop_name)
                description = node_attr.get_metadata(ogn.MetadataKeys.DESCRIPTION)
                if description:
                    prop.metadata[Sdf.PropertySpec.DocumentationKey] = description
                for key in ogn.MetadataKeys.key_names():
                    prop.add_custom_metadata(key, node_attr.get_metadata(key))
        template = self.load_template(props)
        if template is not None:
            return template

        parameter_attrs = []
        relationship_attrs = []
        hidden_props = []

        for prop in props:
            prop_name = prop.prop_name

            if prop.property_type == Usd.Relationship:
                prop.override_display_group("Relationships")
                relationship_attrs.append(prop)
                continue
            if not self.node.get_attribute_exists(prop.prop_name):
                hidden_props.append(prop)
                continue
            is_parameter = False
            attribute = self.node.get_attribute(prop.prop_name)
            if attribute.get_metadata("displayGroup") == "parameters":
                prop.override_display_group("Parameters")

                parameter_attrs.append(prop)
                is_parameter = True

            if attribute.get_metadata(ogn.MetadataKeys.HIDDEN) is not None:
                hidden_props.append(prop)
                continue

            if attribute.get_resolved_type().role == og.AttributeRole.EXECUTION:
                hidden_props.append(prop)
                continue

            if attribute.get_metadata(ogn.MetadataKeys.OBJECT_ID) is not None:
                hidden_props.append(prop)
                continue

            if prop_name.startswith("inputs:"):
                prop_name = prop_name.replace("inputs:", "")
                prop.override_display_name(prop_name)

                if not is_parameter:
                    prop.override_display_group("Inputs")
                    continue

            elif prop_name.startswith("outputs:"):
                prop_name = prop_name.replace("outputs:", "")
                prop.override_display_name(prop_name)
                if not is_parameter:
                    prop.override_display_group("Outputs")
                    continue

            elif prop_name.startswith("state:"):
                prop_name = prop_name.replace("state:", "")
                prop.override_display_name(prop_name)
                if not is_parameter:
                    prop.override_display_group("State")
                    continue

            # Hide everything else
            if not is_parameter:
                hidden_props.append(prop)

        filtered_props = props.copy()
        for hidden_prop in hidden_props:
            filtered_props.remove(hidden_prop)

        self.apply_ogn_metadata_display_names(filtered_props)

        reordered_attrs = filtered_props
        self.move_elements_to_beginning(relationship_attrs, reordered_attrs)
        self.move_elements_to_beginning(parameter_attrs, reordered_attrs)
        return reordered_attrs

    def width_changed_subscribe(self):
        self.window = ui.Workspace.get_window("Property")
        self.window_width_change = True
        if not self.window_width_change_fn_set:
            self.window.set_width_changed_fn(self.width_changed)
            self.window_width_change_fn_set = True

    def list_diff(self, list_a, list_b):
        return [i for i in list_a + list_b if i not in list_a or i not in list_b]

    def apply_ogn_metadata_display_names(self, attrs):
        for attr in attrs:
            if not self.node.get_attribute_exists(attr.attr_name):
                continue
            node_attr = self.node.get_attribute(attr.attr_name)
            if node_attr is None:
                continue
            ogn_display_name = node_attr.get_metadata(ogn.MetadataKeys.UI_NAME)
            if ogn_display_name is not None:
                attr.override_display_name(ogn_display_name)

    def move_elements_to_beginning(self, elements_list, target_list):
        elements_list.reverse()
        for element in elements_list:
            target_list.remove(element)
        for element in elements_list:
            target_list.insert(0, element)

    def build_property_item(self, stage, ui_prop: UsdPropertyUiEntry, prim_paths: List[Sdf.Path]):
        """Override of base - intercepts all attribute widget building"""
        with ui.HStack():
            # override prim paths to build if UsdPropertyUiEntry specifies one
            if ui_prop.prim_paths:
                prim_paths = ui_prop.prim_paths

            # Use supplied build_fn or our customized builder
            build_fn = ui_prop.build_fn if ui_prop.build_fn else OmniGraphPropertiesWidgetBuilder.build
            additional_label_kwargs, additional_widget_kwargs = self.get_additional_kwargs(ui_prop)

            # Handle extended attributes
            if ui_prop.property_type is Usd.Attribute:
                attr_name = ui_prop.attr_name

                # Find a selected node that actually has this attribute in order to interrogate it
                node_with_attr = find_first_node_with_attrib(prim_paths, attr_name)
                if node_with_attr:
                    ogn_attr = node_with_attr.get_attribute(attr_name)

                    is_extended_attribute = ogn_attr.get_extended_type() in (
                        og.ExtendedAttributeType.EXTENDED_ATTR_TYPE_ANY,
                        og.ExtendedAttributeType.EXTENDED_ATTR_TYPE_UNION,
                    )

                    if is_extended_attribute:
                        resolved_type = ogn_attr.get_resolved_type()
                        if resolved_type.base_type == og.BaseDataType.UNKNOWN:
                            # This is an unresolved attribute
                            extended_type = (
                                "any"
                                if ogn_attr.get_extended_type() == og.ExtendedAttributeType.EXTENDED_ATTR_TYPE_ANY
                                else "union"
                            )
                            ui_prop.metadata[
                                OmniGraphPropertiesWidgetBuilder.OVERRIDE_TYPENAME_KEY
                            ] = f"unresolved {extended_type}"
                            ui_prop.metadata[Sdf.PrimSpec.TypeNameKey] = "unknown"
                        else:
                            # Resolved attribute - swap in the resolved type to the ui_prop
                            sdf_type = og.AttributeType.sdf_type_name_from_type(resolved_type)
                            ui_prop.metadata[Sdf.PrimSpec.TypeNameKey] = sdf_type

            models = build_fn(
                stage,
                ui_prop.prop_name,
                ui_prop.metadata,
                ui_prop.property_type,
                prim_paths,
                additional_label_kwargs,
                additional_widget_kwargs,
            )
            if models:
                if not isinstance(models, list):
                    models = [models]
                for model in models:
                    for prim_path in prim_paths:
                        self._models[prim_path.AppendProperty(ui_prop.prop_name)].append(model)

    def get_bundles(self):
        if self.graph_context is None:
            return
        inputs_particles_rel = self.stage.GetRelationshipAtPath(self._payload[-1].pathString + ".inputs:particles")
        if inputs_particles_rel.IsValid():
            input_prims = inputs_particles_rel.GetTargets()
            for input_prim in input_prims:
                self.bundles[input_prim.GetParentPath().name + ".inputs:particles"] = (
                    self.graph_context.get_bundle(input_prim.pathString),
                    [],
                )
        self.bundles["outputs:particles"] = (
            self.graph_context.get_bundle(self._payload[-1].pathString + "/outputs_particles"),
            [],
        )

    def bundle_elements_layout_fn(self, **kwargs):
        with ui.HStack(spacing=HORIZONTAL_SPACING):
            ui.Label(kwargs["name"], name="label", style={"alignment": ui.Alignment.RIGHT_TOP}, width=LABEL_WIDTH)
            ui.Spacer(width=HORIZONTAL_SPACING)
            model = ui.StringField(name="models_readonly", read_only=True, enabled=False).model
            self.bundles[kwargs["bundle_name"]][self.ATTRIB_LIST_INDEX].append(model)

    def display_bundles_content(self):
        self.get_bundles()
        self.bundle_update_sub = (
            omni.kit.app.get_app()
            .get_update_event_stream()
            .create_subscription_to_pop(self.on_bundles_update, name="bundles ui update")
        )
        self.stage_update_sub = (
            omni.usd.get_context()
            .get_stage_event_stream()
            .create_subscription_to_pop(self.on_stage_update, name="bundles ui stage update")
        )

        for bundle_name, bundle_value in self.bundles.items():
            names, _ = bundle_value[self.BUNDLE_INDEX].get_attribute_names_and_types()
            data_count = bundle_value[self.BUNDLE_INDEX].get_attribute_data_count()
            if data_count == 0:
                continue

            with CustomLayoutGroup(bundle_name, collapsed=True):
                for i in range(0, data_count):
                    attr_name = names[i]

                    def fn(*args, n=attr_name, b=bundle_name):
                        self.bundle_elements_layout_fn(name=n, bundle_name=b)

                    CustomLayoutProperty(None, None, build_fn=fn)

    def set_bundle_models_values(self):
        if self.graph_context is None:
            return
        for _bundle_name, bundle_value in self.bundles.items():
            if not bundle_value[1]:
                continue

            _, types = bundle_value[self.BUNDLE_INDEX].get_attribute_names_and_types()
            data = bundle_value[self.BUNDLE_INDEX].get_attribute_data()
            for i, d in enumerate(data):
                with suppress(Exception):
                    attr_type_name = types[i].get_type_name()
                    elem_count = self.graph_context.get_elem_count(d)
                    bundle_value[self.ATTRIB_LIST_INDEX][i].set_value(
                        attr_type_name + " " + str(elem_count) + " elements"
                    )
                    bundle_value[self.ATTRIB_LIST_INDEX][i]._value_changed()  # noqa: PLW0212

    def on_bundles_update(self, event):
        self.time += event.payload["dt"]
        if self.time > 0.25:
            self.set_bundle_models_values()
            self.time = 0

    def on_stage_update(self, event):
        if event.type == int(omni.usd.StageEventType.CLOSING) and self.bundle_update_sub:
            self.bundle_update_sub = None  # unsubscribe from app update stream
