ffmpeg_wrapper/ffwrap/gui/queue_model.py

"""The file queue: Job objects plus a Qt table model / progress delegate."""

from __future__ import annotations

import enum
from dataclasses import dataclass, field
from pathlib import Path

from PySide6.QtCore import QAbstractTableModel, QModelIndex, Qt
from PySide6.QtWidgets import QApplication, QStyle, QStyledItemDelegate, QStyleOptionProgressBar

from ..probe import MediaInfo


class Status(enum.Enum):
    PROBING = "Reading..."
    READY = "Ready"
    INVALID = "Unreadable"
    QUEUED = "Queued"
    RUNNING = "Converting"
    DONE = "Done"
    FAILED = "Failed"
    SKIPPED = "Skipped"
    CANCELLED = "Cancelled"


# Jobs in these states are (re)processed when Start is pressed.
RUNNABLE = {Status.READY, Status.FAILED, Status.CANCELLED, Status.SKIPPED}


@dataclass(eq=False)
class Job:
    path: Path
    info: MediaInfo | None = None
    status: Status = Status.PROBING
    progress: float = 0.0
    message: str = ""
    audio_order: list[int] = field(default_factory=list)
    audio_checked: set[int] = field(default_factory=set)
    output: Path | None = None

    @property
    def selected_audio(self) -> list[int]:
        return [i for i in self.audio_order if i in self.audio_checked]

    def set_audio_selection(self, indices: list[int]) -> None:
        """Check ``indices`` and move them (in order) to the top of the list."""
        rest = [i for i in self.audio_order if i not in indices]
        self.audio_order = list(indices) + rest
        self.audio_checked = set(indices)


def fmt_duration(seconds: float | None) -> str:
    if not seconds:
        return ""
    s = int(round(seconds))
    return f"{s // 3600}:{s % 3600 // 60:02d}:{s % 60:02d}"


COLUMNS = ["File", "Duration", "Video", "Audio", "Subs", "Status", "Progress"]
COL_STATUS, COL_PROGRESS = 5, 6


class QueueModel(QAbstractTableModel):
    def __init__(self, parent=None):
        super().__init__(parent)
        self.jobs: list[Job] = []

    # -- Qt model API -------------------------------------------------------
    def rowCount(self, parent=QModelIndex()):
        return 0 if parent.isValid() else len(self.jobs)

    def columnCount(self, parent=QModelIndex()):
        return 0 if parent.isValid() else len(COLUMNS)

    def headerData(self, section, orientation, role=Qt.DisplayRole):
        if orientation == Qt.Horizontal and role == Qt.DisplayRole:
            return COLUMNS[section]
        return None

    def data(self, index, role=Qt.DisplayRole):
        if not index.isValid():
            return None
        job = self.jobs[index.row()]
        col = index.column()
        info = job.info
        if role == Qt.DisplayRole:
            if col == 0:
                return job.path.name
            if col == 1:
                return fmt_duration(info.duration) if info else ""
            if col == 2:
                return info.video.describe() if info and info.video else ("none" if info else "")
            if col == 3:
                if not info:
                    return ""
                return f"{len(job.audio_checked)} of {len(info.audio)}" if info.audio else "none"
            if col == 4:
                return str(len(info.subtitles)) if info else ""
            if col == COL_STATUS:
                return job.status.value + (f": {job.message}" if job.message else "")
            if col == COL_PROGRESS:
                return None if job.status in (Status.PROBING, Status.INVALID) else job.progress
        elif role == Qt.ToolTipRole:
            if col == 0:
                return str(job.path)
            if col == COL_STATUS:
                tip = job.status.value + (f"\n{job.message}" if job.message else "")
                if job.output:
                    tip += f"\nOutput: {job.output}"
                return tip
        elif role == Qt.UserRole:
            return job
        return None

    # -- helpers ------------------------------------------------------------
    def add(self, jobs: list[Job]) -> None:
        if not jobs:
            return
        start = len(self.jobs)
        self.beginInsertRows(QModelIndex(), start, start + len(jobs) - 1)
        self.jobs.extend(jobs)
        self.endInsertRows()

    def remove_rows(self, rows: list[int]) -> None:
        for row in sorted(set(rows), reverse=True):
            self.beginRemoveRows(QModelIndex(), row, row)
            del self.jobs[row]
            self.endRemoveRows()

    def row_of(self, job: Job) -> int:
        for i, j in enumerate(self.jobs):
            if j is job:
                return i
        return -1

    def refresh(self, job: Job) -> None:
        row = self.row_of(job)
        if row >= 0:
            self.dataChanged.emit(self.index(row, 0), self.index(row, len(COLUMNS) - 1))

    def job_at(self, row: int) -> Job | None:
        return self.jobs[row] if 0 <= row < len(self.jobs) else None


class ProgressDelegate(QStyledItemDelegate):
    def paint(self, painter, option, index):
        value = index.data(Qt.DisplayRole)
        if value is None:
            super().paint(painter, option, index)
            return
        bar = QStyleOptionProgressBar()
        bar.rect = option.rect.adjusted(2, 3, -2, -3)
        bar.minimum, bar.maximum = 0, 1000
        bar.progress = int(value * 1000)
        bar.text = f"{value * 100:.0f}%"
        bar.textVisible = True
        bar.state = option.state
        style = option.widget.style() if option.widget else QApplication.style()
        style.drawControl(QStyle.CE_ProgressBar, bar, painter, option.widget)