ProcessDock/processdock/models.py

from __future__ import annotations

from typing import Any

from PySide6.QtCore import QAbstractTableModel, QModelIndex, QRect, QSortFilterProxyModel, Qt, Signal
from PySide6.QtGui import QColor, QFont, QMouseEvent, QPainter
from PySide6.QtWidgets import QStyledItemDelegate, QStyleOptionViewItem

from processdock.services import matches_process


PROCESS_ROLE = Qt.ItemDataRole.UserRole + 1
SORT_ROLE = Qt.ItemDataRole.UserRole + 2


class ProcessTableModel(QAbstractTableModel):
    columns = ("PID", "进程名称", "监听端口", "CPU", "内存", "用户", "状态", "操作")

    def __init__(self, parent=None) -> None:
        super().__init__(parent)
        self._processes: list[dict[str, Any]] = []

    def set_processes(self, processes: list[dict[str, Any]]) -> None:
        self.beginResetModel()
        self._processes = list(processes)
        self.endResetModel()

    @property
    def processes(self) -> list[dict[str, Any]]:
        return self._processes

    def rowCount(self, parent=QModelIndex()) -> int:  # noqa: N802
        return 0 if parent.isValid() else len(self._processes)

    def columnCount(self, parent=QModelIndex()) -> int:  # noqa: N802
        return 0 if parent.isValid() else len(self.columns)

    def headerData(self, section, orientation, role=Qt.ItemDataRole.DisplayRole):  # noqa: N802
        if role == Qt.ItemDataRole.DisplayRole and orientation == Qt.Orientation.Horizontal:
            return self.columns[section]
        return None

    def data(self, index: QModelIndex, role=Qt.ItemDataRole.DisplayRole):
        if not index.isValid() or not 0 <= index.row() < len(self._processes):
            return None
        process = self._processes[index.row()]
        column = index.column()

        if role == PROCESS_ROLE:
            return process
        if role == SORT_ROLE:
            values = (
                process["pid"],
                process["name"].casefold(),
                process["ports"][0] if process["ports"] else -1,
                process["cpu_percent"],
                process["memory_bytes"],
                process["username"].casefold(),
                process["status_label"],
                process["pid"],
            )
            return values[column]
        if role == Qt.ItemDataRole.DisplayRole:
            ports = ", ".join(str(port) for port in process["ports"]) or "—"
            values = (
                str(process["pid"]),
                process["name"],
                ports,
                f'{process["cpu_percent"]:.1f}%',
                process["memory_text"],
                process["username"],
                process["status_label"],
                "当前应用" if process["is_current"] else "结束",
            )
            return values[column]
        if role == Qt.ItemDataRole.TextAlignmentRole:
            if column in (0, 2, 3, 4, 6, 7):
                return Qt.AlignmentFlag.AlignCenter
            return Qt.AlignmentFlag.AlignVCenter | Qt.AlignmentFlag.AlignLeft
        if role == Qt.ItemDataRole.ForegroundRole:
            if column == 6:
                if process["status"] == "running":
                    return QColor("#129B72")
                if process["status"] in {"stopped", "suspended"}:
                    return QColor("#D88918")
            if column == 2 and process["ports"]:
                return QColor("#5B5BD6")
            if column == 7:
                return QColor("#A73A48" if process["can_terminate"] else "#9AA3B2")
        if role == Qt.ItemDataRole.FontRole:
            font = QFont()
            if column in (0, 2, 3, 4):
                font.setFamily("Cascadia Mono")
                font.setPointSize(9)
            if column == 1:
                font.setWeight(QFont.Weight.DemiBold)
            return font
        if role == Qt.ItemDataRole.ToolTipRole:
            if column == 2 and process["port_details"]:
                return "\n".join(
                    f'{item["protocol"]}  {item["address"]}:{item["port"]}'
                    for item in process["port_details"]
                )
            if column == 1:
                return process["exe"] or process["name"]
            if column == 7 and process["is_current"]:
                return "ProcessDock 不允许结束自身"
        return None


class ProcessFilterProxyModel(QSortFilterProxyModel):
    def __init__(self, parent=None) -> None:
        super().__init__(parent)
        self._query = ""
        self.setDynamicSortFilter(True)
        self.setSortRole(SORT_ROLE)

    def set_query(self, query: str) -> None:
        self.beginFilterChange()
        self._query = query
        self.endFilterChange(QSortFilterProxyModel.Direction.Rows)

    def filterAcceptsRow(self, source_row: int, source_parent: QModelIndex) -> bool:  # noqa: N802
        model = self.sourceModel()
        index = model.index(source_row, 0, source_parent)
        process = model.data(index, PROCESS_ROLE)
        return bool(process and matches_process(process, self._query))

    def lessThan(self, left: QModelIndex, right: QModelIndex) -> bool:  # noqa: N802
        left_value = self.sourceModel().data(left, SORT_ROLE)
        right_value = self.sourceModel().data(right, SORT_ROLE)
        return left_value < right_value


class ActionButtonDelegate(QStyledItemDelegate):
    terminate_requested = Signal(object)

    def paint(self, painter: QPainter, option: QStyleOptionViewItem, index: QModelIndex) -> None:
        process = index.data(PROCESS_ROLE)
        if not process:
            return super().paint(painter, option, index)

        painter.save()
        selected = bool(option.state & option.state.State_Selected)
        if selected:
            painter.fillRect(option.rect, QColor("#EEF0FF"))

        rect = self._button_rect(option.rect)
        enabled = bool(process["can_terminate"])
        fill = QColor("#FFF0F1") if enabled else QColor("#F2F4F7")
        text_color = QColor("#B53C4A") if enabled else QColor("#9AA3B2")
        painter.setRenderHint(QPainter.RenderHint.Antialiasing)
        painter.setPen(Qt.PenStyle.NoPen)
        painter.setBrush(fill)
        painter.drawRoundedRect(rect, 7, 7)
        painter.setPen(text_color)
        font = painter.font()
        font.setPointSize(9)
        font.setWeight(QFont.Weight.DemiBold)
        painter.setFont(font)
        text = "结束" if enabled else "当前应用"
        painter.drawText(rect, Qt.AlignmentFlag.AlignCenter, text)
        painter.restore()

    def editorEvent(self, event, model, option, index) -> bool:  # noqa: N802
        process = index.data(PROCESS_ROLE)
        if (
            process
            and process["can_terminate"]
            and isinstance(event, QMouseEvent)
            and event.type() == QMouseEvent.Type.MouseButtonRelease
            and event.button() == Qt.MouseButton.LeftButton
            and self._button_rect(option.rect).contains(event.position().toPoint())
        ):
            self.terminate_requested.emit(process)
            return True
        return super().editorEvent(event, model, option, index)

    @staticmethod
    def _button_rect(cell: QRect) -> QRect:
        width, height = 58, 30
        return QRect(
            cell.center().x() - width // 2,
            cell.center().y() - height // 2,
            width,
            height,
        )