import time
from PySide6.QtWidgets import (
    QWidget, QGraphicsView, QGraphicsScene, QListWidget, QListWidgetItem,
    QPushButton, QHBoxLayout, QVBoxLayout, QDialog, QTextEdit
)
from PySide6.QtGui import (
    QPainter, QPen, QColor, QPainterPath, QTransform, QMouseEvent
)
from PySide6.QtCore import Qt, QRectF
from PySide6.QtWidgets import QGraphicsItem

from api.execution_context import ExecutionContext
from api.workflow import Workflow
from api.scheduler import Scheduler


# PORT ITEM
class PortItem(QGraphicsItem):
    R = 6

    def __init__(self, parent_node, name, is_output, index):
        super().__init__(parent_node)
        self.parent_node = parent_node
        self.name = name
        self.is_output = is_output
        self.index = index
        self.setFlag(QGraphicsItem.ItemIsSelectable)

    def boundingRect(self):
        return QRectF(-self.R, -self.R, 2*self.R, 2*self.R)

    def paint(self, painter, option, widget=None):
        painter.setRenderHint(QPainter.Antialiasing)
        color = QColor(34,139,34) if self.is_output else QColor(70,130,180)
        painter.setPen(QPen(Qt.black, 1))
        painter.setBrush(color)
        painter.drawEllipse(self.boundingRect())


# CONNECTION ITEM
class ConnectionItem(QGraphicsItem):
    def __init__(self, out_port, in_port, editor):
        super().__init__()
        self.out_port = out_port
        self.in_port = in_port
        self.editor = editor
        self.setZValue(-1)
        self._rect = QRectF()
        self._path = QPainterPath()
        self.setFlag(QGraphicsItem.ItemIsSelectable)

    def update_path(self):
        p1 = self.out_port.scenePos()
        p2 = self.in_port.scenePos()

        new_rect = QRectF(p1, p2).normalized().adjusted(-20, -20, 20, 20)
        if new_rect != self._rect:
            self.prepareGeometryChange()
            self._rect = new_rect

        path = QPainterPath(p1)
        dx = (p2.x() - p1.x()) * 0.5
        path.cubicTo(p1.x() + dx, p1.y(), p2.x() - dx, p2.y(), p2.x(), p2.y())
        self._path = path

    def boundingRect(self):
        return self._rect

    def paint(self, painter, option, widget=None):
        self.update_path()
        painter.setRenderHint(QPainter.Antialiasing)

        if self.isSelected():
            painter.setPen(QPen(QColor(255, 0, 0), 3))
        else:
            painter.setPen(QPen(QColor(50, 50, 50), 2))

        painter.drawPath(self._path)

    def mousePressEvent(self, event):
        if event.button() == Qt.LeftButton:
            self.editor._remove_connection(self)
        super().mousePressEvent(event)


# NODE ITEM
class NodeItem(QGraphicsItem):
    W = 200
    H = 80

    def __init__(self, module, title, editor):
        super().__init__()
        self.module = module
        self.title = title
        self.editor = editor
        self.running = False

        self.in_ports = []
        self.out_ports = []

        self.setFlags(QGraphicsItem.ItemIsMovable | QGraphicsItem.ItemIsSelectable)
        self._build_ports()

    def highlight(self, active):
        self.running = active
        self.update()

    def _build_ports(self):
        y = 30
        for i, p in enumerate(self.module.in_ports):
            port = PortItem(self, p.name, False, i)
            port.setPos(0, y)
            self.in_ports.append(port)
            y += 20

        y = 30
        for i, p in enumerate(self.module.out_ports):
            port = PortItem(self, p.name, True, i)
            port.setPos(self.W, y)
            self.out_ports.append(port)
            y += 20

    def boundingRect(self):
        return QRectF(0, 0, self.W, self.H)

    def paint(self, painter, option, widget=None):
        painter.setRenderHint(QPainter.Antialiasing)

        if self.running:
            painter.setBrush(QColor("#c8f7c5"))
        else:
            painter.setBrush(QColor("#f0f0f0"))

        painter.setPen(QPen(Qt.black, 1))
        painter.drawRect(self.boundingRect())
        painter.drawText(10, 20, self.title)

    def itemChange(self, change, value):
        if change == QGraphicsItem.ItemPositionHasChanged:
            self.editor._update_connections_for_node(self)
        return super().itemChange(change, value)


# CUSTOM VIEW
class NodeGraphicsView(QGraphicsView):
    def __init__(self, scene, editor):
        super().__init__(scene)
        self.editor = editor
        self.setRenderHint(QPainter.Antialiasing)

    def mousePressEvent(self, event: QMouseEvent):
        pos = self.mapToScene(event.pos())
        item = self.scene().itemAt(pos, QTransform())

        port = self.editor._find_port(item)
        if port:
            if port.is_output:
                self.editor.dragging_port = port
                return
            else:
                if self.editor.dragging_port is not None:
                    self.editor._connect(self.editor.dragging_port, port)
                    self.editor.dragging_port = None
                    return

        super().mousePressEvent(event)


# NODE EDITOR
class NodeEditor(QWidget):
    def __init__(self, module_catalog):
        super().__init__()
        self.setWindowTitle("Node Editor")

        self.module_catalog = module_catalog

        self.scene = QGraphicsScene()
        self.view = NodeGraphicsView(self.scene, self)

        self.left = QListWidget()
        for name in module_catalog:
            QListWidgetItem(name, self.left)

        self.log = QListWidget()

        self.run_btn = QPushButton("Run")
        self.stop_btn = QPushButton("Stop")
        self.resume_btn = QPushButton("Resume")
        self.cancel_btn = QPushButton("Cancel")
        self.show_dag_btn = QPushButton("Show DAG")

        self.add_btn = QPushButton("Add Node")

        # Layout
        main_layout = QHBoxLayout(self)
        left_layout = QVBoxLayout()
        left_layout.addWidget(self.left)
        left_layout.addWidget(self.add_btn)
        left_layout.addWidget(self.run_btn)
        left_layout.addWidget(self.stop_btn)
        left_layout.addWidget(self.resume_btn)
        left_layout.addWidget(self.cancel_btn)
        left_layout.addWidget(self.show_dag_btn)
        left_layout.addWidget(self.log)

        main_layout.addLayout(left_layout)
        main_layout.addWidget(self.view)

        self.left.setMaximumWidth(220)
        main_layout.setStretch(1, 1)

        # Events
        self.add_btn.clicked.connect(self.add_node)
        self.run_btn.clicked.connect(self.run_workflow)
        self.stop_btn.clicked.connect(self.stop_workflow)
        self.resume_btn.clicked.connect(self.resume_workflow)
        self.cancel_btn.clicked.connect(self.cancel_workflow)
        self.show_dag_btn.clicked.connect(self.show_dag)

        self.nodes = []
        self.connections = []
        self.dragging_port = None

        self.scheduler = None
        self.ctx = None
        self.execution_order = []
        self.current_index = 0
        self.workflow_running = False
        self.workflow_paused = False

    # LOGGING
    def _log(self, text, color="black"):
        item = QListWidgetItem(text)
        item.setForeground(QColor(color))
        self.log.addItem(item)
        self.log.scrollToBottom()

    # ADD NODE
    def add_node(self):
        sel = self.left.currentItem()
        if not sel:
            return

        name = sel.text()
        module = self.module_catalog[name]()

        node = NodeItem(module, name, self)
        node.setPos(len(self.nodes)*40, len(self.nodes)*40)
        self.scene.addItem(node)

        self.nodes.append(node)
        self._log(f"Node added: {name}", "blue")

    # FIND PORT
    def _find_port(self, item):
        while item:
            if isinstance(item, PortItem):
                return item
            item = item.parentItem()
        return None

    # REMOVE CONNECTION
    def _remove_connection(self, conn):
        if conn in self.connections:
            self.scene.removeItem(conn)
            self.connections.remove(conn)
            self._log("Connection removed", "red")

    # CONNECT WITH TYPE CHECK
    def _connect(self, out_port_item, in_port_item):
        out_port = out_port_item.parent_node.module.out_ports[out_port_item.index]
        in_port = in_port_item.parent_node.module.in_ports[in_port_item.index]

        # Only output → input
        if not out_port_item.is_output or in_port_item.is_output:
            self._log("❌ Invalid direction", "red")
            return

        # Only one connection per input
        for c in self.connections:
            if c.in_port is in_port_item:
                self._log("❌ Input already connected", "red")
                return

        # Type checking
        try:
            out_port.connect(in_port)
        except Exception as e:
            self._log(f"❌ Type mismatch: {e}", "red")
            return

        # Success
        conn = ConnectionItem(out_port_item, in_port_item, self)
        self.scene.addItem(conn)
        self.connections.append(conn)

        self._log(
            f"✔ Connected: {out_port_item.parent_node.title}.{out_port.name} → "
            f"{in_port_item.parent_node.title}.{in_port.name}",
            "green"
        )

    # UPDATE CONNECTIONS ON MOVE
    def _update_connections_for_node(self, node):
        for c in self.connections:
            if c.out_port.parent_node is node or c.in_port.parent_node is node:
                c.update_path()

    # BUILD GRAPH
    def _build_graph(self):
        graph = {i: [] for i in range(len(self.nodes))}

        for c in self.connections:
            out_idx = self.nodes.index(c.out_port.parent_node)
            in_idx = self.nodes.index(c.in_port.parent_node)
            graph[out_idx].append(in_idx)

        return graph

    # TOPOLOGICAL SORT
    def _topological_sort(self, graph):
        indeg = {k: 0 for k in graph}
        for u in graph:
            for v in graph[u]:
                indeg[v] += 1

        q = [u for u in graph if indeg[u] == 0]
        order = []

        while q:
            u = q.pop(0)
            order.append(u)
            for v in graph[u]:
                indeg[v] -= 1
                if indeg[v] == 0:
                    q.append(v)

        if len(order) != len(graph):
            return None

        return order

    # SHOW DAG
    def show_dag(self):
        graph = self._build_graph()

        text = ""
        for src, dsts in graph.items():
            src_name = self.nodes[src].title
            if dsts:
                for d in dsts:
                    text += f"{src_name} → {self.nodes[d].title}\n"
            else:
                text += f"{src_name} → (no outputs)\n"

        dlg = QDialog(self)
        dlg.setWindowTitle("DAG Structure")
        layout = QVBoxLayout(dlg)
        txt = QTextEdit()
        txt.setReadOnly(True)
        txt.setText(text)
        layout.addWidget(txt)
        dlg.resize(400, 300)
        dlg.exec()

    # WORKFLOW EXECUTION
    def run_workflow(self):
        self._log("Workflow started", "green")

        self.ctx = ExecutionContext("gui_run")
        self.scheduler = Scheduler()

        # Build workflow
        wf = Workflow("gui_workflow")
        for n in self.nodes:
            wf.add_child(n.module)

        # Build graph
        graph = self._build_graph()
        order = self._topological_sort(graph)

        if order is None:
            self._log("❌ Cycle detected in graph", "red")
            return

        self.execution_order = order
        self.current_index = 0
        self.workflow_running = True
        self.workflow_paused = False

        self._execute_next()

    def stop_workflow(self):
        self.workflow_paused = True
        self._log("Paused", "orange")

    def resume_workflow(self):
        self.workflow_paused = False
        self._log("Resumed", "green")
        self._execute_next()

    def cancel_workflow(self):
        self.workflow_running = False
        self._log("Cancelled", "red")

    # EXECUTE NEXT NODE
    def _execute_next(self):
        if not self.workflow_running or self.workflow_paused:
            return

        if self.current_index >= len(self.execution_order):
            self._log("Workflow finished", "green")
            self.workflow_running = False
            return

        idx = self.execution_order[self.current_index]
        node = self.nodes[idx]
        module = node.module

        node.highlight(True)
        self._log(f"Running: {node.title}", "blue")

        start = time.time()

        try:
            self.scheduler.run(module, self.ctx)
            end = time.time()
            duration = round(end - start, 4)

            self._log(f"Finished: {node.title} (time: {duration}s)", "green")
            self._show_node_output(module)

        except Exception as e:
            self._log(f"❌ Error in {node.title}: {e}", "red")

        node.highlight(False)

        self.current_index += 1
        self._execute_next()

    # SHOW NODE OUTPUT (SUMMARY)
    def _show_node_output(self, module):
        self._log("Output:", "purple")

        for out_port in module.out_ports:
            try:
                asset = self.ctx.get(out_port)
                data = asset.data

                summary = self._summarize(data)
                self._log(f"  {out_port.name}: {summary}", "black")

            except Exception as e:
                self._log(f"  {out_port.name}: <error reading output>", "red")

    def _summarize(self, data):
        text = str(data)
        if len(text) > 200:
            return text[:200] + " ... (truncated)"
        return text
