import sys
from PySide6.QtWidgets import (
    QMainWindow, QWidget, QVBoxLayout, QTextEdit,
    QPushButton, QLabel, QSplitter
)
from PySide6.QtCore import Qt
from PySide6 import QtGui
from NodeGraphQt import NodeGraph, BaseNode

from core.workflow import Graph
from nodes.image_reader import ImageReader
from nodes.feature_extractor import FeatureExtractor
from nodes.feature_writer import FeatureWriter


# GUI node wrappers around core nodes

class ImageReaderNode(BaseNode):
    __identifier__ = "radiuma.gui"
    NODE_NAME = "ImageReader"

    def __init__(self):
        super().__init__()
        self.add_output("image")
        self.core_node = ImageReader()


class FeatureExtractorNode(BaseNode):
    __identifier__ = "radiuma.gui"
    NODE_NAME = "FeatureExtractor"

    def __init__(self):
        super().__init__()
        self.add_input("image")
        self.add_output("features")
        self.core_node = FeatureExtractor()


class FeatureWriterNode(BaseNode):
    __identifier__ = "radiuma.gui"
    NODE_NAME = "FeatureWriter"

    def __init__(self):
        super().__init__()
        self.add_input("features")
        self.core_node = FeatureWriter()


class VisualEditor(QMainWindow):
    # Visual workflow editor with dynamic semantic validation

    def __init__(self):
        super().__init__()
        self.setWindowTitle("Radiuma R&D – Semantic Workflow Editor")
        self.resize(1200, 800)

        self.graph_engine = Graph()
        self.graph = NodeGraph()
        self.viewer = self.graph.widget

        self.graph.register_node(ImageReaderNode)
        self.graph.register_node(FeatureExtractorNode)
        self.graph.register_node(FeatureWriterNode)

        self.log_panel = QTextEdit()
        self.log_panel.setReadOnly(True)

        btn_reader = QPushButton("Add ImageReader")
        btn_extractor = QPushButton("Add FeatureExtractor")
        btn_writer = QPushButton("Add FeatureWriter")

        btn_reader.clicked.connect(
            lambda: self.graph.create_node("radiuma.gui.ImageReaderNode")
        )
        btn_extractor.clicked.connect(
            lambda: self.graph.create_node("radiuma.gui.FeatureExtractorNode")
        )
        btn_writer.clicked.connect(
            lambda: self.graph.create_node("radiuma.gui.FeatureWriterNode")
        )

        left_panel = QWidget()
        left_layout = QVBoxLayout(left_panel)
        left_layout.addWidget(QLabel("Nodes"))
        left_layout.addWidget(btn_reader)
        left_layout.addWidget(btn_extractor)
        left_layout.addWidget(btn_writer)
        left_layout.addStretch()

        top_splitter = QSplitter(Qt.Horizontal)
        top_splitter.addWidget(left_panel)
        top_splitter.addWidget(self.viewer)

        main_splitter = QSplitter(Qt.Vertical)
        main_splitter.addWidget(top_splitter)
        main_splitter.addWidget(self.log_panel)
        main_splitter.setSizes([600, 200])

        self.setCentralWidget(main_splitter)

        self.graph.port_connected.connect(self._on_port_connected)
        self.graph.port_disconnected.connect(self._on_port_disconnected)

    def log(self, msg, level="INFO"):
        self.log_panel.append(f"[{level}] {msg}")

    def _on_port_connected(self, port_a, port_b):
        node_a = port_a.node()
        node_b = port_b.node()

        core_a = getattr(node_a, "core_node", None)
        core_b = getattr(node_b, "core_node", None)

        if not core_a or not core_b:
            self._color_connection(port_a, port_b, (200, 200, 200))
            self.log("Missing core node mapping", "ERROR")
            return

        # Normalize out → in
        if port_a.type_() == "out":
            out_node, out_port = core_a, port_a.name()
            in_node, in_port = core_b, port_b.name()
            gui_out, gui_in = port_a, port_b
        else:
            out_node, out_port = core_b, port_b.name()
            in_node, in_port = core_a, port_a.name()
            gui_out, gui_in = port_b, port_a

        ok, msg = self.graph_engine.can_connect(
            out_node, out_port, in_node, in_port
        )

        if ok:
            self._color_connection(gui_out, gui_in, (46, 204, 113))
            self.log(msg, "INFO")
        else:
            self._color_connection(gui_out, gui_in, (231, 76, 60))
            self.log(msg, "ERROR")

    def _on_port_disconnected(self, port_a, port_b):
        self._color_connection(port_a, port_b, (180, 180, 180))
        self.log("Disconnected ports", "INFO")

    def _color_connection(self, port_a, port_b, color):
        # Patch connection paint dynamically
        scene = self.graph._viewer.scene()
        for item in scene.items():
            if hasattr(item, "port1") and hasattr(item, "port2"):
                if {item.port1, item.port2} == {port_a, port_b}:
                    setattr(item, "_custom_color", color)
                    if not hasattr(item, "_patched"):
                        original_paint = item.paint

                        def patched_paint(painter, option, widget=None,
                                          _orig=original_paint, _item=item):
                            pen = painter.pen()
                            pen.setColor(QtGui.QColor(*_item._custom_color))
                            painter.setPen(pen)
                            _orig(painter, option, widget)

                        item.paint = patched_paint
                        item._patched = True
                    item.update()
                    return