#!/usr/bin/env python3

# https://stackoverflow.com/questions/26227885/drag-and-drop-rows-within-qtablewidget

import sys

from PyQt6.QtWidgets import (
    QTableWidget,
    QAbstractItemView,
    QTableWidgetItem,
    QWidget,
    QHBoxLayout,
    QApplication,
)
from PyQt6.QtCore import Qt


class TableWidgetDragRows(QTableWidget):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

        self.setDragEnabled(True)
        self.setAcceptDrops(True)
        self.viewport().setAcceptDrops(True)
        self.setDragDropOverwriteMode(False)
        self.setDropIndicatorShown(True)

        self.setSelectionMode(QAbstractItemView.SelectionMode.ExtendedSelection)
        self.setSelectionBehavior(QAbstractItemView.SelectionBehavior.SelectRows)
        self.setDragDropMode(QAbstractItemView.DragDropMode.InternalMove)
        self.DropIndicatorPosition(QAbstractItemView.DropIndicatorPosition.BelowItem)

    def dropEvent(self, event):
        if event.source() == self:
            rows = set([mi.row() for mi in self.selectedIndexes()])
            targetRow = self.indexAt(event.position().toPoint()).row()
            if self.dropIndicatorPosition() == QAbstractItemView.DropIndicatorPosition.BelowItem:
                targetRow += 1
            rows.discard(targetRow)
            rows = sorted(rows)
            if not rows:
                return
            if targetRow == -1:
                targetRow = self.rowCount()
            for _ in range(len(rows)):
                self.insertRow(targetRow)
            rowMapping = dict()  # Src row to target row.
            for idx, row in enumerate(rows):
                if row < targetRow:
                    rowMapping[row] = targetRow + idx
                else:
                    rowMapping[row + len(rows)] = targetRow + idx
            colCount = self.columnCount()
            for srcRow, tgtRow in sorted(rowMapping.items()):
                for col in range(0, colCount):
                    self.setItem(tgtRow, col, self.takeItem(srcRow, col))
            for row in reversed(sorted(rowMapping.keys())):
                self.removeRow(row)
            event.accept()
            return


class Window(QWidget):
    def __init__(self):
        super(Window, self).__init__()

        layout = QHBoxLayout()
        self.setLayout(layout)

        self.table_widget = TableWidgetDragRows()
        layout.addWidget(self.table_widget)

        # setup table widget
        self.table_widget.setColumnCount(2)
        self.table_widget.setHorizontalHeaderLabels(["Type", "Name"])

        items = [
            ("Red", "Toyota"),
            ("Blue", "RV"),
            ("Green", "Beetle"),
            ("Silver", "Chevy"),
            ("Black", "BMW"),
        ]
        self.table_widget.setRowCount(len(items))
        for i, (color, model) in enumerate(items):
            item_flags = Qt.ItemFlag.ItemIsSelectable | Qt.ItemFlag.ItemIsDragEnabled
            colour_item = QTableWidgetItem(color)
            colour_item.setFlags(item_flags)
            model_item = QTableWidgetItem(model)
            model_item.setFlags(item_flags)
            self.table_widget.setItem(i, 0, QTableWidgetItem(color))
            self.table_widget.setItem(i, 1, QTableWidgetItem(model))

        self.resize(400, 400)
        self.show()


if __name__ == "__main__":
    app = QApplication(sys.argv)
    window = Window()
    sys.exit(app.exec())
