import sys

from PyQt5.QtWidgets import QApplication, QGraphicsView, QGraphicsScene, QPushButton, QGraphicsProxyWidget
from PyQt5.QtGui import QPen
from PyQt5.QtCore import Qt

from Desktop.data.connect import connect, Department, init_db


class DepartmentGraph(QGraphicsView):
    def __init__(self, parent=None):
        super().__init__(parent)
        self.scene = QGraphicsScene(self)
        self.setScene(self.scene)
        with connect() as session:
            self.departments = session.query(Department).all()  # Список департаментов с их иерархией
        self.nodes = {}
        self.render_graph()

    def render_graph(self):
        root_departments = [d for d in self.departments if d.parent_id is None]
        x, y = 50, 50  # Начальные координаты
        for root in root_departments:
            self.add_department_node(root, x, y)
            x += 200  # Расстояние между корневыми узлами

    def add_department_node(self, department, x, y, parent_item=None):
        # Создаем кнопку
        button = QPushButton(department.title)
        proxy = QGraphicsProxyWidget()
        proxy.setWidget(button)
        proxy.setPos(x, y)
        self.scene.addItem(proxy)

        # Сохраняем узел
        self.nodes[department.id] = proxy

        # Если есть родитель, добавляем линию
        if parent_item:
            parent_center = parent_item.sceneBoundingRect().center()
            current_center = proxy.sceneBoundingRect().center()
            line = self.scene.addLine(parent_center.x(), parent_center.y(), current_center.x(), current_center.y(),
                                      QPen(Qt.black))

        # Рекурсивно добавляем подразделы
        child_departments = [d for d in self.departments if d.parent_id == department.id]
        child_x = x - len(child_departments) * 100 // 2
        for child in child_departments:
            self.add_department_node(child, child_x, y + 100, proxy)
            child_x += 200


if __name__ == "__main__":
    init_db()

    app = QApplication(sys.argv)
    window = DepartmentGraph()
    window.setWindowTitle("Иерархия департаментов")
    window.resize(800, 600)
    window.show()

    sys.exit(app.exec_())