import pymoku
from pymoku.applets import doasync
from pymoku.sources import RoutingException

from NodeGraphQt import (NodeGraph, BaseNode)
from NodeGraphQt.base.factory import NodeFactory

from NodeGraphQt.qgraphics.node_port_in import PortInputNodeItem

from PySide6 import QtGui
from PySide6.QtCore import Signal
from PySide6.QtWidgets import QMessageBox


COLOR2X = (128, 128, 0)
COLOR4X = (255, 0, 0)


class MokuInputNode(BaseNode):
    def __init__(self, source):
        super().__init__(PortInputNodeItem)
        self.source = source
        self.NODE_NAME = source.name

        color = None
        if source.bus_width == 2:
            color = COLOR2X
        elif source.bus_width == 4:
            color = COLOR4X

        self.add_output(multi_output=True, display_name=False, color=color)


class MokuOutputNode(BaseNode):
    def __init__(self, inp):
        super().__init__(PortInputNodeItem)
        self.inp = inp
        self.NODE_NAME = inp.name

        color = None
        if inp.bus_width == 2:
            color = COLOR2X
        elif inp.bus_width == 4:
            color = COLOR4X

        self.add_input(multi_input=False, display_name=False, color=color)


class SlotNode(BaseNode):
    def __init__(self):
        super().__init__()
        self.slot_inst = None

    def set_slot(self, slot_inst):
        self.slot_inst = slot_inst

        slot_inst.attach()  # TODO
        # TODO we assume the index of the input/output matches the index of the port
        for inp in slot_inst.inputs():
            color = None
            if inp.bus_width == 2:
                color = COLOR2X
            elif inp.bus_width == 4:
                color = COLOR4X
            self.add_input(inp.name, color=color)

        for outp in slot_inst.outputs():
            color = None
            if outp.bus_width == 2:
                color = COLOR2X
            elif outp.bus_width == 4:
                color = COLOR4X
            self.add_output(outp.name, color=color)

    def get_source_from_port(self, port):
        idx = self._outputs.index(port)
        return self.slot_inst.outputs()[idx]

    def get_input_from_port(self, port):
        idx = self._inputs.index(port)
        return self.slot_inst.inputs()[idx]


class PluginNodeFactory(NodeFactory):
    def __init__(self):
        self.moku = None

    def set_moku(self, moku):
        self.moku = moku

    def register_node(self, node, alias=None):
        raise NotImplementedError()

    @property
    def names(self):
        if self.moku is None:
            return {}

        names = {}
        func_ids = self.moku.function_ids_for_platform(self.moku.platform_sha)

        for f in func_ids:
            name = self._node_name(f)
            if name in names:
                names[name].append(f)
            else:
                names[name] = [f]

        return names

    @property
    def aliases(self):
        return {}

    @property
    def nodes(self):
        if self.moku is None:
            return {}
        func_ids = self.moku.function_ids_for_platform(self.moku.platform_sha)
        return {_id: SlotNode for _id in func_ids}

    def _node_name(self, node_type):
        if (plugin := self.moku.plugin_for_function(node_type)) is None:
            return node_type
        try:
            return plugin.get_name(node_type)
        except Exception:
            return node_type

    def create_node_instance(self, node_type=None):
        node_type = node_type.replace('.', ':')
        inst = SlotNode()
        inst.NODE_NAME = self._node_name(node_type)
        inst._id = node_type
        return inst


class MokuNodeGraph(NodeGraph):
    slot_removed = Signal(pymoku.Slot)

    def __init__(self):
        super().__init__(node_factory=PluginNodeFactory())
        self.moku = None

        self.set_acyclic(False)

        self.port_connected.connect(self._port_connected)
        self.port_disconnected.connect(self._port_disconnected)
        self.property_changed.connect(print)

        self.context_menu().add_command('Delete', self.delete_key, shortcut=QtGui.QKeySequence.Backspace)

        self.auto_layout_nodes()
        self.clear_selection()
        self.fit_to_selection()

    def set_moku(self, moku):
        self.moku = None  # don't want to clear all the slots
        self.delete_nodes(self.all_nodes())
        self.moku = moku

        if self.moku is None or self.moku.num_slots == 0:
            return

        self.node_factory.set_moku(moku)

        with self.moku.cached_state():
            y = 0.0
            for src in moku.platform_sources():
                self.add_node(MokuInputNode(src), pos=[0.0, y], selected=False, push_undo=False)
                y += 80.0

            y = 0.0
            for inp in moku.platform_outputs():
                self.add_node(MokuOutputNode(inp), pos=[800, y], selected=False, push_undo=False)
                y += 80.0

            slot_nodes = []
            y = 0.0
            for slot_inst in self.moku._slots:
                if slot_inst is not None and slot_inst.ID != 'pymoku:empty':
                    node = self.node_factory.create_node_instance(slot_inst.ID)
                    self.add_node(node, pos=[400, y], selected=False, push_undo=False)
                    # some connections to other insturment nodes won't be possible
                    # but later nodes will complete them
                    node.set_slot(slot_inst)
                    slot_nodes.append(node)
                    # TODO maybe persist/relaod the current layout?
                    y += 100.0 + len(slot_inst.inputs()) * 20.0

            for node in slot_nodes:
                # defer so all nodes exist
                self.connect_slot_nodes(node)

        self.clear_selection()
        self.fit_to_selection()

    def _to_input_source(self, in_port, out_port):
        if isinstance(in_port.node(), MokuOutputNode):
            inp = in_port.node().inp
            src = out_port.node().get_source_from_port(out_port)
        elif isinstance(out_port.node(), MokuInputNode):
            src = out_port.node().source
            inp = in_port.node().get_input_from_port(in_port)
        else:  # Instr to Instr
            inp = in_port.node().get_input_from_port(in_port)
            src = out_port.node().get_source_from_port(out_port)
        return inp, src

    def _port_connected(self, in_port, out_port):
        inp, src = self._to_input_source(in_port, out_port)

        @doasync
        def a(inp, src):
            inp.set_source(src)
        a.error.connect(lambda exc, i=in_port, o=out_port: self._route_error(exc, i, o))
        a(inp, src)

    def _route_error(self, exc, in_port, out_port):
        in_port.disconnect_from(out_port, push_undo=False, emit_signal=False)
        msg = f'{str(exc)}'
        if isinstance(exc, RoutingException) and exc.conflicts:
            msg += '<br><br>Conflicts:'
            for conf in exc.conflicts:
                msg += f'<br>{conf.name}'
        QMessageBox.critical(None, 'Routing Conflict', msg)

    def _port_disconnected(self, in_port, out_port):
        inp, src = self._to_input_source(in_port, out_port)
        inp.set_source(None)

    def _register_builtin_nodes(self):
        pass

    def get_node_for_slot(self, slot_inst):
        """
        Returns the SlotNode for the given slot.
        """
        for node in self.all_nodes():
            if isinstance(node, SlotNode) and node.slot_inst is slot_inst:
                return node
        return None

    def get_port_for_source(self, src):
        # try for platform sources first
        node = self.get_node_by_name(src.name)
        if node is not None:
            return node.output(0)

        # otherwise it must be an slot source
        node = self.get_node_for_slot(src.slot())
        if node is not None:
            if src not in src.slot_inst.outputs():
                return None

            return node.output(src.slot_inst.outputs().index(src))

        # the slot node might not exist yet
        return None

    def connect_slot_nodes(self, slot_node):
        with slot_node.slot_inst.cached_registers(), self.moku.cached_state():
            for port, inp in zip(slot_node.input_ports(), slot_node.slot_inst.inputs()):
                src = inp.get_attached()
                if src is None:
                    continue
                src_port = self.get_port_for_source(src)
                if src_port is None:
                    continue
                src_port.connect_to(port, push_undo=False, emit_signal=False)

            for node in self.get_nodes_by_type('OdenGraphQt.nodes.MokuOutputNode'):
                src = node.inp.get_attached()

                if hasattr(src, 'slot') and src.slot() is slot_node.slot_inst:
                    src_port = self.get_port_for_source(src)
                    if src_port is not None:
                        src_port.connect_to(node.get_input(0), push_undo=False, emit_signal=False)

    def delete_key(self, graph):
        nodes = [node for node in self.selected_nodes() if isinstance(node, SlotNode)]
        self.delete_nodes(nodes)

    def delete_nodes(self, nodes):
        @doasync
        def clear(slot):
            if self.moku is not None:
                self.moku.clear_slot(slot)

        for node in nodes:
            if isinstance(node, SlotNode) and node.slot_inst is not None:
                clear(node.slot_inst.slot)
                self.slot_removed.emit(node.slot_inst)

        super().delete_nodes(nodes)
