from NodeGraphQt import NodesTreeWidget
from NodeGraphQt.custom_widgets.nodes_tree import (
    _BaseNodeTreeItem, TYPE_NODE, TYPE_CATEGORY, URN_SCHEME, MIME_TYPE)
from PySide6 import QtCore


class SlotFunctionsTreeWidget(NodesTreeWidget):
    def _build_tree(self):
        self.clear()
        categories = set()
        node_types = {}
        for name, node_ids in self._factory.names.items():
            for nid in node_ids:
                if nid == 'pymoku:empty':
                    continue
                categories.add(':'.join(nid.split(':')[:-1]))
                node_types[nid] = name

        self._category_items = {}
        for category in sorted(categories):
            if category in self._custom_labels.keys():
                label = self._custom_labels[category]
            else:
                label = '{}'.format(category)
            cat_item = _BaseNodeTreeItem(self, [label], type=TYPE_CATEGORY)
            cat_item.setFirstColumnSpanned(True)
            cat_item.setFlags(QtCore.Qt.ItemIsEnabled)
            cat_item.setSizeHint(0, QtCore.QSize(100, 26))
            self.addTopLevelItem(cat_item)
            cat_item.setExpanded(True)
            self._category_items[category] = cat_item

        for node_id, node_name in node_types.items():
            category = ':'.join(node_id.split(':')[:-1])
            category_item = self._category_items[category]

            item = _BaseNodeTreeItem(category_item, [node_name], type=TYPE_NODE)
            item.setToolTip(0, node_id)
            item.setSizeHint(0, QtCore.QSize(100, 26))

            category_item.addChild(item)

    def mimeData(self, items):
        node_ids = ["node:{}".format(i.toolTip(0).replace(':', '.')) for i in items]
        node_urn = URN_SCHEME + ';'.join(node_ids)
        mime_data = QtCore.QMimeData()
        mime_data.setData(MIME_TYPE, QtCore.QByteArray(node_urn.encode()))
        return mime_data
