adam/ui/studio.py

from __future__ import annotations

from dataclasses import asdict
from datetime import datetime
import json
import time
from pathlib import Path
from uuid import uuid4

from PySide6.QtCore import QSize, Qt, QThread, QTimer, QUrl, Signal, QObject, QRunnable, QThreadPool
from PySide6.QtGui import QBrush, QDesktopServices, QIcon, QImage, QImageReader, QPixmap
from PySide6.QtWidgets import (
    QAbstractItemView,
    QComboBox,
    QDialog,
    QDialogButtonBox,
    QFileDialog,
    QFrame,
    QGridLayout,
    QHBoxLayout,
    QHeaderView,
    QLabel,
    QLineEdit,
    QListWidget,
    QListWidgetItem,
    QMessageBox,
    QPlainTextEdit,
    QProgressBar,
    QPushButton,
    QSpinBox,
    QTabWidget,
    QTableWidget,
    QTableWidgetItem,
    QVBoxLayout,
    QWidget,
)

from adam.assets import Asset, AssetRegistry
from adam.config import ConfigManager
from adam.dataset_registry import DatasetRegistry
from adam.job_manager import JobManager
from adam.models import Job
from adam.eve import EveResult, EveVisionModel, save_eve_results
from adam.ui.asset_refresh import asset_refresh
from adam.studio import (
    PreviewEvaluation,
    StudioStore,
    TrainingRecipe,
    caption_path,
    checkpoint_files,
    exact_duplicate_groups,
)
from adam.ui.theme import COLORS
from adam.ollama import OllamaClient
from adam.video_captioning import draft_image_caption


def _card() -> QFrame:
    frame = QFrame()
    frame.setProperty("card", True)
    return frame


def _title(text: str) -> QLabel:
    label = QLabel(text)
    label.setObjectName("CardTitle")
    return label


def _header(title: str, subtitle: str) -> QWidget:
    widget = QWidget()
    layout = QVBoxLayout(widget)
    layout.setContentsMargins(0, 0, 0, 10)
    heading = QLabel(title)
    heading.setObjectName("PageTitle")
    detail = QLabel(subtitle)
    detail.setProperty("muted", True)
    detail.setWordWrap(True)
    layout.addWidget(heading)
    layout.addWidget(detail)
    return widget


def _thumbnail(path: Path, width: int, height: int) -> QPixmap:
    """Decode close to display size instead of loading a full-resolution image."""
    reader = QImageReader(str(path))
    reader.setAutoTransform(True)
    source_size = reader.size()
    if source_size.isValid():
        source_size.scale(QSize(width, height), Qt.KeepAspectRatio)
        reader.setScaledSize(source_size)
    image = reader.read()
    return QPixmap.fromImage(image) if not image.isNull() else QPixmap()


class ImageScanWorker(QThread):
    scanned = Signal(object, int)

    def __init__(self, folder: str, token: int, limit: int = 2500) -> None:
        super().__init__()
        self.folder = folder
        self.token = token
        self.limit = limit

    def run(self) -> None:
        root = Path(self.folder).expanduser()
        paths: list[Path] = []
        if root.is_dir():
            for path in root.rglob("*"):
                if self.isInterruptionRequested():
                    return
                if path.is_file() and path.suffix.casefold() in {
                    ".png", ".jpg", ".jpeg", ".webp", ".bmp", ".gif"
                }:
                    paths.append(path)
                    if len(paths) >= self.limit:
                        break
        if not self.isInterruptionRequested():
            self.scanned.emit(sorted(paths), self.token)


class _ImageCaptionSignals(QObject):
    finished = Signal(str, str)


class _ImageCaptionWorker(QRunnable):
    def __init__(self, client: OllamaClient, path: Path, context: str, mode: str, signals: _ImageCaptionSignals):
        super().__init__()
        self.client, self.path, self.context, self.mode, self.signals = client, path, context, mode, signals

    def run(self) -> None:
        try:
            self.signals.finished.emit(draft_image_caption(self.client, self.path, self.context, mode=self.mode), "")
        except Exception as exc:
            self.signals.finished.emit("", str(exc))


class _ImageBatchSignals(QObject):
    progress = Signal(int, int, str)
    caption_ready = Signal(str, str, str)
    finished = Signal(int, int)


class _ImageCaptionBatch(QRunnable):
    """Caption selected images serially so only one full image is sent at a time."""
    def __init__(self, client: OllamaClient, paths: list[Path], mode: str, context: str, signals: _ImageBatchSignals):
        super().__init__()
        self.client, self.paths, self.mode, self.context, self.signals = client, paths, mode, context, signals

    def run(self) -> None:
        successes = failures = 0
        for index, path in enumerate(self.paths, 1):
            self.signals.progress.emit(index, len(self.paths), path.name)
            try:
                context = self.context
                if self.mode == "context":
                    try:
                        context = caption_path(path).read_text(encoding="utf-8").strip()
                    except OSError:
                        context = ""
                    if not context:
                        raise ValueError("Context Caption needs an existing caption for this image.")
                caption = draft_image_caption(self.client, path, context, mode=self.mode)
                self.signals.caption_ready.emit(str(path.resolve()), caption, "")
                successes += 1
            except Exception as exc:
                self.signals.caption_ready.emit(str(path.resolve()), "", str(exc))
                failures += 1
        self.signals.finished.emit(successes, failures)


class ImageCaptionBatchDialog(QDialog):
    """Review serially generated image captions and save only checked drafts."""
    def __init__(self, parent, client: OllamaClient, paths: list[Path], mode: str, context: str):
        super().__init__(parent)
        self.setWindowTitle("AI image caption drafts")
        self.resize(960, 700)
        self.paths = paths
        self.mode = mode
        self.context = context
        self.drafts: dict[str, str] = {}
        self.errors: dict[str, str] = {}
        self.loading_caption = False
        self.signals = _ImageBatchSignals()
        layout = QVBoxLayout(self)
        intro = QLabel(
            f"Drafting {len(paths)} image captions one at a time in {mode.title()} mode. "
            "Review and edit the results, then save the checked captions."
        )
        intro.setWordWrap(True)
        layout.addWidget(intro)
        content = QHBoxLayout()
        self.items = QListWidget()
        self.items.setMinimumWidth(280)
        for path in paths:
            item = QListWidgetItem(f"Waiting · {path.name}")
            item.setData(Qt.UserRole, str(path.resolve()))
            self.items.addItem(item)
        self.items.currentItemChanged.connect(self._select_item)
        content.addWidget(self.items, 1)
        right = QVBoxLayout()
        self.preview = QLabel("Waiting for the first image…")
        self.preview.setAlignment(Qt.AlignCenter)
        self.preview.setMinimumHeight(280)
        self.preview.setStyleSheet(f"background: #050d14; border: 1px solid {COLORS['border']}; border-radius: 8px;")
        self.caption = QPlainTextEdit()
        self.caption.setPlaceholderText("Generated caption appears here for review.")
        self.caption.setEnabled(False)
        self.caption.textChanged.connect(self._caption_changed)
        right.addWidget(self.preview, 2)
        right.addWidget(self.caption, 1)
        content.addLayout(right, 2)
        layout.addLayout(content, 1)
        self.status = QLabel("Starting caption requests…")
        self.status.setWordWrap(True)
        layout.addWidget(self.status)
        buttons = QHBoxLayout()
        check_all = QPushButton("Check all drafts")
        check_all.clicked.connect(lambda: self._set_all_checked(True))
        clear_checks = QPushButton("Clear checks")
        clear_checks.clicked.connect(lambda: self._set_all_checked(False))
        self.save_button = QPushButton("Save checked captions")
        self.save_button.clicked.connect(self._save_checked)
        self.save_button.setEnabled(False)
        close = QPushButton("Close")
        close.clicked.connect(self.reject)
        buttons.addWidget(check_all)
        buttons.addWidget(clear_checks)
        buttons.addStretch(1)
        buttons.addWidget(self.save_button)
        buttons.addWidget(close)
        layout.addLayout(buttons)
        self.signals.progress.connect(self._progress)
        self.signals.caption_ready.connect(self._caption_ready)
        self.signals.finished.connect(self._finished)
        self.worker = _ImageCaptionBatch(client, paths, mode, context, self.signals)
        QThreadPool.globalInstance().start(self.worker)
        if self.items.count():
            self.items.setCurrentRow(0)

    def _progress(self, index: int, total: int, name: str):
        self.status.setText(f"Drafting image {index} of {total}: {name}")

    def _caption_ready(self, path: str, caption: str, error: str):
        self.errors[path] = error
        if caption:
            self.drafts[path] = caption
        for index in range(self.items.count()):
            item = self.items.item(index)
            if str(item.data(Qt.UserRole)) != path:
                continue
            if error:
                item.setText(f"Could not caption · {Path(path).name}")
                item.setToolTip(error)
            else:
                item.setText(f"Ready to review · {Path(path).name}")
                item.setFlags(item.flags() | Qt.ItemIsUserCheckable)
                item.setCheckState(Qt.Unchecked)
                self.save_button.setEnabled(True)
            if self.items.currentItem() is item:
                self._select_item(item, None)
            break

    def _select_item(self, item, _previous):
        self.loading_caption = True
        path = str(item.data(Qt.UserRole)) if item else ""
        pixmap = _thumbnail(Path(path), 520, 320) if path and Path(path).is_file() else QPixmap()
        self.preview.setPixmap(pixmap)
        self.preview.setText("") if not pixmap.isNull() else self.preview.setText("Image preview unavailable")
        self.caption.setPlainText(self.drafts.get(path, ""))
        self.caption.setEnabled(path in self.drafts)
        self.loading_caption = False

    def _caption_changed(self):
        if self.loading_caption:
            return
        item = self.items.currentItem()
        if item:
            path = str(item.data(Qt.UserRole))
            if path in self.drafts:
                self.drafts[path] = self.caption.toPlainText().strip()

    def _set_all_checked(self, checked: bool):
        for index in range(self.items.count()):
            item = self.items.item(index)
            if str(item.data(Qt.UserRole)) in self.drafts:
                item.setCheckState(Qt.Checked if checked else Qt.Unchecked)

    def _finished(self, successes: int, failures: int):
        self.status.setText(f"Caption batch finished · drafts ready: {successes} · errors: {failures}. Review and save the checked captions.")

    def _save_checked(self):
        saved = 0
        try:
            for index in range(self.items.count()):
                item = self.items.item(index)
                path = str(item.data(Qt.UserRole))
                caption = self.drafts.get(path, "").strip()
                if item.checkState() != Qt.Checked or not caption:
                    continue
                target = Path(path).with_suffix(".txt")
                temporary = target.with_suffix(".txt.adam-tmp")
                temporary.write_text(caption + "\n", encoding="utf-8")
                temporary.replace(target)
                saved += 1
            self.status.setText(f"Saved {saved} reviewed caption(s). Other drafts remain unsaved.")
        except OSError as exc:
            QMessageBox.warning(self, "Captions could not be saved", f"Saved {saved} caption(s) before the error.\n\n{exc}")


class EveReviewWorker(QThread):
    progress = Signal(int, str)
    completed = Signal(object)
    failed = Signal(str)

    def __init__(
        self,
        paths: list[Path],
        good_references: list[Path],
        bad_references: list[Path],
        keep_threshold: float,
        reject_threshold: float,
    ) -> None:
        super().__init__()
        self.paths = paths
        self.good_references = good_references
        self.bad_references = bad_references
        self.keep_threshold = keep_threshold
        self.reject_threshold = reject_threshold

    def run(self) -> None:
        try:
            self.progress.emit(1, "Loading EVE's local vision model…")

            def on_progress(done: int, total: int) -> None:
                if self.isInterruptionRequested():
                    raise RuntimeError("EVE review cancelled.")
                percent = 5 + int(done / max(1, total) * 94)
                self.progress.emit(percent, f"EVE analyzed {done} of {total} images…")

            results = EveVisionModel().review(
                self.paths,
                self.good_references,
                self.bad_references,
                keep_threshold=self.keep_threshold,
                reject_threshold=self.reject_threshold,
                progress=on_progress,
            )
            self.progress.emit(100, "EVE finished sorting the dataset.")
            self.completed.emit(results)
        except Exception as exc:
            self.failed.emit(str(exc))


class EveReviewDialog(QDialog):
    applied = Signal(object)

    def __init__(
        self,
        root_path: Path,
        dataset_path: str,
        paths: list[Path],
        store: StudioStore,
        parent: QWidget | None = None,
    ) -> None:
        super().__init__(parent)
        self.root_path = root_path
        self.dataset_path = dataset_path
        self.paths = paths
        self.store = store
        self.results: list[EveResult] = []
        self.worker: EveReviewWorker | None = None
        self.setWindowTitle("EVE · AI Dataset Review")
        self.setMinimumSize(980, 700)
        root = QVBoxLayout(self)
        root.addWidget(_header(
            "EVE · AI Dataset Review",
            "Give EVE a few good examples and optional bad examples. EVE proposes Keep, Reject, or Uncertain; nothing changes until you apply the review.",
        ))

        references = QHBoxLayout()
        self.good_references = self._reference_panel(
            references, "GOOD REFERENCES", "Add good images…", self._add_good_references
        )
        self.bad_references = self._reference_panel(
            references, "BAD REFERENCES (OPTIONAL)", "Add bad images…", self._add_bad_references
        )
        root.addLayout(references)

        controls = QHBoxLayout()
        self.keep_threshold = QSpinBox(); self.keep_threshold.setRange(51, 99); self.keep_threshold.setValue(75); self.keep_threshold.setSuffix("%")
        self.reject_threshold = QSpinBox(); self.reject_threshold.setRange(1, 49); self.reject_threshold.setValue(25); self.reject_threshold.setSuffix("%")
        self.analyze_button = QPushButton("Analyze dataset with EVE")
        self.analyze_button.setProperty("primary", True)
        self.analyze_button.clicked.connect(self._analyze)
        controls.addWidget(QLabel("Keep at or above")); controls.addWidget(self.keep_threshold)
        controls.addWidget(QLabel("Reject at or below")); controls.addWidget(self.reject_threshold)
        controls.addStretch(); controls.addWidget(self.analyze_button)
        root.addLayout(controls)
        self.progress = QProgressBar(); self.progress.setRange(0, 100); self.progress.setValue(0)
        self.status = QLabel(
            "Tip: 3–10 varied good references work best. Bad references help EVE distinguish visually similar mistakes."
        )
        self.status.setProperty("muted", True); self.status.setWordWrap(True)
        root.addWidget(self.progress); root.addWidget(self.status)

        self.tabs = QTabWidget()
        self.result_lists: dict[str, QListWidget] = {}
        for decision, label in (("keep", "KEEP"), ("reject", "REJECT"), ("unreviewed", "UNCERTAIN")):
            gallery = QListWidget()
            gallery.setViewMode(QListWidget.IconMode)
            gallery.setIconSize(QSize(120, 90)); gallery.setGridSize(QSize(165, 145))
            gallery.setResizeMode(QListWidget.Adjust)
            gallery.setSelectionMode(QAbstractItemView.ExtendedSelection)
            gallery.itemDoubleClicked.connect(lambda item: QDesktopServices.openUrl(QUrl.fromLocalFile(str(item.data(Qt.UserRole)))))
            self.result_lists[decision] = gallery
            self.tabs.addTab(gallery, label)
        root.addWidget(self.tabs, 1)

        moves = QHBoxLayout()
        select_all = QPushButton("Select all in current group")
        clear_selection = QPushButton("Clear selection")
        to_keep = QPushButton("Move selected to Keep")
        to_reject = QPushButton("Move selected to Reject")
        to_uncertain = QPushButton("Move selected to Uncertain")
        select_all.clicked.connect(self._select_all_current)
        clear_selection.clicked.connect(self._clear_current_selection)
        to_keep.clicked.connect(lambda: self._move_selected("keep"))
        to_reject.clicked.connect(lambda: self._move_selected("reject"))
        to_uncertain.clicked.connect(lambda: self._move_selected("unreviewed"))
        moves.addWidget(select_all); moves.addWidget(clear_selection)
        moves.addWidget(to_keep); moves.addWidget(to_reject); moves.addWidget(to_uncertain); moves.addStretch()
        root.addLayout(moves)

        buttons = QDialogButtonBox(QDialogButtonBox.Close)
        self.close_button = buttons.button(QDialogButtonBox.Close)
        self.apply_button = QPushButton("Apply EVE review")
        self.apply_button.setProperty("primary", True); self.apply_button.setEnabled(False)
        self.apply_button.clicked.connect(self._apply)
        buttons.addButton(self.apply_button, QDialogButtonBox.AcceptRole)
        buttons.rejected.connect(self.reject)
        root.addWidget(buttons)

    def _reference_panel(self, row: QHBoxLayout, title: str, button_text: str, callback) -> QListWidget:
        frame = _card(); layout = QVBoxLayout(frame); layout.addWidget(_title(title))
        listing = QListWidget(); listing.setMaximumHeight(115)
        button = QPushButton(button_text); button.clicked.connect(callback)
        clear = QPushButton("Clear"); clear.clicked.connect(listing.clear)
        actions = QHBoxLayout(); actions.addWidget(button); actions.addWidget(clear)
        layout.addWidget(listing); layout.addLayout(actions); row.addWidget(frame, 1)
        return listing

    @staticmethod
    def _reference_paths(listing: QListWidget) -> list[Path]:
        return [Path(str(listing.item(index).data(Qt.UserRole))) for index in range(listing.count())]

    def _add_references(self, listing: QListWidget) -> None:
        selected, _ = QFileDialog.getOpenFileNames(
            self, "Choose EVE reference images", self.dataset_path,
            "Images (*.png *.jpg *.jpeg *.webp *.bmp *.gif)",
        )
        existing = {str(path) for path in self._reference_paths(listing)}
        for raw_path in selected:
            path = str(Path(raw_path).resolve())
            if path in existing:
                continue
            item = QListWidgetItem(Path(path).name); item.setData(Qt.UserRole, path)
            listing.addItem(item); existing.add(path)

    def _add_good_references(self) -> None:
        self._add_references(self.good_references)

    def _add_bad_references(self) -> None:
        self._add_references(self.bad_references)

    def _analyze(self) -> None:
        good = self._reference_paths(self.good_references)
        if not good:
            QMessageBox.information(self, "Good references required", "Add at least one good reference image for EVE.")
            return
        if self.reject_threshold.value() >= self.keep_threshold.value():
            QMessageBox.warning(self, "Check thresholds", "Reject confidence must be lower than Keep confidence.")
            return
        self.analyze_button.setEnabled(False); self.apply_button.setEnabled(False)
        self.close_button.setEnabled(False)
        self.progress.setValue(0); self.status.setText("EVE is starting. The first run may download its vision model once.")
        self.worker = EveReviewWorker(
            self.paths, good, self._reference_paths(self.bad_references),
            self.keep_threshold.value() / 100, self.reject_threshold.value() / 100,
        )
        self.worker.progress.connect(self._progress)
        self.worker.completed.connect(self._completed)
        self.worker.failed.connect(self._failed)
        self.worker.finished.connect(self._worker_finished)
        self.worker.start()

    def _progress(self, percent: int, message: str) -> None:
        self.progress.setValue(percent); self.status.setText(message)

    def _completed(self, results: object) -> None:
        if not isinstance(results, list):
            self._failed("EVE returned an invalid review.")
            return
        self.results = results
        save_eve_results(self.root_path, self.dataset_path, self.results)
        self._rebuild_results()
        self.apply_button.setEnabled(True)
        counts = {key: sum(result.suggestion == key for result in self.results) for key in self.result_lists}
        self.status.setText(
            f"EVE proposes {counts['keep']} Keep, {counts['reject']} Reject, and {counts['unreviewed']} Uncertain. Review both sides before applying."
        )

    def _failed(self, message: str) -> None:
        self.status.setText(f"EVE could not finish: {message}")
        QMessageBox.warning(self, "EVE review stopped", message)

    def _worker_finished(self) -> None:
        self.analyze_button.setEnabled(True)
        self.close_button.setEnabled(True)
        if self.worker:
            self.worker.deleteLater()
        self.worker = None

    def _rebuild_results(self) -> None:
        for listing in self.result_lists.values():
            listing.clear()
        for result in self.results:
            path = Path(result.path)
            decision_score = result.match_score if result.suggestion == "keep" else 1.0 - result.match_score if result.suggestion == "reject" else result.match_score
            text = f"{path.name}\n{decision_score * 100:.0f}% " + ("match" if result.suggestion != "reject" else "reject confidence")
            item = QListWidgetItem(QIcon(_thumbnail(path, 120, 90)), text)
            item.setData(Qt.UserRole, result.path)
            self.result_lists[result.suggestion].addItem(item)
        self._update_tab_labels()

    def _update_tab_labels(self) -> None:
        labels = {"keep": "KEEP", "reject": "REJECT", "unreviewed": "UNCERTAIN"}
        for index, key in enumerate(("keep", "reject", "unreviewed")):
            self.tabs.setTabText(index, f"{labels[key]} ({self.result_lists[key].count()})")

    def _current_result_list(self) -> QListWidget:
        return self.tabs.currentWidget()

    def _select_all_current(self) -> None:
        self._current_result_list().selectAll()

    def _clear_current_selection(self) -> None:
        self._current_result_list().clearSelection()

    def _move_selected(self, decision: str) -> None:
        source = self._current_result_list()
        selected_items = source.selectedItems()
        if not selected_items:
            return
        destination = self.result_lists[decision]
        if source is destination:
            return
        selected_paths = {str(item.data(Qt.UserRole)) for item in selected_items}
        for result in self.results:
            if result.path in selected_paths:
                result.suggestion = decision
                result.decision_confidence = 1.0
        # Preserve the existing thumbnails and transfer only the selected items.
        # This avoids decoding the full dataset again after every manual edit.
        selected_rows = sorted((source.row(item) for item in selected_items), reverse=True)
        moved_items = [source.takeItem(row) for row in selected_rows]
        for item in reversed(moved_items):
            destination.addItem(item)
            item.setSelected(True)
        self._update_tab_labels()

    def _apply(self) -> None:
        counts = {key: sum(result.suggestion == key for result in self.results) for key in self.result_lists}
        answer = QMessageBox.question(
            self, "Apply EVE review",
            f"Apply {counts['keep']} Keep and {counts['reject']} Reject decisions?\n\n"
            f"The {counts['unreviewed']} uncertain images will remain unreviewed. Rejected files are not moved until you choose Exclude rejected.",
        )
        if answer != QMessageBox.Yes:
            return
        decisions = {result.path: result.suggestion for result in self.results}
        changed = self.store.apply_decisions(self.dataset_path, decisions)
        save_eve_results(self.root_path, self.dataset_path, self.results)
        self.applied.emit(decisions)
        self.status.setText(f"Applied EVE review ({changed} decisions changed). You can continue reviewing manually.")

    def closeEvent(self, event) -> None:
        if self.worker and self.worker.isRunning():
            self.worker.requestInterruption()
            self.status.setText("EVE is stopping after the current image. The window will be safe to close when analysis ends.")
            event.ignore()
            return
        super().closeEvent(event)


class DatasetLocationDialog(QDialog):
    def __init__(
        self,
        root_path: Path,
        config: ConfigManager,
        assets: AssetRegistry,
        parent: QWidget | None = None,
    ) -> None:
        super().__init__(parent)
        self.registry = DatasetRegistry(root_path, config)
        self.assets = assets
        self.setWindowTitle("Remembered Dataset Locations")
        self.setMinimumSize(740, 420)
        root = QVBoxLayout(self)
        root.addWidget(_header(
            "Remembered dataset locations",
            "Add dataset parent folders here so ADAM Remote can find them without exposing arbitrary PC browsing.",
        ))
        self.list = QListWidget()
        root.addWidget(self.list, 1)
        actions = QHBoxLayout()
        add = QPushButton("Add location...")
        remove = QPushButton("Remove selected")
        refresh = QPushButton("Refresh")
        add.setProperty("primary", True)
        add.clicked.connect(self._add)
        remove.clicked.connect(self._remove)
        refresh.clicked.connect(self.refresh)
        actions.addWidget(add)
        actions.addWidget(remove)
        actions.addWidget(refresh)
        actions.addStretch()
        root.addLayout(actions)
        buttons = QDialogButtonBox(QDialogButtonBox.Close)
        buttons.rejected.connect(self.reject)
        root.addWidget(buttons)
        self.refresh()

    def refresh(self) -> None:
        self.registry.load()
        self.registry.discover_into_assets(self.assets, persist=True)
        self.list.clear()
        for location in self.registry.known_locations():
            status = "available" if Path(location.path).is_dir() else "unavailable"
            label = f"{location.name}\n{status} - {location.source} - {location.path}"
            item = QListWidgetItem(label)
            item.setData(Qt.UserRole, location.id)
            self.list.addItem(item)

    def _add(self) -> None:
        selected = QFileDialog.getExistingDirectory(
            self, "Remember dataset location", str(Path.home())
        )
        if not selected:
            return
        try:
            self.registry.register_location(selected, name=Path(selected).name, source="user")
        except ValueError as exc:
            QMessageBox.warning(self, "Location not saved", str(exc))
            return
        self.refresh()

    def _remove(self) -> None:
        item = self.list.currentItem()
        if not item:
            return
        location_id = str(item.data(Qt.UserRole) or "")
        self.registry.remove_location(location_id)
        self.refresh()


class DatasetReviewTab(QWidget):
    def __init__(self, assets: AssetRegistry, store: StudioStore, config: ConfigManager) -> None:
        super().__init__()
        self.assets = assets
        self.store = store
        self.config = config
        self.registry = DatasetRegistry(assets.path.parent.parent, config)
        self.paths: list[Path] = []
        self.dataset_path = ""
        self._load_index = 0
        self._load_token = 0
        self._requested_row = 0
        self._scan_workers: set[ImageScanWorker] = set()
        self._caption_workers: list[tuple[_ImageCaptionWorker, _ImageCaptionSignals]] = []

        root = QVBoxLayout(self)
        root.setContentsMargins(10, 14, 10, 10)
        top = QVBoxLayout()
        dataset_row = QHBoxLayout()
        review_actions = QHBoxLayout()
        self.dataset = QComboBox()
        self.dataset.setMinimumWidth(300)
        browse = QPushButton("Open another dataset…")
        browse.clicked.connect(self._browse)
        locations = QPushButton("Remember locations…")
        locations.clicked.connect(self._manage_locations)
        refresh = QPushButton("Refresh")
        refresh.clicked.connect(self.refresh)
        duplicates = QPushButton("Check duplicates")
        duplicates.clicked.connect(self._duplicates)
        keep_all = QPushButton("Keep all images")
        keep_all.setProperty("primary", True)
        keep_all.setToolTip(
            "Accept every image in this dataset, then reject only the individual images you do not want."
        )
        keep_all.clicked.connect(self._keep_all_images)
        eve_review = QPushButton("EVE AI Review…")
        eve_review.setToolTip(
            "Sort this dataset from visual reference images, then review EVE's Keep, Reject, and Uncertain groups."
        )
        eve_review.clicked.connect(self._open_eve_review)
        apply_rejected = QPushButton("Exclude rejected")
        apply_rejected.setToolTip(
            "Move rejected images out of the training dataset into ADAM's recoverable quarantine."
        )
        apply_rejected.clicked.connect(self._apply_rejected)
        restore_rejected = QPushButton("Restore excluded")
        restore_rejected.clicked.connect(self._restore_rejected)
        dataset_row.addWidget(QLabel("Dataset"))
        dataset_row.addWidget(self.dataset, 1)
        dataset_row.addWidget(browse)
        dataset_row.addWidget(locations)
        dataset_row.addWidget(refresh)
        review_actions.addWidget(duplicates)
        review_actions.addWidget(keep_all)
        review_actions.addWidget(eve_review)
        review_actions.addWidget(apply_rejected)
        review_actions.addWidget(restore_rejected)
        review_actions.addStretch()
        top.addLayout(dataset_row)
        top.addLayout(review_actions)
        root.addLayout(top)

        self.summary = QLabel("Choose a dataset to begin reviewing it.")
        self.summary.setProperty("muted", True)
        root.addWidget(self.summary)

        body = QHBoxLayout()
        self.gallery = QListWidget()
        self.gallery.setViewMode(QListWidget.IconMode)
        self.gallery.setIconSize(QPixmap(150, 110).size())
        self.gallery.setGridSize(QPixmap(178, 158).size())
        self.gallery.setResizeMode(QListWidget.Adjust)
        self.gallery.setSelectionMode(QAbstractItemView.ExtendedSelection)
        self.gallery.setToolTip("Select multiple images with Ctrl or Shift. Keep, Reject, and AI Draft Captions apply to the selection.")
        self.gallery.currentRowChanged.connect(self._selected)
        body.addWidget(self.gallery, 3)

        detail = _card()
        detail.setMinimumWidth(330)
        detail_layout = QVBoxLayout(detail)
        detail_layout.addWidget(_title("IMAGE REVIEW"))
        self.preview = QLabel("Select an image")
        self.preview.setAlignment(Qt.AlignCenter)
        self.preview.setMinimumHeight(230)
        self.preview.setStyleSheet(
            f"background: #050d14; border: 1px solid {COLORS['border']}; border-radius: 8px;"
        )
        self.file_label = QLabel()
        self.file_label.setWordWrap(True)
        self.file_label.setProperty("muted", True)
        self.caption = QPlainTextEdit()
        self.caption.setPlaceholderText("Caption text stored beside the image…")
        self.caption.setMaximumHeight(120)
        detail_layout.addWidget(self.preview)
        detail_layout.addWidget(self.file_label)
        detail_layout.addWidget(QLabel("Caption"))
        self.caption_mode = QComboBox()
        self.caption_mode.addItem("Simple Caption", "simple")
        self.caption_mode.addItem("Detailed Caption", "detailed")
        self.caption_mode.addItem("Context Caption", "context")
        self.caption_mode.setToolTip(
            "Simple: a few plain sentences. Detailed: a longer visual description. "
            "Context: use the current caption text as context and retain it at the start."
        )
        detail_layout.addWidget(self.caption_mode)
        self.caption_context = QLineEdit()
        self.caption_context.setPlaceholderText("Optional trigger word or context for Simple and Detailed modes")
        detail_layout.addWidget(self.caption_context)
        self.caption_mode.currentIndexChanged.connect(
            lambda _index: self.caption_context.setVisible(self.caption_mode.currentData() != "context")
        )
        self.caption_context.setVisible(True)
        detail_layout.addWidget(self.caption)
        buttons = QGridLayout()
        keep = QPushButton("Keep selected")
        keep.setProperty("primary", True)
        reject = QPushButton("Reject selected")
        reject.setProperty("danger", True)
        restore = QPushButton("Unreview selected")
        save_caption = QPushButton("Save caption")
        draft_caption = QPushButton("AI Draft Captions")
        draft_caption.setObjectName("aiDraftCaptionsButton")
        draft_caption.setToolTip("Describe the selected image with the configured Ollama vision model. Review the draft before saving.")
        keep.clicked.connect(lambda: self._decide("keep"))
        reject.clicked.connect(lambda: self._decide("reject"))
        restore.clicked.connect(lambda: self._decide("unreviewed"))
        save_caption.clicked.connect(self._save_caption)
        draft_caption.clicked.connect(self._ai_draft_caption)
        buttons.addWidget(keep, 0, 0)
        buttons.addWidget(reject, 0, 1)
        buttons.addWidget(restore, 1, 0)
        buttons.addWidget(save_caption, 1, 1)
        buttons.addWidget(draft_caption, 2, 0, 1, 2)
        detail_layout.addLayout(buttons)
        detail_layout.addStretch()
        body.addWidget(detail, 2)
        root.addLayout(body, 1)

        self.dataset.currentIndexChanged.connect(self.refresh)
        self.reload_assets()

    def reload_assets(self) -> None:
        if not hasattr(self, "_asset_refresh"):
            self._asset_refresh = asset_refresh(self.assets, self.config)
            self._asset_refresh.changed.connect(self.reload_assets)
        self._asset_refresh.request()
        signature = tuple((asset.id, asset.name, asset.path) for asset in self.assets.assets if asset.kind == "dataset")
        if signature == getattr(self, "_dataset_choices_signature", None):
            return
        self._dataset_choices_signature = signature
        current = self.dataset.currentData()
        self.dataset.blockSignals(True)
        self.dataset.clear()
        for asset in self.assets.assets:
            if asset.kind == "dataset" and Path(asset.path).is_dir():
                self.dataset.addItem(asset.name, asset.path)
        index = self.dataset.findData(current)
        if index >= 0:
            self.dataset.setCurrentIndex(index)
        self.dataset.blockSignals(False)
        self.refresh()

    def _browse(self) -> None:
        selected = QFileDialog.getExistingDirectory(
            self, "Choose an image dataset", self.dataset_path or str(Path.home())
        )
        if not selected:
            return
        index = self.dataset.findData(selected)
        if index < 0:
            asset = self.assets.register(
                kind="dataset",
                name=Path(selected).name,
                path=selected,
                metadata={"dataset_registry_source": "studio"},
            )
            self.registry.record_for_path(asset.path)
            self.dataset.addItem(asset.name, asset.path)
            index = self.dataset.count() - 1
        self.dataset.setCurrentIndex(index)

    def _manage_locations(self) -> None:
        dialog = DatasetLocationDialog(
            self.assets.path.parent.parent,
            self.config,
            self.assets,
            self,
        )
        dialog.exec()
        self._asset_refresh.request(force=True)
        self.reload_assets()

    def refresh(self) -> None:
        previous_path = self.dataset_path
        self.dataset_path = str(self.dataset.currentData() or "")
        if self.dataset_path != previous_path:
            self._requested_row = 0
        self.gallery.clear()
        self._load_token += 1
        token = self._load_token
        self._load_index = 0
        for worker in self._scan_workers:
            worker.requestInterruption()
        if not self.dataset_path:
            self.paths = []
            self.summary.setText("Choose a dataset to begin reviewing it.")
            self.preview.setText("No supported images found")
            return
        self.summary.setText("Scanning dataset…")
        worker = ImageScanWorker(self.dataset_path, token)
        self._scan_workers.add(worker)
        worker.scanned.connect(self._scan_finished)
        worker.finished.connect(
            lambda worker=worker: self._scan_workers.discard(worker)
        )
        worker.finished.connect(worker.deleteLater)
        worker.start()

    def _scan_finished(self, paths: object, token: int) -> None:
        if token != self._load_token or not isinstance(paths, list):
            return
        self.paths = paths
        self._update_review_summary(loading=True)
        if not self.paths:
            self.preview.setText("No supported images found")
            return
        QTimer.singleShot(0, lambda: self._load_next_thumbnail(token))

    def _update_review_summary(self, *, loading: bool = False) -> None:
        """Update counts without rescanning files or rebuilding thumbnails."""
        review = self.store.review(self.dataset_path)
        kept = rejected = 0
        for path in self.paths:
            decision = review.decisions.get(str(path.resolve()), "unreviewed") if review else "unreviewed"
            kept += decision == "keep"
            rejected += decision == "reject"
        reviewed = kept + rejected
        captions = sum(caption_path(path).is_file() for path in self.paths)
        self.summary.setText(
            f"{len(self.paths)} images · {reviewed} reviewed · {kept} kept · "
            f"{rejected} rejected · {captions} captions"
            + (" · Loading thumbnails…" if loading else "")
        )

    @staticmethod
    def _style_review_item(item: QListWidgetItem, path: Path, decision: str) -> None:
        prefix = {"keep": "✓ ", "reject": "× ", "unreviewed": ""}[decision]
        item.setText(prefix + path.name)
        if decision == "reject":
            item.setForeground(Qt.red)
        elif decision == "keep":
            item.setForeground(Qt.green)
        else:
            item.setForeground(QBrush())

    def _load_next_thumbnail(self, token: int) -> None:
        """Decode one image per event-loop turn so large datasets stay responsive."""
        if token != self._load_token or self._load_index >= len(self.paths):
            if token == self._load_token:
                self.summary.setText(self.summary.text().replace(" · Loading thumbnails…", ""))
            return
        path = self.paths[self._load_index]
        review = self.store.review(self.dataset_path) if self.dataset_path else None
        decision = review.decisions.get(str(path.resolve()), "unreviewed") if review else "unreviewed"
        item = QListWidgetItem(path.name)
        item.setData(Qt.UserRole, str(path))
        pixmap = _thumbnail(path, 150, 110)
        if not pixmap.isNull():
            item.setIcon(QIcon(pixmap))
        self._style_review_item(item, path, decision)
        self.gallery.addItem(item)
        target_row = min(self._requested_row, len(self.paths) - 1)
        if self._load_index == target_row:
            self.gallery.setCurrentRow(target_row)
        self._load_index += 1
        QTimer.singleShot(0, lambda: self._load_next_thumbnail(token))

    def _current_path(self) -> Path | None:
        item = self.gallery.currentItem()
        return Path(str(item.data(Qt.UserRole))) if item else None

    def _selected(self, _row: int) -> None:
        path = self._current_path()
        if not path:
            return
        pixmap = _thumbnail(path, 310, 260)
        self.preview.setPixmap(
            pixmap
        )
        self.file_label.setText(str(path))
        try:
            text = caption_path(path).read_text(encoding="utf-8")
        except OSError:
            text = ""
        self.caption.setPlainText(text)
        self.caption_context.clear()

    def _decide(self, decision: str) -> None:
        if not self.dataset_path:
            return
        selected = self.gallery.selectedItems()
        if not selected and self.gallery.currentItem():
            selected = [self.gallery.currentItem()]
        if not selected:
            return
        for item in selected:
            path = Path(str(item.data(Qt.UserRole)))
            self.store.set_decision(self.dataset_path, str(path), decision)
            self._style_review_item(item, path, decision)
        self._update_review_summary()
        if len(selected) == 1 and decision in {"keep", "reject"} and self.gallery.currentRow() + 1 < self.gallery.count():
            row = self.gallery.currentRow()
            self.gallery.setCurrentRow(row + 1)
            self.gallery.scrollToItem(self.gallery.currentItem())

    def _save_caption(self) -> None:
        path = self._current_path()
        if not path:
            return
        try:
            caption_path(path).write_text(
                self.caption.toPlainText().strip() + "\n", encoding="utf-8"
            )
        except OSError as exc:
            QMessageBox.warning(self, "Caption not saved", str(exc))
            return
        self.summary.setText(self.summary.text() + " · Caption saved")

    def _ai_draft_caption(self) -> None:
        selected_items = self.gallery.selectedItems()
        if not selected_items and self.gallery.currentItem():
            selected_items = [self.gallery.currentItem()]
        paths = [Path(str(item.data(Qt.UserRole))) for item in selected_items]
        if not paths:
            QMessageBox.information(self, "Select images", "Select one or more images before drafting captions.")
            return
        model = str(self.config.get("ollama_model", "")).strip()
        if not model:
            QMessageBox.warning(self, "Ollama model not configured", "Choose an Ollama model in Settings → Planning model first.")
            return
        client = OllamaClient(
            str(self.config.get("ollama_url", "http://localhost:11434")), model,
            timeout=600,
            chat_max_tokens=min(512, int(self.config.get("ollama_chat_max_tokens", 512))),
            chat_response_length="balanced",
        )
        mode = str(self.caption_mode.currentData())
        if len(paths) > 1:
            dialog = ImageCaptionBatchDialog(self, client, paths, mode, self.caption_context.text())
            dialog.exec()
            self._update_review_summary()
            return
        path = paths[0]
        button = self.findChild(QPushButton, "aiDraftCaptionsButton")
        if button:
            button.setEnabled(False)
            button.setText("Drafting…")
        signals = _ImageCaptionSignals()
        context = self.caption.toPlainText() if mode == "context" else self.caption_context.text()
        worker = _ImageCaptionWorker(client, path, context, mode, signals)
        self._caption_workers.append((worker, signals))
        signals.finished.connect(
            lambda caption, error, worker=worker, signals=signals, path=path:
                self._image_caption_ready(worker, signals, path, caption, error)
        )
        QThreadPool.globalInstance().start(worker)

    def _image_caption_ready(self, worker, signals, path: Path, caption: str, error: str) -> None:
        self._caption_workers = [(w, s) for w, s in self._caption_workers if w is not worker]
        button = self.findChild(QPushButton, "aiDraftCaptionsButton")
        if button:
            button.setEnabled(True)
            button.setText("AI Draft Captions")
        if error:
            QMessageBox.warning(self, "Caption draft failed", error)
        elif self._current_path() == path:
            self.caption.setPlainText(caption)
            self.summary.setText(self.summary.text() + " · Draft ready to review")

    def _apply_rejected(self) -> None:
        if not self.dataset_path:
            return
        review = self.store.review(self.dataset_path)
        count = sum(value == "reject" for value in review.decisions.values())
        if not count:
            QMessageBox.information(
                self, "No rejected images", "Mark images as rejected before excluding them."
            )
            return
        answer = QMessageBox.question(
            self,
            "Exclude rejected images",
            f"Move {count} rejected image(s) and their captions out of this training dataset?\n\n"
            "They remain recoverable with Restore excluded.",
        )
        if answer != QMessageBox.Yes:
            return
        moved = self.store.apply_rejections(self.dataset_path)
        self._requested_row = 0
        self.refresh()
        self.summary.setText(
            f"Excluded {moved} rejected image(s) from training. They remain recoverable."
        )

    def _keep_all_images(self) -> None:
        if not self.dataset_path or not self.paths:
            QMessageBox.information(
                self, "No images", "Choose a dataset and wait for its images to finish loading."
            )
            return
        answer = QMessageBox.question(
            self,
            "Keep all images",
            f"Mark all {len(self.paths)} images in this dataset as kept?\n\n"
            "You can still reject individual images afterward.",
        )
        if answer != QMessageBox.Yes:
            return
        changed = self.store.set_all_decisions(
            self.dataset_path, self.paths, "keep"
        )
        for row in range(self.gallery.count()):
            item = self.gallery.item(row)
            self._style_review_item(item, Path(str(item.data(Qt.UserRole))), "keep")
        self._update_review_summary()
        self.summary.setText(
            self.summary.text()
            + (f" · All images kept ({changed} changed)" if changed else " · All images already kept")
        )

    def _open_eve_review(self) -> None:
        if not self.dataset_path or not self.paths:
            QMessageBox.information(
                self, "No dataset ready", "Choose a dataset and wait for its image scan to finish."
            )
            return
        root_path = self.store.path.parent.parent
        dialog = EveReviewDialog(root_path, self.dataset_path, self.paths, self.store, self)
        dialog.applied.connect(self._eve_decisions_applied)
        dialog.exec()

    def _eve_decisions_applied(self, decisions: object) -> None:
        if not isinstance(decisions, dict):
            return
        for row in range(self.gallery.count()):
            item = self.gallery.item(row)
            path = Path(str(item.data(Qt.UserRole)))
            decision = str(decisions.get(str(path.resolve()), "unreviewed"))
            self._style_review_item(item, path, decision)
        self._update_review_summary()
        self.summary.setText(self.summary.text() + " · EVE review applied")

    def _restore_rejected(self) -> None:
        if not self.dataset_path:
            return
        restored = self.store.restore_rejections(self.dataset_path)
        self._requested_row = 0
        self.refresh()
        self.summary.setText(
            f"Restored {restored} excluded image(s)."
            if restored
            else "No excluded images were available to restore."
        )

    def _duplicates(self) -> None:
        groups = exact_duplicate_groups(self.paths)
        exact_members = {value for group in groups for value in group}
        hashes: list[tuple[Path, int]] = []
        for path in self.paths[:500]:
            image = QImage(str(path))
            if image.isNull():
                continue
            sample = image.convertToFormat(QImage.Format_Grayscale8).scaled(
                8, 8, Qt.IgnoreAspectRatio, Qt.SmoothTransformation
            )
            values = [
                sample.pixelColor(x, y).red() for y in range(8) for x in range(8)
            ]
            average = sum(values) / len(values)
            bits = 0
            for index, value in enumerate(values):
                if value >= average:
                    bits |= 1 << index
            hashes.append((path, bits))
        near: list[tuple[str, str]] = []
        for index, (first_path, first_hash) in enumerate(hashes):
            for second_path, second_hash in hashes[index + 1 :]:
                if str(first_path) in exact_members and str(second_path) in exact_members:
                    continue
                if (first_hash ^ second_hash).bit_count() <= 5:
                    near.append((first_path.name, second_path.name))
                    if len(near) >= 20:
                        break
            if len(near) >= 20:
                break
        if not groups and not near:
            QMessageBox.information(
                self,
                "Duplicate check",
                "No exact or visually similar duplicate candidates were found.",
            )
            return
        lines = [
            " = ".join(Path(value).name for value in group) for group in groups[:20]
        ]
        lines.extend(f"≈ {first} / {second}" for first, second in near)
        scope = (
            " Visual similarity checked the first 500 images."
            if len(self.paths) > 500
            else ""
        )
        QMessageBox.warning(
            self,
            "Duplicate candidates found",
            f"{len(groups)} exact group(s) and {len(near)} visually similar "
            f"candidate pair(s).{scope}\n\n" + "\n".join(lines),
        )


class ExperimentsTab(QWidget):
    def __init__(
        self, jobs: JobManager, assets: AssetRegistry, store: StudioStore
    ) -> None:
        super().__init__()
        self.jobs = jobs
        self.assets = assets
        self.store = store
        root = QVBoxLayout(self)
        root.setContentsMargins(10, 14, 10, 10)
        hint = QLabel(
            "Select one run for details or two runs to compare their recipes and outcomes."
        )
        hint.setProperty("muted", True)
        root.addWidget(hint)
        self.table = QTableWidget(0, 7)
        self.table.setHorizontalHeaderLabels(
            ["RUN", "PROJECT", "TRAINER", "EPOCHS", "STATUS", "CREATED", "OUTPUT"]
        )
        self.table.setSelectionBehavior(QAbstractItemView.SelectRows)
        self.table.setSelectionMode(QAbstractItemView.ExtendedSelection)
        self.table.setEditTriggers(QAbstractItemView.NoEditTriggers)
        self.table.verticalHeader().hide()
        header = self.table.horizontalHeader()
        header.setSectionResizeMode(1, QHeaderView.Stretch)
        for column in (0, 2, 3, 4, 5, 6):
            header.setSectionResizeMode(column, QHeaderView.ResizeToContents)
        self.table.itemSelectionChanged.connect(self._show_selection)
        root.addWidget(self.table, 3)
        detail = _card()
        detail_layout = QVBoxLayout(detail)
        detail_layout.addWidget(_title("RUN COMPARISON"))
        self.comparison = QPlainTextEdit()
        self.comparison.setReadOnly(True)
        self.comparison.setMaximumHeight(180)
        detail_layout.addWidget(self.comparison)
        actions = QHBoxLayout()
        self.best = QPushButton("Mark model as best")
        self.open = QPushButton("Open output")
        self.recipe = QPushButton("Save as reusable recipe")
        self.best.clicked.connect(self._toggle_best)
        self.open.clicked.connect(self._open_output)
        self.recipe.clicked.connect(self._save_recipe)
        actions.addWidget(self.best)
        actions.addWidget(self.recipe)
        actions.addWidget(self.open)
        detail_layout.addLayout(actions)
        root.addWidget(detail, 1)
        self.refresh_timer = QTimer(self)
        self.refresh_timer.setSingleShot(True)
        self.refresh_timer.setInterval(1000)
        self.refresh_timer.timeout.connect(self._refresh_if_visible)
        jobs.job_created.connect(self._schedule_refresh)
        jobs.job_updated.connect(self._schedule_refresh)
        self.refresh()

    def _schedule_refresh(self, _job: Job) -> None:
        if self.isVisible() and not self.refresh_timer.isActive():
            self.refresh_timer.start()

    def _refresh_if_visible(self) -> None:
        if self.isVisible():
            self.refresh()

    @staticmethod
    def _training(job: Job) -> tuple[str, int]:
        for step in job.plan.steps:
            if step.tool_id.endswith("_trainer"):
                return step.tool_id.removesuffix("_trainer").upper(), int(
                    step.arguments.get("epochs", 0) or 0
                )
        return "—", 0

    def refresh(self) -> None:
        selected = {job.id for job in self._selected_jobs()}
        self.table.setRowCount(len(self.jobs.jobs))
        for row, job in enumerate(self.jobs.jobs):
            trainer, epochs = self._training(job)
            values = [
                job.id,
                job.plan.project_name,
                trainer,
                str(epochs or "—"),
                job.status.value,
                self._date(job.created_at),
                "Ready" if job.output_folder else "—",
            ]
            for column, value in enumerate(values):
                item = QTableWidgetItem(value)
                if column != 1:
                    item.setTextAlignment(Qt.AlignCenter)
                self.table.setItem(row, column, item)
            if job.id in selected:
                self.table.selectRow(row)
        self._show_selection()

    def _selected_jobs(self) -> list[Job]:
        rows = sorted({index.row() for index in self.table.selectionModel().selectedRows()})
        result = []
        for row in rows:
            item = self.table.item(row, 0)
            if item:
                try:
                    result.append(self.jobs.get(item.text()))
                except KeyError:
                    pass
        return result

    def _show_selection(self) -> None:
        jobs = self._selected_jobs()
        if not jobs:
            self.comparison.setPlainText("Select a run to inspect it.")
        else:
            blocks = []
            for job in jobs[:2]:
                trainer, epochs = self._training(job)
                duration = self._duration(job)
                blocks.append(
                    f"{job.plan.project_name} ({job.id})\n"
                    f"  Trainer: {trainer} · Epochs: {epochs or '—'}\n"
                    f"  Status: {job.status.value} · Progress: {job.progress}%\n"
                    f"  Duration: {duration} · Output: {job.output_folder or 'Not recorded'}"
                )
            self.comparison.setPlainText("\n\n".join(blocks))
        one = len(jobs) == 1
        self.open.setEnabled(one and bool(jobs[0].output_folder))
        self.recipe.setEnabled(one and self._training(jobs[0])[0] != "—")
        model = self._model_for_job(jobs[0]) if one else None
        self.best.setEnabled(model is not None)
        if model:
            self.best.setText(
                "Unmark best" if model.id in self.store.best_models else "Mark model as best"
            )

    def _model_for_job(self, job: Job) -> Asset | None:
        if not job.output_folder:
            return None
        output = Path(job.output_folder)
        return next(
            (
                asset
                for asset in self.assets.assets
                if asset.kind == "model"
                and (
                    Path(asset.path) == output
                    or output in Path(asset.path).parents
                    or Path(asset.path) in output.parents
                )
            ),
            None,
        )

    def _toggle_best(self) -> None:
        jobs = self._selected_jobs()
        model = self._model_for_job(jobs[0]) if len(jobs) == 1 else None
        if model:
            self.store.toggle_best(model.id)
            self._show_selection()

    def _open_output(self) -> None:
        jobs = self._selected_jobs()
        if len(jobs) == 1 and jobs[0].output_folder:
            QDesktopServices.openUrl(QUrl.fromLocalFile(jobs[0].output_folder))

    def _save_recipe(self) -> None:
        jobs = self._selected_jobs()
        if len(jobs) != 1:
            return
        job = jobs[0]
        for step in job.plan.steps:
            if step.tool_id.endswith("_trainer"):
                self.store.add_recipe(
                    TrainingRecipe(
                        name=f"{job.plan.project_name} recipe",
                        trainer=step.tool_id.removesuffix("_trainer"),
                        epochs=int(step.arguments.get("epochs", 100) or 100),
                        base_model=str(step.arguments.get("base_model", "")),
                        notes=f"Created from job {job.id}",
                    )
                )
                QMessageBox.information(
                    self, "Recipe saved", "The run settings are now reusable in Recipes."
                )
                return

    @staticmethod
    def _date(value: str) -> str:
        try:
            return datetime.fromisoformat(value).astimezone().strftime("%b %d · %H:%M")
        except ValueError:
            return value[:16]

    @staticmethod
    def _duration(job: Job) -> str:
        if not job.started_at or not job.ended_at:
            return "In progress" if job.started_at else "Not started"
        try:
            seconds = int(
                (
                    datetime.fromisoformat(job.ended_at)
                    - datetime.fromisoformat(job.started_at)
                ).total_seconds()
            )
            return f"{seconds // 3600}h {(seconds % 3600) // 60}m" if seconds >= 3600 else f"{seconds // 60}m {seconds % 60}s"
        except ValueError:
            return "Unknown"


class PreviewLabTab(QWidget):
    plan_requested = Signal(str)

    def __init__(self, assets: AssetRegistry, store: StudioStore) -> None:
        super().__init__()
        self.assets = assets
        self.store = store
        self._output_paths: list[Path] = []
        self._output_index = 0
        self._output_token = 0
        self._output_scan_workers: set[ImageScanWorker] = set()
        root = QHBoxLayout(self)
        root.setContentsMargins(10, 14, 10, 10)
        form = _card()
        form_layout = QGridLayout(form)
        form_layout.addWidget(_title("PROMPT & CHECKPOINT EVALUATION"), 0, 0, 1, 2)
        self.model = QComboBox()
        self.checkpoint = QComboBox()
        self.prompt = QPlainTextEdit()
        self.prompt.setPlaceholderText(
            "A consistent evaluation prompt for comparing model checkpoints…"
        )
        self.prompt.setMaximumHeight(110)
        self.seed = QSpinBox()
        self.seed.setRange(0, 2_147_483_647)
        self.count = QSpinBox()
        self.count.setRange(1, 16)
        self.count.setValue(4)
        self.rating = QSpinBox()
        self.rating.setRange(0, 5)
        self.rating.setSuffix(" / 5")
        self.notes = QPlainTextEdit()
        self.notes.setPlaceholderText("What worked, what drifted, and what to try next…")
        self.notes.setMaximumHeight(100)
        rows = (
            ("Model", self.model),
            ("Checkpoint", self.checkpoint),
            ("Evaluation prompt", self.prompt),
            ("Seed", self.seed),
            ("Preview count", self.count),
            ("Rating", self.rating),
            ("Notes", self.notes),
        )
        for row, (label, widget) in enumerate(rows, 1):
            form_layout.addWidget(QLabel(label), row, 0, Qt.AlignTop)
            form_layout.addWidget(widget, row, 1)
        save = QPushButton("Save evaluation")
        save.clicked.connect(self._save)
        request = QPushButton("Plan preview job")
        request.setProperty("primary", True)
        request.clicked.connect(self._request)
        form_layout.addWidget(save, len(rows) + 1, 0)
        form_layout.addWidget(request, len(rows) + 1, 1)
        root.addWidget(form, 2)

        history = _card()
        history_layout = QVBoxLayout(history)
        history_layout.addWidget(_title("MODEL OUTPUTS"))
        self.outputs = QListWidget()
        self.outputs.setViewMode(QListWidget.IconMode)
        self.outputs.setIconSize(QPixmap(110, 80).size())
        self.outputs.setMaximumHeight(210)
        self.outputs.itemDoubleClicked.connect(
            lambda item: QDesktopServices.openUrl(
                QUrl.fromLocalFile(str(item.data(Qt.UserRole)))
            )
        )
        history_layout.addWidget(self.outputs)
        history_layout.addWidget(_title("EVALUATION HISTORY"))
        self.history = QListWidget()
        history_layout.addWidget(self.history)
        root.addWidget(history, 1)
        self.model.currentIndexChanged.connect(self._reload_checkpoints)
        self.reload_assets()

    def reload_assets(self) -> None:
        signature = (
            tuple((asset.id, asset.name, asset.path, asset.checkpoint) for asset in self.assets.assets if asset.kind == "model"),
            tuple(sorted(self.store.best_models)),
        )
        if signature == getattr(self, "_model_choices_signature", None) and time.monotonic() - self._model_choices_refreshed_at < 30:
            return
        self._model_choices_signature = signature
        self._model_choices_refreshed_at = time.monotonic()
        current = self.model.currentData()
        self.model.blockSignals(True)
        self.model.clear()
        for asset in self.assets.assets:
            if asset.kind == "model":
                star = "★ " if asset.id in self.store.best_models else ""
                self.model.addItem(star + asset.name, asset.id)
        index = self.model.findData(current)
        self.model.setCurrentIndex(index if index >= 0 else 0)
        self.model.blockSignals(False)
        self._reload_checkpoints()
        self._reload_history()

    def _asset(self) -> Asset | None:
        model_id = str(self.model.currentData() or "")
        return next((asset for asset in self.assets.assets if asset.id == model_id), None)

    def _reload_checkpoints(self) -> None:
        self.checkpoint.clear()
        self.outputs.clear()
        self._output_token += 1
        token = self._output_token
        self._output_index = 0
        for worker in self._output_scan_workers:
            worker.requestInterruption()
        asset = self._asset()
        if not asset:
            return
        found = checkpoint_files(asset.path)
        if asset.checkpoint and Path(asset.checkpoint).exists():
            found = [Path(asset.checkpoint), *[path for path in found if str(path) != asset.checkpoint]]
        if not found:
            self.checkpoint.addItem("Latest model output", asset.path)
        else:
            for path in found:
                self.checkpoint.addItem(path.name, str(path))
        worker = ImageScanWorker(asset.path, token, limit=80)
        self._output_scan_workers.add(worker)
        worker.scanned.connect(self._preview_scan_finished)
        worker.finished.connect(
            lambda worker=worker: self._output_scan_workers.discard(worker)
        )
        worker.finished.connect(worker.deleteLater)
        worker.start()

    def _preview_scan_finished(self, paths: object, token: int) -> None:
        if token != self._output_token or not isinstance(paths, list):
            return
        self._output_paths = paths
        QTimer.singleShot(0, lambda: self._load_next_preview(token))

    def _load_next_preview(self, token: int) -> None:
        if token != self._output_token or self._output_index >= len(self._output_paths):
            return
        path = self._output_paths[self._output_index]
        item = QListWidgetItem(path.name)
        item.setData(Qt.UserRole, str(path))
        pixmap = _thumbnail(path, 110, 80)
        if not pixmap.isNull():
            item.setIcon(QIcon(pixmap))
        self.outputs.addItem(item)
        self._output_index += 1
        QTimer.singleShot(0, lambda: self._load_next_preview(token))

    def _save(self) -> None:
        asset = self._asset()
        if not asset:
            QMessageBox.warning(self, "No model", "Register or finish a model first.")
            return
        self.store.add_evaluation(
            PreviewEvaluation(
                model_id=asset.id,
                checkpoint=str(self.checkpoint.currentData() or ""),
                prompt=self.prompt.toPlainText().strip(),
                seed=self.seed.value(),
                rating=self.rating.value(),
                notes=self.notes.toPlainText().strip(),
            )
        )
        self._reload_history()

    def _reload_history(self) -> None:
        self.history.clear()
        names = {asset.id: asset.name for asset in self.assets.assets}
        for evaluation in reversed(self.store.evaluations[-100:]):
            self.history.addItem(
                f"{'★' * evaluation.rating or 'Unrated'} · "
                f"{names.get(evaluation.model_id, 'Unknown model')}\n"
                f"{evaluation.prompt or 'No prompt recorded'}"
            )

    def _request(self) -> None:
        asset = self._asset()
        if not asset:
            return
        prompt = self.prompt.toPlainText().strip()
        checkpoint = str(self.checkpoint.currentData() or "")
        request = (
            f"Generate {self.count.value()} previews for the {asset.name} model"
            + (f" from checkpoint {checkpoint}" if checkpoint else "")
            + (f" using this evaluation prompt: {prompt}" if prompt else "")
            + f". Use seed {self.seed.value()}."
        )
        self.plan_requested.emit(request)


class RecipesTab(QWidget):
    plan_requested = Signal(str)

    def __init__(self, store: StudioStore) -> None:
        super().__init__()
        self.store = store
        root = QHBoxLayout(self)
        root.setContentsMargins(10, 14, 10, 10)
        self.list = QListWidget()
        self.list.currentRowChanged.connect(self._selected)
        root.addWidget(self.list, 1)
        form = _card()
        layout = QGridLayout(form)
        layout.addWidget(_title("REPRODUCIBLE TRAINING RECIPE"), 0, 0, 1, 2)
        self.name = QLineEdit()
        self.trainer = QComboBox()
        self.trainer.addItem("LoRA", "lora")
        self.trainer.addItem("DDPM", "ddpm")
        self.trainer.addItem("Flow Matching", "flow")
        self.trainer.addItem("INRFlow", "inrflow")
        self.trainer.addItem("PixelRow", "pixelrow")
        self.epochs = QSpinBox()
        self.epochs.setRange(1, 100_000)
        self.epochs.setValue(100)
        self.images = QSpinBox()
        self.images.setRange(10, 100_000)
        self.images.setValue(60)
        self.base_model = QLineEdit()
        self.preview_prompt = QLineEdit()
        self.notes = QPlainTextEdit()
        self.notes.setMaximumHeight(100)
        for row, (label, widget) in enumerate(
            (
                ("Name", self.name),
                ("Trainer", self.trainer),
                ("Epochs", self.epochs),
                ("Dataset target", self.images),
                ("Base model", self.base_model),
                ("Preview prompt", self.preview_prompt),
                ("Notes", self.notes),
            ),
            1,
        ):
            layout.addWidget(QLabel(label), row, 0)
            layout.addWidget(widget, row, 1)
        save = QPushButton("Save recipe")
        save.clicked.connect(self._save)
        use = QPushButton("Create model from recipe")
        use.setProperty("primary", True)
        use.clicked.connect(self._use)
        layout.addWidget(save, 8, 0)
        layout.addWidget(use, 8, 1)
        transfer = QHBoxLayout()
        export = QPushButton("Export recipe…")
        import_button = QPushButton("Import recipe…")
        export.clicked.connect(self._export)
        import_button.clicked.connect(self._import)
        transfer.addWidget(import_button)
        transfer.addWidget(export)
        layout.addLayout(transfer, 9, 0, 1, 2)
        root.addWidget(form, 2)
        self.refresh()

    def refresh(self) -> None:
        row = self.list.currentRow()
        self.list.clear()
        for recipe in self.store.recipes:
            item = QListWidgetItem(
                f"{recipe.name}\n{recipe.trainer.upper()} · {recipe.epochs} epochs"
            )
            item.setData(Qt.UserRole, recipe.id)
            self.list.addItem(item)
        if self.list.count():
            self.list.setCurrentRow(max(0, min(row, self.list.count() - 1)))

    def _current(self) -> TrainingRecipe | None:
        item = self.list.currentItem()
        recipe_id = str(item.data(Qt.UserRole)) if item else ""
        return next((recipe for recipe in self.store.recipes if recipe.id == recipe_id), None)

    def _selected(self, _row: int) -> None:
        recipe = self._current()
        if not recipe:
            return
        self.name.setText(recipe.name)
        self.trainer.setCurrentIndex(max(0, self.trainer.findData(recipe.trainer)))
        self.epochs.setValue(recipe.epochs)
        self.images.setValue(recipe.image_count)
        self.base_model.setText(recipe.base_model)
        self.preview_prompt.setText(recipe.preview_prompt)
        self.notes.setPlainText(recipe.notes)

    def _save(self) -> None:
        current = self._current()
        recipe = TrainingRecipe(
            id=current.id if current else uuid4().hex[:10],
            created_at=current.created_at if current else datetime.now().astimezone().isoformat(),
            name=self.name.text().strip() or "Untitled recipe",
            trainer=str(self.trainer.currentData()),
            epochs=self.epochs.value(),
            image_count=self.images.value(),
            base_model=self.base_model.text().strip(),
            preview_prompt=self.preview_prompt.text().strip(),
            notes=self.notes.toPlainText().strip(),
        )
        self.store.add_recipe(recipe)
        self.refresh()

    def _use(self) -> None:
        trainer = str(self.trainer.currentData()).upper()
        self.plan_requested.emit(
            f"Create a {trainer} model using {self.images.value()} images and train "
            f"for {self.epochs.value()} epochs. Use the recipe named "
            f"{self.name.text().strip() or 'Untitled recipe'}."
        )

    def _export(self) -> None:
        recipe = self._current()
        if not recipe:
            QMessageBox.information(self, "No recipe", "Select or save a recipe first.")
            return
        selected, _filter = QFileDialog.getSaveFileName(
            self,
            "Export ADAM recipe",
            f"{recipe.name}.adam-recipe.json",
            "ADAM recipes (*.json)",
        )
        if selected:
            try:
                Path(selected).write_text(
                    json.dumps(
                        {"format": "adam-training-recipe-v1", "recipe": asdict(recipe)},
                        indent=2,
                    ),
                    encoding="utf-8",
                )
            except OSError as exc:
                QMessageBox.warning(self, "Recipe not exported", str(exc))

    def _import(self) -> None:
        selected, _filter = QFileDialog.getOpenFileName(
            self, "Import ADAM recipe", "", "ADAM recipes (*.json)"
        )
        if not selected:
            return
        try:
            payload = json.loads(Path(selected).read_text(encoding="utf-8"))
            if payload.get("format") != "adam-training-recipe-v1":
                raise ValueError("This is not an ADAM training recipe.")
            recipe = TrainingRecipe.from_dict(dict(payload["recipe"]))
            recipe.id = uuid4().hex[:10]
            self.store.add_recipe(recipe)
        except (OSError, ValueError, TypeError, KeyError, json.JSONDecodeError) as exc:
            QMessageBox.warning(self, "Recipe not imported", str(exc))
            return
        self.refresh()


class StudioPage(QWidget):
    plan_requested = Signal(str)

    def __init__(
        self,
        root_path: Path,
        jobs: JobManager,
        assets: AssetRegistry,
        config: ConfigManager,
    ) -> None:
        super().__init__()
        self.store = StudioStore(root_path)
        root = QVBoxLayout(self)
        root.setContentsMargins(24, 20, 24, 17)
        root.setSpacing(8)
        root.addWidget(
            _header(
                "Training studio",
                "Review datasets, compare experiments, evaluate checkpoints, and preserve reproducible recipes.",
            )
        )
        self.tabs = QTabWidget()
        self.datasets = DatasetReviewTab(assets, self.store, config)
        self.experiments = ExperimentsTab(jobs, assets, self.store)
        self.previews = PreviewLabTab(assets, self.store)
        self.recipes = RecipesTab(self.store)
        self.tabs.addTab(self.datasets, "Datasets")
        self.tabs.addTab(self.experiments, "Experiments")
        self.tabs.addTab(self.previews, "Checkpoint Lab")
        self.tabs.addTab(self.recipes, "Recipes")
        root.addWidget(self.tabs, 1)
        self.previews.plan_requested.connect(self.plan_requested)
        self.recipes.plan_requested.connect(self.plan_requested)
        self.tabs.currentChanged.connect(self._refresh_current_tab)

    def refresh(self) -> None:
        self._refresh_current_tab()

    def _refresh_current_tab(self, _index: int = 0) -> None:
        current = self.tabs.currentWidget()
        if current is self.datasets:
            self.datasets.reload_assets()
        elif current is self.experiments:
            self.experiments.refresh()
        elif current is self.previews:
            self.previews.reload_assets()
        elif current is self.recipes:
            self.recipes.refresh()

    def shutdown(self) -> None:
        workers = [
            *self.datasets._scan_workers,
            *self.previews._output_scan_workers,
        ]
        for worker in workers:
            worker.requestInterruption()
        for worker in workers:
            worker.wait(1000)