adam/ui/widgets.py

from __future__ import annotations

from dataclasses import asdict
from datetime import datetime, timezone
import math
from pathlib import Path
import re
from typing import Callable

from PySide6.QtCore import QPointF, QTimer, Qt, Signal
from PySide6.QtGui import QColor, QKeyEvent, QPainter, QPen, QPixmap
from PySide6.QtWidgets import (
    QFrame,
    QHBoxLayout,
    QLabel,
    QProgressBar,
    QPushButton,
    QSizePolicy,
    QTextEdit,
    QVBoxLayout,
    QWidget,
)

from adam.models import ExecutionPlan, Job, JobStatus, PlanStep, StepStatus, SystemSnapshot
from adam.ui.theme import COLORS


def clear_layout(layout: QVBoxLayout | QHBoxLayout) -> None:
    while layout.count():
        item = layout.takeAt(0)
        widget = item.widget()
        child_layout = item.layout()
        if widget:
            widget.deleteLater()
        elif child_layout:
            clear_layout(child_layout)  # type: ignore[arg-type]


class PromptEdit(QTextEdit):
    send_requested = Signal()

    def keyPressEvent(self, event: QKeyEvent) -> None:
        if event.key() in (Qt.Key_Return, Qt.Key_Enter) and not (
            event.modifiers() & Qt.ShiftModifier
        ):
            event.accept()
            self.send_requested.emit()
            return
        super().keyPressEvent(event)


class SparklineWidget(QWidget):
    """A small dependency-free chart for metrics found in trainer logs."""

    def __init__(self, empty_text: str = "Loss appears when reported") -> None:
        super().__init__()
        self.values: list[float] = []
        self.empty_text = empty_text
        self.setMinimumHeight(54)

    def set_values(self, values: list[float]) -> None:
        self.values = values[-100:]
        self.update()

    def paintEvent(self, _event) -> None:
        painter = QPainter(self)
        painter.setRenderHint(QPainter.Antialiasing)
        painter.fillRect(self.rect(), QColor("#050d14"))
        if len(self.values) < 2:
            painter.setPen(QColor(COLORS["muted"]))
            painter.drawText(self.rect(), Qt.AlignCenter, self.empty_text)
            return
        low, high = min(self.values), max(self.values)
        span = max(0.000001, high - low)
        width = max(1, self.width() - 12)
        height = max(1, self.height() - 12)
        points = [
            QPointF(
                6 + index / (len(self.values) - 1) * width,
                6 + (high - value) / span * height,
            )
            for index, value in enumerate(self.values)
        ]
        painter.setPen(QPen(QColor(COLORS["green"]), 2))
        for first, second in zip(points, points[1:]):
            painter.drawLine(first, second)


class ChatBubble(QFrame):
    def __init__(self, text: str, *, user: bool = False, label: str = "") -> None:
        super().__init__()
        self.setProperty("innerCard", True)
        self.setMaximumWidth(760)
        layout = QVBoxLayout(self)
        layout.setContentsMargins(16, 12, 16, 13)
        layout.setSpacing(6)

        self.title = QLabel(label or ("YOU" if user else "ADAM"))
        self.title.setObjectName("CardTitle")
        if user:
            self.title.setStyleSheet(f"color: {COLORS['green']};")
        self.message = QLabel(text)
        self.message.setWordWrap(True)
        self.message.setTextInteractionFlags(Qt.TextSelectableByMouse)
        self.message.setStyleSheet("font-size: 14px; line-height: 1.4;")
        layout.addWidget(self.title)
        layout.addWidget(self.message)
        if user:
            self.setStyleSheet(
                f"QFrame {{ background: #0b1c18; border-color: #1e543e; }}"
            )

    def set_text(self, text: str) -> None:
        self.message.setText(text)

    def set_label(self, text: str) -> None:
        self.title.setText(text)


class GenerationChatCard(QFrame):
    """Inline Command Center preview for one real image-generation job."""

    cancel_requested = Signal(str)
    open_requested = Signal(str)

    def __init__(self, job: Job, atom_path: Path) -> None:
        super().__init__()
        self.job_id = job.id
        self.image_path = ""
        self.image_paths: list[str] = []
        self.image_index = 0
        self._animation_frame = 0
        self._atom = QPixmap(str(atom_path))
        self.setProperty("innerCard", True)
        self.setMaximumWidth(900)
        self.setMinimumWidth(560)

        root = QVBoxLayout(self)
        root.setContentsMargins(16, 13, 16, 14)
        root.setSpacing(10)
        self.title = QLabel("ADAM · GENERATING IMAGE")
        self.title.setObjectName("CardTitle")
        root.addWidget(self.title)

        body = QHBoxLayout()
        body.setSpacing(16)
        preview_column = QVBoxLayout()
        preview_column.setSpacing(5)
        preview_row = QHBoxLayout()
        preview_row.setSpacing(6)
        self.previous_button = QPushButton("‹")
        self.previous_button.setFixedSize(30, 54)
        self.previous_button.setToolTip("Previous generated image")
        self.previous_button.clicked.connect(lambda: self._move_image(-1))
        self.previous_button.hide()
        preview_row.addWidget(self.previous_button, 0, Qt.AlignVCenter)
        self.preview = QLabel()
        self.preview.setAlignment(Qt.AlignCenter)
        self.preview.setFixedSize(300, 190)
        self.preview.setStyleSheet(
            "background: #050c13; border: 1px solid #183a50; border-radius: 8px;"
        )
        preview_row.addWidget(self.preview, 0, Qt.AlignTop)
        self.next_button = QPushButton("›")
        self.next_button.setFixedSize(30, 54)
        self.next_button.setToolTip("Next generated image")
        self.next_button.clicked.connect(lambda: self._move_image(1))
        self.next_button.hide()
        preview_row.addWidget(self.next_button, 0, Qt.AlignVCenter)
        preview_column.addLayout(preview_row)
        self.image_counter = QLabel()
        self.image_counter.setAlignment(Qt.AlignCenter)
        self.image_counter.setProperty("muted", True)
        self.image_counter.hide()
        preview_column.addWidget(self.image_counter)
        body.addLayout(preview_column)

        details = QVBoxLayout()
        details.setSpacing(5)
        arguments = job.plan.steps[0].arguments if job.plan.steps else {}
        prompt = str(arguments.get("prompt", "")).strip() or "Unconditional sample"
        prompt_label = QLabel(prompt)
        prompt_label.setWordWrap(True)
        prompt_label.setStyleSheet("font-size: 14px; font-weight: 600; padding-bottom: 5px;")
        details.addWidget(prompt_label)
        negative_prompt = str(arguments.get("negative_prompt", "")).strip()
        if negative_prompt:
            if len(negative_prompt) > 120:
                negative_prompt = negative_prompt[:117].rstrip() + "…"
            negative_label = QLabel(f"Negative: {negative_prompt}")
            negative_label.setWordWrap(True)
            negative_label.setProperty("muted", True)
            negative_label.setStyleSheet("font-size: 11px; padding-bottom: 4px;")
            details.addWidget(negative_label)
        base_model_path = str(arguments.get("base_model_path", "")).strip()
        lora_name = str(arguments.get("model_name", "")).strip()
        settings = []
        if base_model_path:
            settings.append(("Model", Path(base_model_path).stem))
            if str(arguments.get("model_path", "")).strip() and lora_name:
                settings.append(("LoRA", lora_name))
        else:
            settings.append(("Model", lora_name))
        settings.extend([
            ("Sampler", str(arguments.get("sampler", ""))),
            ("Steps", str(arguments.get("steps", ""))),
            ("Aspect ratio", str(arguments.get("aspect_ratio", ""))),
            ("Images", str(arguments.get("image_count", 1))),
            ("Seed", str(arguments.get("seed", 0) or "Random")),
        ])
        if "cfg_scale" in arguments:
            settings.append(("CFG scale", str(arguments.get("cfg_scale") or "Default")))
        if float(arguments.get("lora_strength", 0) or 0) > 0:
            settings.append(("LoRA strength", str(arguments.get("lora_strength"))))
        if str(arguments.get("reference_image", "")).strip():
            settings.append(("Reference", Path(str(arguments["reference_image"])).name))
        for name, value in settings:
            row = QHBoxLayout()
            key = QLabel(name)
            key.setProperty("muted", True)
            key.setMinimumWidth(88)
            val = QLabel(value)
            val.setTextInteractionFlags(Qt.TextSelectableByMouse)
            row.addWidget(key)
            row.addWidget(val, 1)
            details.addLayout(row)
        details.addStretch()
        body.addLayout(details, 1)
        root.addLayout(body)

        self.status = QLabel("Preparing generator…")
        self.status.setProperty("muted", True)
        root.addWidget(self.status)
        self.progress = QProgressBar()
        self.progress.setRange(0, 100)
        self.progress.setValue(0)
        root.addWidget(self.progress)
        actions = QHBoxLayout()
        self.cancel_button = QPushButton("Cancel generation")
        self.cancel_button.setProperty("danger", True)
        self.cancel_button.clicked.connect(lambda: self.cancel_requested.emit(self.job_id))
        self.open_button = QPushButton("Open image")
        self.open_button.setEnabled(False)
        self.open_button.clicked.connect(lambda: self.open_requested.emit(self.image_path))
        actions.addWidget(self.cancel_button)
        actions.addWidget(self.open_button)
        actions.addStretch()
        root.addLayout(actions)

        self._animation = QTimer(self)
        self._animation.setInterval(55)
        self._animation.timeout.connect(self._animate_atom)
        self._animation.start()
        self._animate_atom()
        self.update_job(job)

    def _animate_atom(self) -> None:
        if self._atom.isNull():
            self.preview.setText("GENERATING\n\nADAM is processing your request…")
            return
        self._animation_frame = (self._animation_frame + 1) % 120
        size = round(154 + 9 * (1 + math.sin(self._animation_frame / 120 * math.tau)) / 2)
        self.preview.setPixmap(self._atom.scaled(size, size, Qt.KeepAspectRatio, Qt.SmoothTransformation))

    def _finished_images(self, job: Job) -> list[Path]:
        if not job.output_folder:
            return []
        folder = Path(job.output_folder)
        if not folder.is_dir():
            return []
        candidates = [
            path for path in folder.iterdir()
            if path.is_file() and path.suffix.casefold() in {".png", ".jpg", ".jpeg", ".webp", ".bmp"}
            and job.id.casefold() in path.name.casefold()
        ]
        return sorted(candidates, key=lambda path: path.name.casefold())

    def _move_image(self, direction: int) -> None:
        if len(self.image_paths) < 2:
            return
        self.image_index = (self.image_index + direction) % len(self.image_paths)
        self._show_current_image()

    def _show_current_image(self) -> None:
        if not self.image_paths:
            return
        self.image_index = max(0, min(self.image_index, len(self.image_paths) - 1))
        self.image_path = self.image_paths[self.image_index]
        pixmap = QPixmap(self.image_path)
        if not pixmap.isNull():
            self.preview.setPixmap(
                pixmap.scaled(self.preview.size(), Qt.KeepAspectRatio, Qt.SmoothTransformation)
            )
        multiple = len(self.image_paths) > 1
        self.previous_button.setVisible(multiple)
        self.next_button.setVisible(multiple)
        self.image_counter.setVisible(multiple)
        self.image_counter.setText(f"Image {self.image_index + 1} of {len(self.image_paths)}")

    def update_job(self, job: Job) -> None:
        if job.id != self.job_id:
            return
        self.progress.setValue(job.progress)
        if job.status == JobStatus.FINISHED:
            self._animation.stop()
            results = self._finished_images(job)
            if results:
                self.image_paths = [str(result) for result in results]
                self.image_index = 0
                self._show_current_image()
                self.open_button.setEnabled(True)
            self.title.setText("ADAM · IMAGE CREATED")
            self.status.setText("Generation complete · saved to Generation History")
            self.cancel_button.hide()
        elif job.status == JobStatus.FAILED:
            self._animation.stop()
            self.preview.setText("Generation failed")
            self.title.setText("ADAM · GENERATION FAILED")
            self.status.setText(job.error or "The generator stopped unexpectedly.")
            self.cancel_button.hide()
        elif job.status in {JobStatus.CANCELLED, JobStatus.INTERRUPTED}:
            self._animation.stop()
            self.preview.setText("Generation cancelled")
            self.title.setText("ADAM · GENERATION CANCELLED")
            self.status.setText("No new image was created.")
            self.cancel_button.hide()
        elif job.status == JobStatus.AWAITING_CONFIRMATION:
            self.status.setText("Waiting for approval in Current Plan")
        elif job.status == JobStatus.QUEUED:
            self.status.setText("Queued · waiting for the generator")
        elif job.status == JobStatus.PAUSED:
            self.status.setText(f"Paused · {job.progress}%")
        else:
            latest = job.logs[-1] if job.logs else "Generating preview…"
            self.status.setText(latest)


class StepRow(QFrame):
    def __init__(self, index: int, step: PlanStep) -> None:
        super().__init__()
        self.index = index
        layout = QHBoxLayout(self)
        layout.setContentsMargins(0, 5, 0, 5)
        layout.setSpacing(10)
        self.marker = QLabel()
        self.marker.setAlignment(Qt.AlignCenter)
        self.marker.setFixedSize(24, 24)
        self.name = QLabel()
        self.name.setStyleSheet("font-weight: 600;")
        self.description = QLabel()
        self.description.setProperty("muted", True)
        self.description.setWordWrap(False)
        self.description.setStyleSheet("font-size: 11px;")
        text_layout = QVBoxLayout()
        text_layout.setSpacing(2)
        text_layout.addWidget(self.name)
        text_layout.addWidget(self.description)
        layout.addWidget(self.marker, 0, Qt.AlignTop)
        layout.addLayout(text_layout, 1)
        self.update_step(step)

    def update_step(self, step: PlanStep) -> None:
        state_color = {
            StepStatus.PENDING: COLORS["muted"],
            StepStatus.RUNNING: COLORS["blue_2"],
            StepStatus.FINISHED: COLORS["green"],
            StepStatus.FAILED: COLORS["red"],
            StepStatus.SKIPPED: COLORS["orange"],
        }[step.status]
        self.marker.setText("✓" if step.status == StepStatus.FINISHED else str(self.index + 1))
        self.marker.setStyleSheet(
            f"border: 1px solid {state_color}; color: {state_color}; "
            "border-radius: 12px; font-size: 11px; font-weight: 700;"
        )
        self.name.setText(step.title)
        description = step.description
        if len(description) > 64:
            description = description[:61].rstrip() + "…"
        self.description.setText(description)
        self.description.setToolTip(step.description)


class PlanPanel(QFrame):
    approved = Signal(str)
    rejected = Signal(str)

    def __init__(self) -> None:
        super().__init__()
        self.setProperty("card", True)
        self.job: Job | None = None
        self._step_rows: list[StepRow] = []
        root = QVBoxLayout(self)
        root.setContentsMargins(17, 16, 17, 16)
        root.setSpacing(10)
        title = QLabel("CURRENT PLAN")
        title.setObjectName("CardTitle")
        root.addWidget(title)
        self.project = QLabel("No plan selected")
        self.project.setStyleSheet("font-size: 17px; font-weight: 650;")
        root.addWidget(self.project)
        self.summary = QLabel(
            "Ask ADAM to collect a dataset, train a LoRA, generate previews, or inspect the GPU."
        )
        self.summary.setProperty("muted", True)
        self.summary.setWordWrap(True)
        root.addWidget(self.summary)

        self.orion = QLabel()
        self.orion.setWordWrap(True)
        self.orion.setStyleSheet(
            f"background: #111b2b; border: 1px solid {COLORS['blue']}; "
            f"border-radius: 7px; padding: 9px; color: {COLORS['blue_2']};"
        )
        self.orion.hide()
        root.addWidget(self.orion)

        self.steps_widget = QWidget()
        self.steps_layout = QVBoxLayout(self.steps_widget)
        self.steps_layout.setContentsMargins(0, 2, 0, 2)
        self.steps_layout.setSpacing(1)
        root.addWidget(self.steps_widget, 1)

        self.confirmation = QLabel()
        self.confirmation.setWordWrap(True)
        self.confirmation.setStyleSheet(
            f"background: #1c1425; border: 1px solid #51316c; "
            f"border-radius: 7px; padding: 9px; color: {COLORS['purple']};"
        )
        self.confirmation.hide()
        root.addWidget(self.confirmation)

        buttons = QHBoxLayout()
        self.cancel_button = QPushButton("Cancel")
        self.cancel_button.setProperty("danger", True)
        self.run_button = QPushButton("Approve and run")
        self.run_button.setProperty("primary", True)
        self.run_button.setStyleSheet(
            f"background-color: {COLORS['blue']}; color: #00101b; "
            f"border: 1px solid {COLORS['blue_2']}; font-weight: 700;"
        )
        self.cancel_button.clicked.connect(self._reject)
        self.run_button.clicked.connect(self._approve)
        buttons.addWidget(self.cancel_button)
        buttons.addWidget(self.run_button, 1)
        root.addLayout(buttons)
        self.cancel_button.hide()
        self.run_button.hide()

    def set_job(self, job: Job | None) -> None:
        previous_id = self.job.id if self.job else None
        self.job = job
        if job is None:
            clear_layout(self.steps_layout)
            self._step_rows = []
            self.project.setText("No plan selected")
            self.summary.setText(
                "Ask ADAM to collect a dataset, train a LoRA, generate previews, or inspect the GPU."
            )
            self.confirmation.hide()
            self.orion.hide()
            self.run_button.hide()
            self.cancel_button.hide()
            return
        plan = job.plan
        self.project.setText(plan.project_name)
        self.summary.setText(plan.summary.split("\n\nORION —", 1)[0])
        review = plan.orion_review
        if review:
            messages = [item.get("message", "") for item in review.get("findings", [])]
            self.orion.setText(
                f"ORION · {review.get('headline', 'PLAN REVIEW').upper()}\n"
                + "\n".join(f"• {message}" for message in messages if message)
            )
            self.orion.setVisible(True)
        else:
            self.orion.hide()
        if previous_id != job.id or len(self._step_rows) != len(plan.steps):
            clear_layout(self.steps_layout)
            self._step_rows = []
            for index, step in enumerate(plan.steps):
                row = StepRow(index, step)
                self._step_rows.append(row)
                self.steps_layout.addWidget(row)
            self.steps_layout.addStretch(1)
        else:
            for row, step in zip(self._step_rows, plan.steps):
                row.update_step(step)
        waiting = job.status == JobStatus.AWAITING_CONFIRMATION
        self.confirmation.setText("CONFIRMATION REQUIRED\n" + plan.confirmation_reason)
        self.confirmation.setVisible(waiting)
        self.run_button.setVisible(waiting)
        self.cancel_button.setVisible(waiting)

    def _approve(self) -> None:
        if self.job:
            self.approved.emit(self.job.id)

    def _reject(self) -> None:
        if self.job:
            self.rejected.emit(self.job.id)


class ActiveJobPanel(QFrame):
    pause_requested = Signal(str)
    resume_requested = Signal(str)
    cancel_requested = Signal(str)
    open_requested = Signal(str)

    def __init__(self) -> None:
        super().__init__()
        self.setProperty("card", True)
        # Training logs arrive continuously and their wrapped text can otherwise
        # change this card's preferred height on every refresh.  Keep the panel
        # at one reserved height so the rest of the command-center layout stays
        # still while a job is running.
        self.setFixedHeight(432)
        self.setSizePolicy(QSizePolicy.Preferred, QSizePolicy.Fixed)
        self.job: Job | None = None
        layout = QVBoxLayout(self)
        layout.setContentsMargins(17, 15, 17, 15)
        layout.setSpacing(8)
        title = QLabel("ACTIVE JOB")
        title.setObjectName("CardTitle")
        layout.addWidget(title)
        row = QHBoxLayout()
        self.name = QLabel("No active work")
        self.name.setStyleSheet("font-weight: 650;")
        self.status = QLabel("IDLE")
        self.status.setProperty("success", True)
        self.status.setStyleSheet("font-size: 11px; font-weight: 700;")
        row.addWidget(self.name, 1)
        row.addWidget(self.status)
        layout.addLayout(row)
        self.detail = QLabel("ADAM is ready.")
        self.detail.setProperty("muted", True)
        self.detail.setWordWrap(True)
        layout.addWidget(self.detail)
        self.atlas = QLabel("ATLAS · Standing by")
        self.atlas.setWordWrap(True)
        self.atlas.setProperty("muted", True)
        self.atlas.setStyleSheet("font-size: 10px;")
        layout.addWidget(self.atlas)
        self.preview_section = QWidget()
        preview_layout = QVBoxLayout(self.preview_section)
        preview_layout.setContentsMargins(0, 2, 0, 2)
        preview_layout.setSpacing(6)
        preview_header = QHBoxLayout()
        self.preview_title = QLabel("LIVE PREVIEW")
        self.preview_title.setObjectName("CardTitle")
        self.preview_next = QLabel("Waiting for first preview")
        self.preview_next.setProperty("muted", True)
        self.preview_next.setStyleSheet("font-size: 10px;")
        preview_header.addWidget(self.preview_title)
        preview_header.addStretch()
        preview_header.addWidget(self.preview_next)
        preview_layout.addLayout(preview_header)
        self.preview_image = QLabel("The first preview will appear here")
        self.preview_image.setAlignment(Qt.AlignCenter)
        self.preview_image.setFixedSize(256, 256)
        self.preview_image.setProperty("muted", True)
        self.preview_image.setStyleSheet(
            f"background: #050d14; border: 1px solid {COLORS['border_bright']}; "
            "border-radius: 8px;"
        )
        preview_layout.addWidget(self.preview_image, 0, Qt.AlignHCenter)
        self.open_preview_button = QPushButton("Open full size  ↗")
        self.open_preview_button.setEnabled(False)
        preview_layout.addWidget(self.open_preview_button)
        self.preview_meta = QLabel("Preview settings will appear when training starts.")
        self.preview_meta.setWordWrap(True)
        self.preview_meta.setProperty("muted", True)
        self.preview_meta.setStyleSheet("font-size: 10px;")
        preview_layout.addWidget(self.preview_meta)
        layout.addWidget(self.preview_section)
        self.preview_section.hide()
        self.progress = QProgressBar()
        self.progress.setRange(0, 100)
        self.progress.setValue(0)
        layout.addWidget(self.progress)
        self.progress_text = QLabel("0%")
        self.progress_text.setProperty("muted", True)
        layout.addWidget(self.progress_text)
        self.timing = QLabel("Elapsed — · ETA —")
        self.timing.setProperty("muted", True)
        layout.addWidget(self.timing)
        self.loss_chart = SparklineWidget()
        self.loss_chart.setMaximumHeight(46)
        layout.addWidget(self.loss_chart)
        self.recent_log = QLabel()
        self.recent_log.setWordWrap(True)
        self.recent_log.setProperty("muted", True)
        self.recent_log.setStyleSheet("font-size: 10px;")
        self.recent_log.setFixedHeight(132)
        self.recent_log.setSizePolicy(QSizePolicy.Preferred, QSizePolicy.Fixed)
        layout.addWidget(self.recent_log)
        self.loss_chart.hide()
        self.recent_log.hide()
        buttons = QHBoxLayout()
        self.pause_button = QPushButton("Pause")
        self.cancel_button = QPushButton("Stop")
        self.cancel_button.setProperty("danger", True)
        self.open_button = QPushButton("Open output")
        buttons.addWidget(self.pause_button)
        buttons.addWidget(self.cancel_button)
        buttons.addWidget(self.open_button, 1)
        layout.addLayout(buttons)
        self.pause_button.clicked.connect(self._pause_or_resume)
        self.cancel_button.clicked.connect(self._cancel)
        self.open_button.clicked.connect(self._open)
        self.open_preview_button.clicked.connect(self._open_preview)
        self.timer = QTimer(self)
        self.timer.timeout.connect(self._refresh_runtime)
        self.timer.start(1000)
        self._update_buttons()

    def set_job(self, job: Job | None) -> None:
        self.job = job
        if not job:
            self.name.setText("No active work")
            self.status.setText("IDLE")
            self.status.setStyleSheet(
                f"font-size: 11px; font-weight: 700; color: {COLORS['green']};"
            )
            self.detail.setText("ADAM is ready.")
            self.atlas.setText("ATLAS · Standing by")
            self.progress.setValue(0)
            self.progress_text.setText("0%")
            self.timing.setText("Elapsed — · ETA —")
            self.loss_chart.set_values([])
            self.loss_chart.hide()
            self.recent_log.clear()
            self.recent_log.hide()
            self.preview_section.hide()
            self.setFixedHeight(432)
            self._update_buttons()
            return
        self.name.setText(job.plan.project_name)
        self.loss_chart.show()
        self.recent_log.show()
        self.status.setText(job.status.value.upper())
        color = {
            JobStatus.RUNNING: COLORS["blue_2"],
            JobStatus.PAUSED: COLORS["orange"],
            JobStatus.FINISHED: COLORS["green"],
            JobStatus.FAILED: COLORS["red"],
            JobStatus.CANCELLED: COLORS["red"],
            JobStatus.INTERRUPTED: COLORS["orange"],
            JobStatus.AWAITING_CONFIRMATION: COLORS["purple"],
        }.get(job.status, COLORS["muted"])
        self.status.setStyleSheet(
            f"font-size: 11px; font-weight: 700; color: {color};"
        )
        if 0 <= job.current_step < len(job.plan.steps):
            active_step = job.plan.steps[job.current_step]
            self.detail.setText(active_step.title)
        elif job.logs:
            self.detail.setText(job.logs[-1].split("] ", 1)[-1])
        else:
            self.detail.setText(job.plan.summary)
        atlas = job.atlas_report
        self.atlas.setText(
            f"ATLAS · {str(atlas.get('severity', 'watching')).upper()} — "
            f"{atlas.get('message', 'Monitoring begins with the training step.')}"
        )
        self.progress.setValue(job.progress)
        self.progress_text.setText(f"{job.progress}% · Job {job.id}")
        losses: list[float] = []
        for line in job.logs:
            match = re.search(
                r"\bloss(?:\s*[:=]|\s+)(-?\d+(?:\.\d+)?(?:e[+-]?\d+)?)",
                line,
                re.IGNORECASE,
            )
            if match:
                try:
                    losses.append(float(match.group(1)))
                except ValueError:
                    pass
        self.loss_chart.set_values(losses)
        self.recent_log.setText("\n".join(job.logs[-3:]))
        self._update_preview(job)
        self._refresh_runtime()
        self.pause_button.setText("Resume" if job.status == JobStatus.PAUSED else "Pause")
        self._update_buttons()

    def _update_preview(self, job: Job) -> None:
        step = (
            job.plan.steps[job.current_step]
            if 0 <= job.current_step < len(job.plan.steps) else None
        )
        is_generation = bool(step and step.arguments.get("preview_interval", 0))
        self.preview_title.setText("LIVE DENOISING PREVIEW" if is_generation else "LIVE TRAINING PREVIEW")
        enabled = bool(step and (
            (step.tool_id.endswith("_trainer") and step.arguments.get("preview_enabled", False))
            or is_generation
        ))
        self.preview_section.setVisible(enabled)
        # A 256px image needs the room previously reserved for the loss chart and
        # recent log. Keeping both would force Qt to compress child layouts into
        # each other in the narrow command-center sidebar.
        self.loss_chart.setVisible(not enabled)
        self.recent_log.setVisible(not enabled)
        self.setFixedHeight(620 if enabled else 432)
        if not enabled or step is None:
            return
        interval = int(step.arguments.get("preview_interval" if is_generation else "preview_every", 5) or 5)
        epochs = int(step.arguments.get("epochs", 0) or 0)
        total = int(step.arguments.get("steps", 0) or 0) if is_generation else epochs
        unit = "step" if is_generation else "epoch"
        self.preview_next.setText(
            f"Latest: step {job.preview_current} of {job.preview_total or total}"
            if is_generation and job.preview_current else
            (f"Next at epoch {job.preview_next_epoch}"
             if job.preview_next_epoch else f"First at {unit} {min(interval, total) if total else interval}")
        )
        path = job.preview_path or ""
        pixmap = QPixmap(path) if path else QPixmap()
        if not pixmap.isNull():
            self.preview_image.setPixmap(
                pixmap.scaled(256, 256, Qt.KeepAspectRatio, Qt.SmoothTransformation)
            )
            self.preview_image.setText("")
            self.open_preview_button.setEnabled(True)
        else:
            self.preview_image.setPixmap(QPixmap())
            self.preview_image.setText("The first preview will appear here")
            self.open_preview_button.setEnabled(False)
        details = []
        if is_generation and job.preview_current:
            details.append(f"Step {job.preview_current} of {job.preview_total or total}")
            if job.preview_image_count:
                details.append(f"Image {job.preview_image_index} of {job.preview_image_count}")
        elif job.preview_epoch:
            details.append(f"Epoch {job.preview_epoch} of {epochs}")
        prompt = job.preview_prompt or str(step.arguments.get("preview_prompt", ""))
        if prompt:
            details.append(f'Prompt: “{prompt}”')
        seed = job.preview_seed if job.preview_seed is not None else step.arguments.get("preview_seed")
        if seed is not None:
            details.append(f"Seed: {seed}")
        steps = job.preview_steps or int(step.arguments.get("preview_steps", 0) or 0)
        if steps:
            details.append(f"Preview steps: {steps}")
        self.preview_meta.setText("  ·  ".join(details) or f"Preview every {interval} {unit}s")

    def _refresh_runtime(self) -> None:
        job = self.job
        if not job or not job.started_at:
            return
        try:
            start = datetime.fromisoformat(job.started_at)
            end = (
                datetime.fromisoformat(job.ended_at)
                if job.ended_at
                else datetime.now(timezone.utc)
            )
            elapsed = max(0, int((end - start).total_seconds()))
        except ValueError:
            return
        elapsed_text = (
            f"{elapsed // 3600}h {(elapsed % 3600) // 60}m"
            if elapsed >= 3600
            else f"{elapsed // 60}m {elapsed % 60}s"
        )
        eta = "—"
        if 0 < job.progress < 100:
            remaining = int(elapsed * (100 - job.progress) / job.progress)
            eta = (
                f"{remaining // 3600}h {(remaining % 3600) // 60}m"
                if remaining >= 3600
                else f"{remaining // 60}m {remaining % 60}s"
            )
        elif job.progress == 100:
            eta = "complete"
        self.timing.setText(f"Elapsed {elapsed_text} · ETA {eta}")

    def _update_buttons(self) -> None:
        running = bool(
            self.job and self.job.status in {JobStatus.RUNNING, JobStatus.PAUSED}
        )
        self.pause_button.setEnabled(running)
        self.cancel_button.setEnabled(running)
        self.open_button.setEnabled(bool(self.job and self.job.output_folder))

    def _pause_or_resume(self) -> None:
        if not self.job:
            return
        if self.job.status == JobStatus.PAUSED:
            self.resume_requested.emit(self.job.id)
        else:
            self.pause_requested.emit(self.job.id)

    def _cancel(self) -> None:
        if self.job:
            self.cancel_requested.emit(self.job.id)

    def _open(self) -> None:
        if self.job and self.job.output_folder:
            self.open_requested.emit(self.job.output_folder)

    def _open_preview(self) -> None:
        if self.job and self.job.preview_path:
            self.open_requested.emit(self.job.preview_path)


class MetricCard(QFrame):
    def __init__(self, label: str, accent: str = COLORS["blue"]) -> None:
        super().__init__()
        self.setProperty("card", True)
        self.accent = accent
        layout = QVBoxLayout(self)
        layout.setContentsMargins(16, 14, 16, 14)
        layout.setSpacing(6)
        self.label = QLabel(label.upper())
        self.label.setObjectName("CardTitle")
        self.value = QLabel("—")
        self.value.setStyleSheet("font-size: 24px; font-weight: 650;")
        self.detail = QLabel("Waiting for sensor data")
        self.detail.setProperty("muted", True)
        self.progress = QProgressBar()
        self.progress.setRange(0, 100)
        self.progress.setStyleSheet(
            f"QProgressBar::chunk {{ background: {self.accent}; }}"
        )
        layout.addWidget(self.label)
        layout.addWidget(self.value)
        layout.addWidget(self.detail)
        layout.addWidget(self.progress)

    def update_metric(self, value: str, detail: str, percent: float) -> None:
        self.value.setText(value)
        self.detail.setText(detail)
        self.progress.setValue(max(0, min(int(percent), 100)))