adam/ui/showcase.py

from __future__ import annotations

from pathlib import Path

from PySide6.QtCore import Qt, QUrl
from PySide6.QtGui import QDesktopServices
from PySide6.QtWidgets import (
    QAbstractItemView, QComboBox, QFrame, QGridLayout, QHBoxLayout, QLabel,
    QLineEdit, QListWidget, QListWidgetItem, QMessageBox, QPushButton,
    QSpinBox, QVBoxLayout, QWidget,
)

from adam.assets import AssetRegistry
from adam.config import ConfigManager
from adam.generations import build_generation_plan, generation_tools
from adam.job_manager import JobManager
from adam.models import Job, JobStatus
from adam.registry import ToolRegistry
from adam.showcase import build_showcase_plan
from adam.ui.asset_refresh import asset_refresh


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


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


class ShowcasePage(QWidget):
    """Create a finished, self-generated DDPM/Flow showcase video."""

    def __init__(self, root: Path, registry: ToolRegistry, jobs: JobManager, assets: AssetRegistry, config: ConfigManager) -> None:
        super().__init__()
        self.root = root.resolve(); self.registry = registry; self.jobs = jobs; self.assets = assets; self.config = config
        self.latest_video: Path | None = None
        layout = QVBoxLayout(self); layout.setContentsMargins(22, 18, 22, 18); layout.setSpacing(12)
        heading = QLabel("Showcase Video"); heading.setObjectName("PageTitle")
        subtitle = QLabel("Generate 12–24 images per DDPM or Flow Matching model, then export the animated request interface as a finished MP4.")
        subtitle.setProperty("muted", True); subtitle.setWordWrap(True)
        layout.addWidget(heading); layout.addWidget(subtitle)
        body = QHBoxLayout(); body.setSpacing(12)
        body.addWidget(self._build_models(), 1); body.addWidget(self._build_settings(), 1)
        layout.addLayout(body, 1)
        self.jobs.job_updated.connect(self._job_updated)
        self._asset_refresh = asset_refresh(self.assets, self.config)
        self._asset_refresh.changed.connect(self._render_models)
        self.refresh(); self._restore()

    def _build_models(self) -> QFrame:
        card = _card(); layout = QVBoxLayout(card); layout.setContentsMargins(16, 15, 16, 16)
        top = QHBoxLayout(); top.addWidget(_title("REQUEST / MODEL ORDER")); top.addStretch()
        refresh = QPushButton("Refresh"); refresh.clicked.connect(lambda: self.refresh(force=True)); top.addWidget(refresh); layout.addLayout(top)
        note = QLabel("Select one or more completed models. Drag entries to set their order in the video. LoRAs are intentionally excluded.")
        note.setProperty("muted", True); note.setWordWrap(True); layout.addWidget(note)
        self.models = QListWidget(); self.models.setSelectionMode(QAbstractItemView.MultiSelection); self.models.setDragDropMode(QAbstractItemView.InternalMove); self.models.setDefaultDropAction(Qt.MoveAction); self.models.setSpacing(5)
        layout.addWidget(self.models, 1)
        buttons = QHBoxLayout(); select_all = QPushButton("Select all"); clear = QPushButton("Clear")
        select_all.clicked.connect(self.models.selectAll); clear.clicked.connect(self.models.clearSelection)
        buttons.addWidget(select_all); buttons.addWidget(clear); buttons.addStretch(); layout.addLayout(buttons)
        return card

    def _build_settings(self) -> QFrame:
        card = _card(); layout = QVBoxLayout(card); layout.setContentsMargins(16, 15, 16, 16); layout.setSpacing(10)
        layout.addWidget(_title("SHOWCASE SETTINGS"))
        grid = QGridLayout(); grid.setVerticalSpacing(8)
        self.video_title = QLineEdit("ADAM Generation Showcase")
        self.images = QSpinBox(); self.images.setRange(12, 24); self.images.setValue(18); self.images.setSuffix(" per model")
        self.duration = QComboBox(); self.duration.addItems(["3 seconds", "4 seconds", "5 seconds"]); self.duration.setCurrentIndex(1)
        self.steps = QSpinBox(); self.steps.setRange(5, 200); self.steps.setValue(30)
        self.ddpm_sampler = QComboBox(); self.ddpm_sampler.addItems(["DDIM", "DDPM"])
        self.flow_sampler = QComboBox(); self.flow_sampler.addItems(["Heun", "Euler"])
        self.aspect = QComboBox(); self.aspect.addItems(["16:9 (Widescreen)", "1:1 (Square)"])
        self.seed = QSpinBox(); self.seed.setRange(1, 2_146_000_000); self.seed.setValue(123456)
        self.resolution = QComboBox(); self.resolution.addItems(["1080p", "720p"])
        fields = [
            ("Video title", self.video_title), ("Images", self.images), ("Each image lasts", self.duration),
            ("Generation steps", self.steps), ("DDPM sampler", self.ddpm_sampler),
            ("Flow method", self.flow_sampler), ("Image aspect", self.aspect),
            ("Starting seed", self.seed), ("Video resolution", self.resolution),
        ]
        for row, (label, widget) in enumerate(fields): grid.addWidget(QLabel(label), row, 0); grid.addWidget(widget, row, 1)
        layout.addLayout(grid)
        self.estimate = QLabel(); self.estimate.setProperty("muted", True); self.estimate.setWordWrap(True); layout.addWidget(self.estimate)
        self.create_button = QPushButton("Generate images and create MP4  →"); self.create_button.setProperty("primary", True); self.create_button.clicked.connect(self._create)
        layout.addWidget(self.create_button)
        self.open_button = QPushButton("Open latest showcase video"); self.open_button.setEnabled(False); self.open_button.clicked.connect(self._open_latest); layout.addWidget(self.open_button)
        self.status = QLabel("Ready"); self.status.setProperty("muted", True); self.status.setWordWrap(True); layout.addWidget(self.status)
        layout.addStretch()
        self.models.itemSelectionChanged.connect(self._update_estimate)
        self.images.valueChanged.connect(self._update_estimate); self.duration.currentIndexChanged.connect(self._update_estimate)
        return card

    @staticmethod
    def _ready(path: Path, trainer: str) -> bool:
        if trainer == "ddpm": return path.is_dir() and (path / "model_index.json").is_file()
        if trainer == "flow": return path.is_dir() and (path / "flow_model_info.json").is_file() and (path / "unet" / "config.json").is_file()
        return False

    def refresh(self, *, force: bool = False) -> None:
        self._asset_refresh.request(force=force)
        self._render_models()

    def _render_models(self) -> None:
        selected = {str(item.data(Qt.UserRole).get("path")) for item in self.models.selectedItems()} if hasattr(self, "models") else set()
        self.models.clear()
        tools = generation_tools(self.registry)
        for asset in self.assets.assets:
            if asset.kind != "model" or asset.trainer not in {"ddpm", "flow"} or not self._ready(Path(asset.path), asset.trainer): continue
            tool = next((value for value in tools if asset.trainer in value.model_trainers), None)
            if not tool: continue
            label = "DDPM" if asset.trainer == "ddpm" else "Flow Matching"
            item = QListWidgetItem(f"{asset.name}  —  {label}")
            item.setData(Qt.UserRole, {"name": asset.name, "path": asset.path, "trainer": asset.trainer, "trainer_label": label, "tool_id": tool.id})
            self.models.addItem(item); item.setSelected(asset.path in selected)
        self.create_button.setEnabled(self.models.count() > 0); self._update_estimate()

    def _duration_seconds(self) -> int:
        return int(self.duration.currentText().split()[0])

    def _update_estimate(self) -> None:
        count = len(self.models.selectedItems()) if hasattr(self, "models") else 0
        images = count * self.images.value() if hasattr(self, "images") else 0
        seconds = images * self._duration_seconds() if images else 0
        self.estimate.setText(f"{images} images total · finished video about {seconds // 60}:{seconds % 60:02d}" if count else "Select at least one model to calculate the video length.")

    def _selected(self) -> list[dict]:
        return [dict(self.models.item(index).data(Qt.UserRole)) for index in range(self.models.count()) if self.models.item(index).isSelected()]

    def _create(self) -> None:
        selected = self._selected()
        if not selected:
            QMessageBox.information(self, "Choose models", "Select at least one completed DDPM or Flow Matching model."); return
        count = self.images.value(); plans = []; settings = []
        for index, entry in enumerate(selected):
            tool = self.registry.get(entry["tool_id"]); trainer = entry["trainer"]
            sampler = self.ddpm_sampler.currentText() if trainer == "ddpm" else self.flow_sampler.currentText()
            plans.append(build_generation_plan(
                tool, model_name=entry["name"], model_path=entry["path"], prompt="",
                image_count=count, steps=self.steps.value(), seed=self.seed.value() + index * count,
                sampler=sampler, aspect_ratio=self.aspect.currentText(),
            ))
            settings.append({**entry, "steps": self.steps.value(), "sampler": sampler, "aspect_ratio": self.aspect.currentText()})
        plan = build_showcase_plan(
            plans, title=self.video_title.text(), display_seconds=self._duration_seconds(),
            resolution=self.resolution.currentText(), model_settings=settings,
        )
        job = self.jobs.submit(plan); self._save(); self.status.setText(f"Showcase job {job.id} queued. ADAM will generate every image before rendering the MP4.")

    def _save(self) -> None:
        self.config.update({"showcase_settings": {"title": self.video_title.text(), "images": self.images.value(), "duration": self._duration_seconds(), "steps": self.steps.value(), "ddpm_sampler": self.ddpm_sampler.currentText(), "flow_sampler": self.flow_sampler.currentText(), "aspect": self.aspect.currentText(), "seed": self.seed.value(), "resolution": self.resolution.currentText()}})

    def _restore(self) -> None:
        saved = self.config.get("showcase_settings", {})
        if not isinstance(saved, dict): return
        self.video_title.setText(str(saved.get("title", self.video_title.text())))
        for spin, key in ((self.images, "images"), (self.steps, "steps"), (self.seed, "seed")):
            try: spin.setValue(max(spin.minimum(), min(int(saved.get(key, spin.value())), spin.maximum())))
            except (TypeError, ValueError): pass
        duration = f"{saved.get('duration', 4)} seconds"
        for combo, value in ((self.duration, duration), (self.ddpm_sampler, saved.get("ddpm_sampler")), (self.flow_sampler, saved.get("flow_sampler")), (self.aspect, saved.get("aspect")), (self.resolution, saved.get("resolution"))):
            if value and combo.findText(str(value)) >= 0: combo.setCurrentText(str(value))
        self._update_estimate()

    def _job_updated(self, job: Job) -> None:
        if job.plan.project_name != "Showcase Video": return
        if job.status == JobStatus.FINISHED:
            videos = sorted((self.root / "data" / "showcase_videos").glob(f"*{job.id}.mp4"))
            self.latest_video = videos[-1] if videos else None; self.open_button.setEnabled(self.latest_video is not None)
            self.status.setText(f"Showcase {job.id} finished successfully." if self.latest_video else f"Showcase {job.id} finished, but its MP4 could not be located.")
        elif job.status == JobStatus.FAILED: self.status.setText(f"Showcase {job.id} failed: {job.error or 'Unknown error'}")
        elif job.status == JobStatus.CANCELLED: self.status.setText(f"Showcase {job.id} was cancelled.")
        else: self.status.setText(f"Showcase {job.id}: {job.status.value} · {job.progress}%")

    def _open_latest(self) -> None:
        if self.latest_video and self.latest_video.is_file(): QDesktopServices.openUrl(QUrl.fromLocalFile(str(self.latest_video)))