# 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 os
from collections import defaultdict
from typing import List, Tuple

import carb
import omni.kit.app
import omni.ui as ui
from omni.kit.widget.graph import GraphNodeDescription, GraphPortDescription
from omni.kit.widget.graph.abstract_graph_node_delegate import AbstractGraphNodeDelegate
from omni.kit.widget.graph.graph_model import GraphModel
from omni.kit.widget.graph.graph_view import GraphView

from .utils import ext_id_to_fullname

# pylint: disable=property-with-parameters


# Colors & Style
BACKGROUND_COLOR = 0xFF34302A
BORDER_DEFAULT = 0xFF232323
CONNECTION = 0xFF80C280
NODE_BACKGROUND = 0xFF141414
PLUGINS_COLOR = 0xFFFFE3C4
LIBRARIES_COLOR = 0xAA44EBE7
MODULES_COLOR = 0xFFAAE3C4

GRAPH_STYLE = {
    "Graph": {"background_color": BACKGROUND_COLOR},
    "Graph.Connection": {"color": CONNECTION, "background_color": CONNECTION, "border_width": 2.0},
    # Node
    "Graph.Node.Background": {"background_color": NODE_BACKGROUND},
    "Graph.Node.Border": {"background_color": BORDER_DEFAULT},
    # Header
    "Graph.Node.Header.Label": {"color": 0xFFEEEEEE, "margin_height": 5.0, "font_size": 14.0},
}

# Constants
MARGIN_WIDTH = 7.5
MARGIN_TOP = 20.0
MARGIN_BOTTOM = 25.0

BORDER_THICKNESS = 3.0
HEADER_HEIGHT = 25.0
MIN_WIDTH = 180.0

CONNECTION_CURVE = 60


class GraphNodeDelegate(AbstractGraphNodeDelegate):
    """
    The delegate with the Omniverse design.
    """

    def __init__(self, scale_factor=1.0):
        self._manager = omni.kit.app.get_app_interface().get_extension_manager()
        self._scale_factor = scale_factor

    def __scale(self, value):
        """Return the value multiplied by global scale multiplier"""
        return value * self._scale_factor

    def set_scale_factor(self, scale_factor):
        """Replace scale factor"""
        self._scale_factor = scale_factor

    def node_background(self, model, node_desc: GraphNodeDescription):
        """Called to create widgets of the node background"""

        # Computed values
        left_right_offset = MARGIN_WIDTH - BORDER_THICKNESS * 0.5

        # Draw a rectangle and a top line
        with ui.HStack():
            # Left offset
            ui.Spacer(width=self.__scale(left_right_offset))

            with ui.VStack():
                ui.Spacer(height=self.__scale(MARGIN_TOP))

                # The node body
                with ui.ZStack():
                    # This trick makes min width
                    ui.Spacer(width=self.__scale(MIN_WIDTH))

                    # Build outer rectangle
                    ui.Rectangle(style_type_name_override="Graph.Node.Border")

                    # Build inner rectangle
                    with ui.VStack():
                        ui.Spacer(height=self.__scale(HEADER_HEIGHT))
                        with ui.HStack():
                            ui.Spacer(width=self.__scale(BORDER_THICKNESS))
                            ui.Rectangle(style_type_name_override="Graph.Node.Background")
                            ui.Spacer(width=self.__scale(BORDER_THICKNESS))
                        ui.Spacer(height=self.__scale(BORDER_THICKNESS))

                ui.Spacer(height=self.__scale(MARGIN_BOTTOM))

            # Right offset
            ui.Spacer(width=self.__scale(left_right_offset))

    def node_header_input(self, model, node_desc: GraphNodeDescription):
        """Called to create the left part of the header that will be used as input when the node is collapsed"""
        with ui.ZStack(width=self.__scale(8)):
            if node_desc.connected_target:
                # Circle that shows that the port is a target for the connection
                ui.Circle(
                    radius=self.__scale(4),
                    size_policy=ui.CircleSizePolicy.FIXED,
                    style_type_name_override="Graph.Connection",
                    alignment=ui.Alignment.RIGHT_CENTER,
                )

    def node_header_output(self, model, node_desc: GraphNodeDescription):
        """Called to create the right part of the header that will be used as output when the node is collapsed"""
        with ui.ZStack(width=self.__scale(8)):
            if node_desc.connected_source:
                # Circle that shows that the port is a source for the connection
                ui.Circle(
                    radius=self.__scale(4),
                    size_policy=ui.CircleSizePolicy.FIXED,
                    style_type_name_override="Graph.Connection",
                    alignment=ui.Alignment.LEFT_CENTER,
                )

    def node_header(self, model, node_desc: GraphNodeDescription):
        """Called to create widgets of the top of the node"""

        item = model[node_desc.node].item

        def build_list(files, title, color):
            if len(files) > 0:
                style = {"CollapsableFrame": {"color": color, "background_color": NODE_BACKGROUND}}
                with ui.CollapsableFrame(title, collapsed=False, style=style):
                    with ui.VStack(height=0):
                        for p in files:
                            ui.Label("  - " + os.path.basename(p), height=self.__scale(20)).set_tooltip(p)

        # Draw the node name and a bit of space
        with ui.VStack(height=0):
            ui.Spacer(height=self.__scale(18))
            with ui.HStack(height=0):
                ui.Spacer(width=self.__scale(18))
                title = item.id
                ui.Label(title, style_type_name_override="Graph.Node.Header.Label")
            build_list(item.plugins, "carb plugins", PLUGINS_COLOR)
            build_list(item.libraries, "shared libraries", LIBRARIES_COLOR)
            build_list(item.modules, "python modules", MODULES_COLOR)

            ui.Spacer(height=self.__scale(55))

    def node_footer(self, model, node_desc: GraphNodeDescription):
        pass

    def port_input(self, model, node_desc: GraphNodeDescription, port_desc: GraphPortDescription):
        pass

    def port_output(self, model, node_desc: GraphNodeDescription, port_desc: GraphPortDescription):
        pass

    def port(self, model, node_desc: GraphNodeDescription, port_desc: GraphPortDescription):
        pass

    def connection(self, model, source, target):
        """Called to create the connection between ports"""
        # If the connection is reversed, we need to mirror tangents
        connection_direction_is_same_as_flow = source.level >= target.level
        reversed_tangent = -1.0 if connection_direction_is_same_as_flow else 1.0
        ui.FreeBezierCurve(
            target.widget,
            source.widget,
            start_tangent_width=ui.Percent(CONNECTION_CURVE * reversed_tangent),
            end_tangent_width=ui.Percent(-CONNECTION_CURVE * reversed_tangent),
            style_type_name_override="Graph.Connection",
        )


class ExtGraphItem:
    def __init__(self, manager, ext_id):
        info = manager.get_extension_dict(ext_id)

        # Extension dependencies:
        self.deps = info.get("state/dependencies", [])
        self.port_deps = [d + "/out" for d in self.deps]

        # Extension useful info (id, version, plugins, modules etc)
        state_dict = info.get("state", {})
        native_dict = state_dict.get("native", {})
        self.id = ext_id
        self.plugins = native_dict.get("plugins", [])
        self.libraries = native_dict.get("libraries", [])
        self.modules = state_dict.get("python", {}).get("modules", [])

    def __lt__(self, other):
        return self.id < other.id


class Model(GraphModel):
    def __init__(self, ext_id: str = None):
        super().__init__()

        # build graph out of enabled extensions
        manager = omni.kit.app.get_app_interface().get_extension_manager()
        exts = manager.get_extensions()
        enabled_exts = [e["id"] for e in exts if e["enabled"]]
        self._nodes = {ext_id: ExtGraphItem(manager, ext_id) for ext_id in enabled_exts}

        # If ext_id was passed filter out extension that are not reachable from this one. To show only its graph:
        self._root_ext_id = ext_id
        if self._root_ext_id:
            self._filter_out_unreachable_nodes(ext_id)

    def _filter_out_unreachable_nodes(self, ext_id):
        visited = set()
        q = []

        root = self._nodes.get(ext_id, None)
        if not root:
            carb.log_error(f"Failure to filter dep graph, can't find ext node: {ext_id}")
            return

        q.append(root)
        visited.add(root)

        while len(q) > 0:
            node = q.pop()
            for d in node.deps:
                child = self._nodes[d]
                if child not in visited:
                    visited.add(child)
                    q.append(child)

        self._nodes = {ext_id: item for ext_id, item in self._nodes.items() if item in visited}

    @property
    def expansion_state(self, item=None):
        return self.ExpansionState.CLOSED

    @property
    def nodes(self, item=None):
        return {e.id for e in self._nodes.values()}

    @property
    def name(self, item=None):
        return self._nodes[item].id

    @property
    def item(self, item=None):
        return self._nodes[item]

    @property
    def ports(self, item=None):
        return [item + "/in", item + "/out"]

    @property
    def inputs(self, item):
        if item.endswith("/in"):
            item = self._nodes.get(item[:-3], None)
            return item.port_deps if item else []
        return None

    @property
    def outputs(self, item):
        outputs = [] if item.endswith("/out") else None
        return outputs

    def copy_as_graphviz(self):

        output = ""

        output += """
digraph Extensions {
    graph [outputMode=nodesfirst rankdir=LR]
    node [shape=box style=filled]
        """

        for _, node in self._nodes.items():
            label = node.id + "\\n-------------------\\n"

            def gen(files, title):
                if len(files) == 0:
                    return ""
                return f"{title}:\\l" + "".join([f" * {os.path.basename(f)}\\l" for f in files])

            label += gen(node.plugins, "plugins")
            label += gen(node.libraries, "libraries")
            label += gen(node.modules, "modules")
            output += f'\t\t"{node.id}" [label="{label}"]\n'

            for d in node.deps:
                output += f'\t\t"{d}" -> "{node.id}"\n'

        output += "}\n"

        print(output)
        omni.kit.clipboard.copy(output)


class ExtsListView:
    def __init__(self, model, ext_id):
        self._model = model
        self._ext_id = ext_id

        # gather reverse dependencies from both local extensions and registry extensions
        registry_exts = []
        unique_registry_exts = set()
        manager = omni.kit.app.get_app_interface().get_extension_manager()
        for ext in manager.get_extensions():
            ext_name = ext["name"]
            registry_exts.append(ext)
            unique_registry_exts.add(ext_name)
        for ext in manager.get_registry_extensions():
            ext_id = ext["id"]
            ext_name = ext["name"]
            if ext_name not in unique_registry_exts:
                # grab latest version of each extension
                versions = manager.fetch_extension_versions(ext_name)
                if len(versions) > 0 and versions[0]["id"] == ext_id:
                    registry_exts.append(ext)
                    unique_registry_exts.add(ext_name)
        self._reverse_dep = self._get_reverse_dependencies(manager, registry_exts)

        with ui.ScrollingFrame(
            horizontal_scrollbar_policy=ui.ScrollBarPolicy.SCROLLBAR_ALWAYS_OFF,
            vertical_scrollbar_policy=ui.ScrollBarPolicy.SCROLLBAR_ALWAYS_ON,
        ):
            with ui.ZStack():
                ui.Rectangle(style_type_name_override="ExtensionDescription.ContentBackground")
                with ui.HStack():
                    ui.Spacer(width=10)
                    with ui.VStack(height=0):
                        ui.Spacer(height=10)
                        direct = self._model._nodes[self._ext_id].deps if self._ext_id else ()
                        ui.Label(f"- Direct Dependencies ({len(direct)}):", height=20)
                        for ext in sorted(direct):
                            ui.Label(ext)
                        ui.Spacer(height=10)
                        indirect = [
                            node
                            for node in self._model._nodes.values()
                            if node.id not in direct and node.id != self._ext_id
                        ]
                        ui.Label(f"- All Indirect Dependencies ({len(indirect)}):", height=20)
                        for node in sorted(indirect):
                            ui.Label(node.id)
                    ui.Spacer(width=10)
                    with ui.VStack(height=10):
                        ui.Spacer(height=10)
                        direct = self._reverse_dep[0]
                        ui.Label(f"- Direct Reverse Dependencies ({len(direct)}):", height=20)
                        for item in sorted(direct):
                            ui.Label(item)
                        ui.Spacer(height=10)
                        indirect = [
                            ext
                            for ext in self._reverse_dep[1]
                            if ext not in direct and ext != ext_id_to_fullname(self._ext_id)
                        ]
                        ui.Label(f"- All Indirect Reverse Dependencies ({len(indirect)}):", height=20)
                        for item in sorted(indirect):
                            ui.Label(item)

    def _get_reverse_dependencies(self, manager, exts) -> Tuple[List, List]:
        dependents = defaultdict(set)
        unique_exts_first_order = set()
        unique_exts = set()
        max_depth = 40

        for ext in exts:
            ext_id = ext["id"]
            ext_name = ext["name"]
            info = manager.get_extension_dict(ext_id)
            if not info:
                info = manager.get_registry_extension_dict(ext_id)
            if info:
                deps = info.get("dependencies", [])
                for dep_name in deps:
                    dependents[dep_name].add(ext_name)

        def recurse(ext_name: str, cur_depth: int):
            if cur_depth < max_depth:
                if cur_depth == 1:
                    unique_exts_first_order.add(ext_name)
                unique_exts.add(ext_name)
                for dep_name in dependents[ext_name]:
                    recurse(dep_name, cur_depth + 1)

        if self._ext_id:
            info = manager.get_extension_dict(self._ext_id)
            ext_name = info.get("package/name", "")
            recurse(ext_name, 0)

        # returns a tuple, 1st order deps only and all reverse deps
        return list(unique_exts_first_order), list(unique_exts)


class ExtsGraphWidget:
    """Extensions graph window"""

    def __init__(self, ext_id):
        self._ext_id = ext_id
        self._view_index = 0
        with ui.ZStack():
            self._delegate = GraphNodeDelegate()
            self._model = Model(self._ext_id)
            # TODO(anov): remove pan_x, pan_y hardcoded after graph auto align fixed.

            self._graph_frame = ui.Frame()
            raster_nodes = carb.settings.get_settings().get("/exts/omni.kit.window.extensions/raster_nodes")
            with self._graph_frame:
                self._graph_view = GraphView(
                    delegate=self._delegate,
                    model=self._model,
                    style=GRAPH_STYLE,
                    zoom=0.5,
                    pan_x=1000,
                    pan_y=400,
                    raster_nodes=raster_nodes,
                )

            self._list_frame = ui.Frame()
            with self._list_frame:
                self._list_view = ExtsListView(self._model, self._ext_id)

            self._list_frame.visible = False

            with ui.VStack():
                ui.Spacer()
                with ui.HStack(height=0):
                    ui.Spacer()
                    with ui.VStack(width=200, height=0, content_clipping=True):
                        ui.Button(
                            "Toggle View",
                            width=200,
                            height=20,
                            mouse_pressed_fn=lambda *_: self._toggle_view(),
                        )
                        ui.Button(
                            "Copy as Graphviz .dot",
                            width=200,
                            height=20,
                            mouse_pressed_fn=lambda *_: self._model.copy_as_graphviz(),
                        )
                    ui.Spacer(width=10)
                ui.Spacer(height=10)

    def _toggle_view(self):
        self._view_index = (self._view_index + 1) % 2
        self._list_frame.visible = self._view_index == 1
        self._graph_frame.visible = self._view_index == 0

    def destroy(self):
        self._delegate = None
        self._model = None
        self._list_frame = None
        self._list_view = None
        self._graph_frame = None
        self._graph_view = None


class ExtsGraphWindow:
    """Extensions graph window"""

    def __init__(self):
        self._frame = ui.Frame()
        self._frame.visible = False
        self._graph_widget = None
        self.select_ext(None)

        manager = omni.kit.app.get_app_interface().get_extension_manager()

        change_stream = manager.get_change_event_stream()
        self._change_script_sub = change_stream.create_subscription_to_pop(
            lambda _: self._refresh(), name="ExtsGraphWindow watch for exts"
        )

    def destroy(self):
        self._frame = None
        if self._graph_widget:
            self._graph_widget.destroy()
        self._change_script_sub = None

    def _refresh(self):
        self.set_visible(self._frame.visible)

    def select_ext(self, ext_id):
        self._ext_id = ext_id

    def set_visible(self, visible: bool):
        if self._frame.visible == visible:
            if visible:
                # recreate if already shown
                self.set_visible(False)
            else:
                # if not shown already -> do nothing.
                return

        if not visible:
            self._frame.visible = False
        else:
            self._build()
            self._frame.visible = True

    def _build(self):
        with self._frame:
            if self._graph_widget:
                self._graph_widget.destroy()
            self._graph_widget = ExtsGraphWidget(self._ext_id)
