adam/job_manager.py

from __future__ import annotations

import json
import logging
import math
import re
import threading
import time
from datetime import datetime, timedelta, timezone
from pathlib import Path
from typing import Any

from PySide6.QtCore import QCoreApplication, QObject, QThread, QTimer, Signal

from adam.executor import ToolAdjustmentRequested, ToolCancelled, ToolExecutionError, ToolExecutor
from adam.assets import AssetRegistry
from adam.atlas import AtlasSupervisor
from adam.experiment_tracker import ExperimentStore
from adam.models import ExecutionPlan, Job, JobStatus, StepStatus, utc_now
from adam.nova import evaluate_job_output
from adam.training_assistant import append_preflight_summary


def _safe_int(value: Any) -> int:
    try:
        return max(0, int(value or 0))
    except (TypeError, ValueError):
        return 0


class JobWorker(QThread):
    event = Signal(object)

    def __init__(self, job: Job, executor: ToolExecutor) -> None:
        super().__init__()
        self.job = job
        self.executor = executor
        self.cancel_event = threading.Event()
        self.run_event = threading.Event()
        self.run_event.set()
        self.adjustment_event = threading.Event()
        self.adjustment_request: dict[str, Any] = {}

    def pause(self) -> None:
        self.run_event.clear()

    def resume(self) -> None:
        self.run_event.set()

    def cancel(self) -> None:
        self.cancel_event.set()
        self.run_event.set()

    def request_adjustment(self, updates: dict[str, Any]) -> None:
        self.adjustment_request.clear()
        self.adjustment_request.update(updates)
        self.adjustment_event.set()
        self.run_event.set()

    def run(self) -> None:
        total_steps = len(self.job.plan.steps)
        try:
            for index, step in enumerate(self.job.plan.steps):
                if self.cancel_event.is_set():
                    raise ToolCancelled("Job cancelled by user.")
                self.event.emit(
                    {
                        "type": "step_started",
                        "index": index,
                        "message": step.title,
                    }
                )

                preview_state = {"epoch": 0, "path": ""}
                last_progress_emit = {"time": 0.0, "overall": -1, "message": ""}
                progress_samples: list[dict[str, Any]] = []

                def on_progress(percent: int, message: str, step_index: int = index, **details: Any) -> None:
                    overall = int(((step_index + percent / 100) / total_steps) * 100)
                    now = time.monotonic()
                    progress_eta = self._estimate_step_eta(details, progress_samples, now)
                    changed = overall != last_progress_emit["overall"] or message != last_progress_emit["message"]
                    terminal = percent >= 100 or overall >= 100
                    if not terminal and (not changed or now - last_progress_emit["time"] < 0.25):
                        return
                    last_progress_emit.update({"time": now, "overall": overall, "message": message})
                    self.event.emit(
                        {
                            "type": "progress",
                            "step_percent": percent,
                            "overall": overall,
                            "message": message,
                            **progress_eta,
                        }
                    )
                    self._discover_external_preview(step, message, preview_state)

                def on_log(message: str) -> None:
                    self.event.emit({"type": "log", "message": message})

                def on_preview(preview: dict[str, Any]) -> None:
                    preview_state["epoch"] = int(preview.get("epoch", 0) or 0)
                    preview_state["path"] = str(preview.get("path", ""))
                    self.event.emit({"type": "preview", **preview})

                result = self.executor.execute(
                    step.tool_id,
                    step.arguments,
                    job_id=self.job.id,
                    cancel_event=self.cancel_event,
                    run_event=self.run_event,
                    progress_callback=on_progress,
                    log_callback=on_log,
                    preview_callback=on_preview,
                    adjustment_event=self.adjustment_event,
                    adjustment_request=self.adjustment_request,
                )
                self.event.emit(
                    {
                        "type": "step_finished",
                        "index": index,
                        "result": result,
                    }
                )
            self.event.emit({"type": "completed"})
        except ToolCancelled as exc:
            self.event.emit({"type": "cancelled", "message": str(exc)})
        except ToolAdjustmentRequested as exc:
            self.event.emit({"type": "adjustment_ready", "message": str(exc), **exc.details})
        except Exception as exc:
            self.event.emit(
                {
                    "type": "failed",
                    "message": str(exc),
                    "exception": type(exc).__name__,
                }
            )

    def _discover_external_preview(
        self, step: Any, message: str, state: dict[str, Any]
    ) -> None:
        """Discover conventional preview files from any registered trainer."""
        if not step.tool_id.endswith("_trainer") or not step.arguments.get("preview_enabled", False):
            return
        match = re.search(r"\bepoch\s+(\d+)(?:\s+(?:of|/)|/)?", message, re.I)
        if not match:
            return
        epoch = int(match.group(1))
        interval = max(1, int(step.arguments.get("preview_every", 5) or 5))
        if epoch % interval or epoch == int(state.get("epoch", 0)):
            return
        output = Path(str(step.arguments.get("output_dir", ""))).expanduser()
        if not output.is_dir():
            return
        try:
            candidates = [
                path for path in output.rglob("*")
                if path.is_file() and path.suffix.lower() in {".png", ".jpg", ".jpeg", ".webp", ".bmp"}
                and any(token in path.name.lower() for token in ("preview", "sample", "epoch"))
            ]
            latest = max(candidates, key=lambda path: path.stat().st_mtime) if candidates else None
        except OSError:
            latest = None
        if latest and str(latest) != state.get("path"):
            state["epoch"] = epoch
            state["path"] = str(latest)
            self.event.emit({
                "type": "preview", "path": str(latest), "epoch": epoch,
                "next_epoch": min(int(step.arguments.get("epochs", epoch + interval)), epoch + interval),
                "prompt": str(step.arguments.get("preview_prompt", "")),
                "seed": step.arguments.get("preview_seed"),
                "steps": int(step.arguments.get("preview_steps", 0) or 0),
            })

    @staticmethod
    def _estimate_step_eta(
        details: dict[str, Any],
        samples: list[dict[str, Any]],
        now: float,
    ) -> dict[str, Any]:
        """Estimate remaining runtime from real step cadence instead of percent alone."""
        current = _safe_int(details.get("current_step", details.get("step", details.get("current"))))
        total = _safe_int(details.get("total_steps", details.get("total")))
        unit = str(details.get("unit", "step") or "step")
        epoch = _safe_int(details.get("epoch"))
        total_epochs = _safe_int(details.get("total_epochs"))
        if details.get("reset_eta"):
            samples.clear()
        if (not current or not total) and epoch and total_epochs:
            current, total, unit = epoch, total_epochs, "epoch"
        payload: dict[str, Any] = {
            "progress_current": current,
            "progress_total": total,
            "progress_unit": unit,
        }
        if not current or not total or current >= total:
            return payload
        last = samples[-1] if samples else None
        if last and (
            int(last.get("total", total)) != total
            or str(last.get("unit", unit)) != unit
            or current < last["current"]
        ):
            # A different tqdm operation or a restarted counter needs a fresh
            # cadence; carrying the old rate creates wildly incorrect ETAs.
            samples.clear()
            last = None
        if last and current <= last["current"]:
            return payload
        samples.append({"time": now, "current": float(current), "total": float(total), "unit": unit})
        del samples[:-25]
        if len(samples) < 2:
            return payload

        first = samples[0]
        elapsed = now - first["time"]
        completed = current - int(first["current"])
        if completed <= 0 or elapsed <= 0:
            return payload
        lifetime_seconds_per_unit = elapsed / completed

        recent = samples[-8:]
        recent_first = recent[0]
        recent_completed = current - int(recent_first["current"])
        recent_elapsed = now - recent_first["time"]
        recent_seconds_per_unit = (
            recent_elapsed / recent_completed
            if recent_completed > 0 and recent_elapsed > 0
            else lifetime_seconds_per_unit
        )
        seconds_per_unit = (recent_seconds_per_unit * 0.65) + (lifetime_seconds_per_unit * 0.35)
        remaining = max(0, int(round((total - current) * seconds_per_unit)))
        if remaining:
            payload["eta_seconds"] = remaining
            payload["progress_rate"] = 1 / seconds_per_unit if seconds_per_unit > 0 else 0.0
            payload["estimated_completion_at"] = (
                datetime.now(timezone.utc) + timedelta(seconds=remaining)
            ).isoformat()
        return payload


class JobManager(QObject):
    job_created = Signal(object)
    job_updated = Signal(object)
    log_added = Signal(str, str)
    active_changed = Signal(object)
    notification = Signal(str, str)

    def __init__(
        self,
        root: Path,
        executor: ToolExecutor,
        logger: logging.Logger,
        config: Any | None = None,
    ) -> None:
        super().__init__()
        self.root = root
        self.executor = executor
        self.logger = logger
        self.config = config if config is not None else {}
        self.jobs_path = root / "data" / "jobs.json"
        self.assets = AssetRegistry(root)
        self.atlas = AtlasSupervisor(config)
        self.experiments = ExperimentStore(root)
        self.jobs: list[Job] = []
        self._queue: list[str] = []
        self._worker: JobWorker | None = None
        self._active_job: Job | None = None
        self._last_snapshot = None
        self._pending_update_job_ids: set[str] = set()
        self._pending_update_timer_active = False
        self._load()
        self._schedule_timer = QTimer(self)
        self._schedule_timer.timeout.connect(self._release_due_scheduled)
        if QCoreApplication.instance() is not None:
            self._schedule_timer.start(15_000)
            QTimer.singleShot(0, self._release_due_scheduled)
            if self._queue:
                QTimer.singleShot(0, self._start_next)

    @property
    def active_job(self) -> Job | None:
        return self._active_job

    @property
    def pending_jobs(self) -> list[Job]:
        return [
            job for job in self.jobs
            if job.status in {JobStatus.AWAITING_CONFIRMATION, JobStatus.SCHEDULED, JobStatus.QUEUED}
        ]

    def submit(self, plan: ExecutionPlan, scheduled_for: str | None = None) -> Job:
        # Every entry point must review a plan before its queue state is chosen.
        # UI and Remote may prepare it earlier to keep filesystem work off Qt.
        append_preflight_summary(plan, self.config)
        is_future = self._is_future(scheduled_for)
        status = (
            JobStatus.AWAITING_CONFIRMATION
            if plan.requires_confirmation
            else JobStatus.SCHEDULED if is_future else JobStatus.QUEUED
        )
        job = Job(plan=plan, status=status, scheduled_for=scheduled_for if is_future else None)
        self.jobs.insert(0, job)
        self._append_log(job, f"Plan created: {plan.summary}")
        if plan.requires_confirmation:
            self._append_log(job, "Waiting for user confirmation.")
            if is_future:
                self._append_log(job, f"Requested start time: {self._display_time(scheduled_for)}.")
        elif is_future:
            self._append_log(job, f"Scheduled for {self._display_time(scheduled_for)}.")
        else:
            self._queue.append(job.id)
        self._save()
        self.job_created.emit(job)
        self.job_updated.emit(job)
        if not plan.requires_confirmation:
            self._start_next()
        return job

    def confirm(self, job_id: str) -> None:
        job = self.get(job_id)
        if job.status != JobStatus.AWAITING_CONFIRMATION:
            return
        job.status = JobStatus.SCHEDULED if self._is_future(job.scheduled_for) else JobStatus.QUEUED
        self._append_log(job, "Plan approved by user.")
        if job.status == JobStatus.SCHEDULED:
            self._append_log(job, f"Training will become eligible at {self._display_time(job.scheduled_for)}.")
        else:
            self._queue.append(job.id)
        self._save()
        self.job_updated.emit(job)
        self._start_next()

    def reject(self, job_id: str) -> None:
        job = self.get(job_id)
        if job.status != JobStatus.AWAITING_CONFIRMATION:
            return
        job.status = JobStatus.CANCELLED
        job.ended_at = utc_now()
        self._append_log(job, "Plan cancelled before execution.")
        self._save()
        self.job_updated.emit(job)

    def pause(self, job_id: str) -> None:
        job = self.get(job_id)
        if job is self._active_job and job.status == JobStatus.RUNNING and self._worker:
            self._worker.pause()
            job.status = JobStatus.PAUSED
            self._append_log(job, "Job paused.")
            self._save()
            self.job_updated.emit(job)

    def resume(self, job_id: str) -> None:
        job = self.get(job_id)
        if job is self._active_job and job.status == JobStatus.PAUSED and self._worker:
            self._worker.resume()
            job.status = JobStatus.RUNNING
            self._append_log(job, "Job resumed.")
            self._save()
            self.job_updated.emit(job)

    def cancel(self, job_id: str) -> None:
        job = self.get(job_id)
        if job is self._active_job and self._worker:
            self._append_log(job, "Cancellation requested.")
            self._worker.cancel()
            self._save()
            self.job_updated.emit(job)
            return
        if job.id in self._queue:
            self._queue.remove(job.id)
        if job.status in {
            JobStatus.SCHEDULED,
            JobStatus.QUEUED,
            JobStatus.AWAITING_CONFIRMATION,
            JobStatus.DRAFT,
        }:
            job.status = JobStatus.CANCELLED
            job.ended_at = utc_now()
            self._append_log(job, "Job cancelled.")
            self._save()
            self.job_updated.emit(job)

    def request_training_adjustment(self, job_id: str, updates: dict[str, Any]) -> None:
        """Apply safe DDPM settings after the current epoch and resume automatically."""
        job = self.get(job_id)
        if job is not self._active_job or job.status not in {JobStatus.RUNNING, JobStatus.PAUSED} or not self._worker:
            raise ValueError("Only the active training job can be adjusted.")
        if not (0 <= job.current_step < len(job.plan.steps)):
            raise ValueError("The active training step is unavailable.")
        step = job.plan.steps[job.current_step]
        if step.tool_id != "ddpm_trainer":
            raise ValueError("Safe epoch-boundary adjustment currently supports DDPM training.")
        if step.arguments.get("progressive_stages"):
            raise ValueError(
                "Change batch settings before starting a progressive run; each stage manages its own saved handoff."
            )
        allowed = {"batch_size", "training_intensity", "gradient_accumulation_steps"}
        cleaned = {key: int(value) for key, value in updates.items() if key in allowed}
        if not cleaned or not 1 <= cleaned.get("batch_size", 1) <= 64 \
                or not 10 <= cleaned.get("training_intensity", 100) <= 100 \
                or not 1 <= cleaned.get("gradient_accumulation_steps", 1) <= 64:
            raise ValueError("The requested training settings are outside ADAM's safe range.")
        previous = {key: step.arguments.get(key) for key in cleaned}
        if all(previous[key] == value for key, value in cleaned.items()):
            raise ValueError("Those settings are already active.")
        job.status = JobStatus.RUNNING
        self._append_log(job, f"Adjustment queued for the end of this epoch: {cleaned}.")
        self._worker.request_adjustment(cleaned)
        self._save()
        self.job_updated.emit(job)

    def safer_vram_retry(self, job_id: str) -> Job:
        """Create a checkpoint-aware DDPM retry with a smaller physical batch."""
        original = self.get(job_id)
        if original.status != JobStatus.FAILED or not self._looks_like_vram_failure(original):
            raise ValueError("This job did not fail with a recognizable VRAM error.")
        plan = ExecutionPlan.from_dict(original.to_dict()["plan"])
        start_index = max(0, min(original.current_step, len(plan.steps) - 1))
        plan.steps = plan.steps[start_index:]
        step = plan.steps[0]
        old_batch = max(1, int(step.arguments.get("batch_size", 1)))
        if old_batch <= 1:
            raise ValueError("Batch size is already 1; lower resolution or enable other memory-saving options.")
        original_epochs = max(1, int(step.arguments.get("epochs", 1)))
        resume_note = self._prepare_ddpm_resume(step.arguments, step.tool_id)
        remaining_epochs = max(1, int(step.arguments.get("epochs", original_epochs)))
        completed_epochs = max(0, original_epochs - remaining_epochs) if resume_note else 0
        new_batch = max(1, old_batch // 2)
        old_accumulation = max(1, int(step.arguments.get("gradient_accumulation_steps", 1)))
        step.arguments["batch_size"] = new_batch
        step.arguments["gradient_accumulation_steps"] = min(64, old_accumulation * max(1, math.ceil(old_batch / new_batch)))
        if completed_epochs:
            step.arguments["completed_epochs"] = completed_epochs
        for item in plan.steps:
            item.status = StepStatus.PENDING
        plan.id = original.plan.id + "-vram-retry"
        plan.created_at = utc_now()
        plan.requires_confirmation = True
        plan.confirmation_reason = "VRAM recovery reduced the physical batch and preserved the effective batch with gradient accumulation."
        retry = self.submit(plan)
        self._append_log(retry, f"VRAM recovery changed batch {old_batch} → {new_batch} and gradient accumulation {old_accumulation} → {step.arguments['gradient_accumulation_steps']}.")
        if resume_note:
            self._append_log(retry, resume_note)
        return retry

    @staticmethod
    def _looks_like_vram_failure(job: Job) -> bool:
        text = "\n".join([job.error or "", *job.logs[-100:]]).lower()
        return any(token in text for token in ("out of memory", "cuda oom", "cuda error: out of memory"))

    @staticmethod
    def _is_future(value: str | None) -> bool:
        if not value:
            return False
        try:
            scheduled = datetime.fromisoformat(value)
            if scheduled.tzinfo is None:
                scheduled = scheduled.astimezone()
            return scheduled.astimezone(timezone.utc) > datetime.now(timezone.utc)
        except (TypeError, ValueError):
            return False

    @staticmethod
    def _display_time(value: str | None) -> str:
        try:
            return datetime.fromisoformat(str(value)).astimezone().strftime("%b %d at %I:%M %p")
        except ValueError:
            return str(value or "the requested time")

    def _release_due_scheduled(self) -> None:
        released: list[Job] = []
        for job in reversed(self.jobs):
            if job.status == JobStatus.SCHEDULED and not self._is_future(job.scheduled_for):
                job.status = JobStatus.QUEUED
                self._queue.append(job.id)
                self._append_log(job, "Scheduled start time reached; waiting for the training slot.")
                released.append(job)
        if not released:
            return
        self._save()
        for job in released:
            self.job_updated.emit(job)
        self._start_next()

    def get(self, job_id: str) -> Job:
        for job in self.jobs:
            if job.id == job_id:
                return job
        raise KeyError(f"Unknown job: {job_id}")

    def retry(self, job_id: str) -> Job:
        """Create an approval-gated retry, resuming interrupted DDPM work when possible."""
        original = self.get(job_id)
        plan = ExecutionPlan.from_dict(original.to_dict()["plan"])
        resume_note = ""
        if original.status == JobStatus.INTERRUPTED:
            start_index = max(0, min(original.current_step, len(plan.steps) - 1))
            plan.steps = plan.steps[start_index:]
            if plan.steps:
                step = plan.steps[0]
                resume_note = self._prepare_ddpm_resume(step.arguments, step.tool_id)
        for step in plan.steps:
            step.status = StepStatus.PENDING
        plan.id = original.plan.id + "-retry"
        plan.created_at = utc_now()
        plan.requires_confirmation = True
        plan.confirmation_reason = (
            "This retries a previous job. Review paths, checkpoints, and settings "
            "because files or available resources may have changed."
        )
        retry = self.submit(plan)
        if resume_note:
            self._append_log(retry, resume_note)
        return retry

    @staticmethod
    def _prepare_ddpm_resume(arguments: dict[str, Any], tool_id: str) -> str:
        """Attach the newest complete Accelerate checkpoint to a DDPM retry."""
        if tool_id != "ddpm_trainer":
            return ""
        output = Path(str(arguments.get("output_dir", ""))).expanduser()
        dataset = Path(str(arguments.get("dataset_dir", ""))).expanduser()
        try:
            checkpoints = sorted(
                (path for path in output.glob("checkpoint-*") if path.is_dir()
                 and (path / "unet" / "diffusion_pytorch_model.safetensors").is_file()
                 and (path / "optimizer.bin").is_file() and (path / "scheduler.bin").is_file()),
                key=lambda path: int(path.name.rsplit("-", 1)[-1]),
            )
            image_count = sum(1 for item in dataset.iterdir() if item.is_file() and item.suffix.lower() in {".png", ".jpg", ".jpeg", ".webp", ".bmp"})
            batch_size = max(1, int(arguments.get("batch_size", 1)))
            completed_epochs = int(checkpoints[-1].name.rsplit("-", 1)[-1]) // max(1, math.ceil(image_count / batch_size))
            remaining_epochs = int(arguments.get("epochs", 0)) - completed_epochs
        except (IndexError, OSError, ValueError, TypeError):
            return ""
        if remaining_epochs <= 0:
            return ""
        arguments["resume_from"] = str(checkpoints[-1])
        arguments["epochs"] = remaining_epochs
        return f"Resuming interrupted DDPM training from {checkpoints[-1].name} (about epoch {completed_epochs}; {remaining_epochs} epochs remaining)."

    def end_task(self, job_id: str) -> bool:
        """Acknowledge an interrupted job and leave it inactive in history.

        This is intentionally limited to interrupted jobs: active work must still use
        ``cancel`` so its worker receives the cancellation signal.
        """
        job = self.get(job_id)
        if job.status != JobStatus.INTERRUPTED:
            return False
        job.status = JobStatus.CANCELLED
        job.ended_at = utc_now()
        job.logs = [
            line
            for line in job.logs
            if not line.startswith("[startup] Previous session ended before this job.")
        ]
        self._append_log(job, "Interrupted job ended by user; no retry is pending.")
        self._save()
        self.job_updated.emit(job)
        return True

    def remove_completed_or_failed(self) -> int:
        """Remove completed and failed history records without touching output files."""
        removable = {JobStatus.FINISHED, JobStatus.FAILED}
        before = len(self.jobs)
        self.jobs = [job for job in self.jobs if job.status not in removable]
        removed = before - len(self.jobs)
        if removed:
            self._save()
        return removed

    def _start_next(self) -> None:
        if self._worker is not None and self._worker.isRunning():
            return
        while self._queue:
            job_id = self._queue.pop(0)
            job = self.get(job_id)
            if job.status != JobStatus.QUEUED:
                continue
            self._active_job = job
            job.status = JobStatus.RUNNING
            job.started_at = utc_now()
            self._append_log(job, f"Job {job.id} started.")
            self._worker = JobWorker(job, self.executor)
            self._worker.event.connect(self._handle_event)
            self._worker.finished.connect(self._worker_finished)
            self._save()
            self.active_changed.emit(job)
            self.job_updated.emit(job)
            self.notification.emit("Training started" if self._has_training(job) else "Job started", job.plan.project_name)
            self._worker.start()
            return
        self._active_job = None
        self.active_changed.emit(None)

    def _handle_event(self, event: dict[str, Any]) -> None:
        job = self._active_job
        if job is None:
            return
        event_type = event.get("type")
        if event_type == "step_started":
            index = int(event["index"])
            job.current_step = index
            job.plan.steps[index].status = StepStatus.RUNNING
            if job.plan.steps[index].tool_id.endswith("_trainer"):
                job.preview_path = None
                job.preview_epoch = 0
                job.preview_next_epoch = 0
                job.preview_prompt = ""
                job.preview_seed = None
                job.preview_steps = 0
                job.preview_kind = "training"
                job.preview_current = 0
                job.preview_total = 0
                job.preview_image_index = 0
                job.preview_image_count = 0
                job.eta_seconds = None
                job.estimated_completion_at = None
                job.progress_current = 0
                job.progress_total = 0
                job.progress_rate = 0.0
                job.progress_unit = "step"
            self._append_log(job, f"Starting: {job.plan.steps[index].title}")
        elif event_type == "progress":
            job.progress = int(event["overall"])
            job.eta_seconds = _safe_int(event.get("eta_seconds")) or None
            job.estimated_completion_at = str(event.get("estimated_completion_at") or "") or None
            job.progress_current = _safe_int(event.get("progress_current"))
            job.progress_total = _safe_int(event.get("progress_total"))
            job.progress_rate = float(event.get("progress_rate", 0.0) or 0.0)
            job.progress_unit = str(event.get("progress_unit", "step") or "step")
            message = str(event["message"])
            if message and (not job.logs or message not in job.logs[-1]):
                self._append_log(job, message)
            self._schedule_job_update(job)
            return
        elif event_type == "log":
            self._append_log(job, str(event["message"]))
            self._schedule_job_update(job)
            return
        elif event_type == "preview":
            job.preview_path = str(event.get("path", "")) or None
            job.preview_epoch = int(event.get("epoch", 0) or 0)
            job.preview_next_epoch = int(event.get("next_epoch", 0) or 0)
            job.preview_prompt = str(event.get("prompt", ""))
            seed = event.get("seed")
            job.preview_seed = int(seed) if seed is not None else None
            job.preview_steps = int(event.get("steps", 0) or 0)
            job.preview_kind = str(event.get("kind", "training"))
            job.preview_current = int(event.get("current", 0) or 0)
            job.preview_total = int(event.get("total", 0) or 0)
            job.preview_image_index = int(event.get("image_index", 0) or 0)
            job.preview_image_count = int(event.get("image_count", 0) or 0)
            label = "Denoising" if job.preview_kind == "generation" else "Training"
            position = f" step {job.preview_current}" if job.preview_current else f" epoch {job.preview_epoch}"
            self._append_log(job, f"{label} preview updated at{position}.")
            self._schedule_job_update(job)
            return
        elif event_type == "step_finished":
            index = int(event["index"])
            job.plan.steps[index].status = StepStatus.FINISHED
            result = event.get("result") or {}
            self.assets.ingest_result(result)
            if result.get("output_folder"):
                job.output_folder = str(result["output_folder"])
            if job.plan.steps[index].tool_id.endswith("_trainer"):
                evaluation = evaluate_job_output(job)
                if evaluation:
                    evaluation["step"] = job.plan.steps[index].title
                    evaluation["model_name"] = str(result.get("model_name", ""))
                    reports = list(job.nova_report.get("evaluations", []))
                    reports.append(evaluation)
                    job.nova_report = {
                        "agent": "NOVA", "evaluations": reports, "latest": evaluation,
                    }
                    self._append_log(
                        job,
                        f"NOVA — {evaluation['status']}: {evaluation['summary']}",
                    )
            self._append_log(job, f"Finished: {job.plan.steps[index].title}")
        elif event_type == "completed":
            job.status = JobStatus.FINISHED
            job.progress = 100
            job.eta_seconds = 0
            job.estimated_completion_at = utc_now()
            job.ended_at = utc_now()
            self._append_log(job, "Job finished successfully.")
            self._record_experiment(job)
            self.notification.emit("Job complete", job.plan.project_name)
        elif event_type == "cancelled":
            job.status = JobStatus.CANCELLED
            job.ended_at = utc_now()
            self._append_log(job, str(event.get("message", "Job cancelled.")))
            for step in job.plan.steps:
                if step.status == StepStatus.RUNNING:
                    step.status = StepStatus.SKIPPED
            self._record_experiment(job)
            self.notification.emit("Job cancelled", job.plan.project_name)
        elif event_type == "adjustment_ready":
            index = max(0, min(job.current_step, len(job.plan.steps) - 1))
            remaining_steps = job.plan.steps[index:]
            step = remaining_steps[0]
            updates = dict(event.get("updates") or {})
            step.arguments.update(updates)
            checkpoint = str(event.get("checkpoint", ""))
            completed_epochs = max(0, int(event.get("completed_epochs", 0) or 0))
            total_epochs = max(1, int(step.arguments.get("epochs", 1)))
            step.arguments["epochs"] = max(1, total_epochs - completed_epochs)
            step.arguments["completed_epochs"] = completed_epochs
            if checkpoint:
                step.arguments["resume_from"] = checkpoint
            for pending in remaining_steps:
                pending.status = StepStatus.PENDING
            job.plan.steps = remaining_steps
            job.current_step = -1
            job.status = JobStatus.QUEUED
            job.eta_seconds = None
            job.estimated_completion_at = None
            self._queue.insert(0, job.id)
            self._append_log(job, f"Epoch {completed_epochs} checkpoint is complete. Restarting with {updates}.")
            self.notification.emit("Training settings ready", "Restarting from the completed epoch checkpoint.")
        elif event_type == "failed":
            job.status = JobStatus.FAILED
            job.ended_at = utc_now()
            job.error = str(event.get("message", "Unknown error"))
            if 0 <= job.current_step < len(job.plan.steps):
                job.plan.steps[job.current_step].status = StepStatus.FAILED
            self._append_log(job, f"Failed: {job.error}")
            self.logger.error(
                "Job %s failed (%s): %s",
                job.id,
                event.get("exception"),
                job.error,
            )
            self._record_experiment(job)
            self.notification.emit("Job failed", job.error)
        self._save()
        self.job_updated.emit(job)

    def _schedule_job_update(self, job: Job) -> None:
        self._pending_update_job_ids.add(job.id)
        if QCoreApplication.instance() is None:
            self._flush_pending_job_updates()
            return
        if self._pending_update_timer_active:
            return
        self._pending_update_timer_active = True
        QTimer.singleShot(300, self._flush_pending_job_updates)

    def _flush_pending_job_updates(self) -> None:
        if not self._pending_update_job_ids:
            self._pending_update_timer_active = False
            return
        pending_ids = list(self._pending_update_job_ids)
        self._pending_update_job_ids.clear()
        self._pending_update_timer_active = False
        self._save()
        for job_id in pending_ids:
            try:
                self.job_updated.emit(self.get(job_id))
            except KeyError:
                continue

    def _worker_finished(self) -> None:
        self._worker = None
        self._active_job = None
        self.active_changed.emit(None)
        self._start_next()

    def supervise(self, snapshot: Any) -> None:
        """Let ATLAS inspect the active training run and apply critical pauses."""
        self._last_snapshot = snapshot
        job = self._active_job
        if job is None or job.status != JobStatus.RUNNING:
            return
        decision = self.atlas.observe(job, snapshot)
        previous = str(job.atlas_report.get("message", ""))
        unchanged = (
            previous == decision.message
            and job.atlas_report.get("severity") == decision.severity
            and job.atlas_report.get("action") == decision.action
        )
        if unchanged:
            return
        job.atlas_report = {
            "agent": "ATLAS",
            "severity": decision.severity,
            "message": decision.message,
            "action": decision.action,
            "updated_at": utc_now(),
        }
        if decision.message != previous and decision.severity in {"warning", "critical"}:
            self._append_log(job, f"ATLAS {decision.severity.upper()} — {decision.message}")
            self.notification.emit(f"ATLAS {decision.severity}", decision.message)
        self._save()
        self.job_updated.emit(job)
        if decision.action == "pause" and job.status == JobStatus.RUNNING:
            self.pause(job.id)

    def _record_experiment(self, job: Job) -> None:
        if not self._has_training(job):
            return
        try:
            self.experiments.record_job(job, self._last_snapshot)
        except Exception as exc:
            self.logger.warning("Experiment tracking failed for %s: %s", job.id, exc)

    def _append_log(self, job: Job, message: str) -> None:
        timestamp = datetime.now().strftime("%H:%M:%S")
        line = f"[{timestamp}] {message}"
        job.logs.append(line)
        job.logs = job.logs[-1000:]
        self.log_added.emit(job.id, line)
        self.logger.info("Job %s | %s", job.id, message)

    @staticmethod
    def _has_training(job: Job) -> bool:
        return any(step.tool_id.endswith("trainer") for step in job.plan.steps)

    def _load(self) -> None:
        if not self.jobs_path.exists():
            return
        try:
            payload = json.loads(self.jobs_path.read_text(encoding="utf-8"))
            self.jobs = [Job.from_dict(item) for item in payload.get("jobs", [])]
        except (OSError, ValueError, TypeError, KeyError, json.JSONDecodeError):
            self.jobs = []
            return
        for job in self.jobs:
            if job.status in {JobStatus.RUNNING, JobStatus.PAUSED}:
                job.status = JobStatus.INTERRUPTED
                job.ended_at = utc_now()
                job.logs.append(
                    "[startup] Previous session ended before this job. "
                    "Review it before retrying."
                )
            elif job.status == JobStatus.QUEUED:
                self._queue.append(job.id)
                if not any("Queued job restored" in line for line in job.logs[-5:]):
                    job.logs.append("[startup] Queued job restored and will run when ADAM is ready.")
        self._save()

    def _save(self) -> None:
        self.jobs_path.parent.mkdir(parents=True, exist_ok=True)
        temporary = self.jobs_path.with_suffix(".tmp")
        temporary.write_text(
            json.dumps({"jobs": [job.to_dict() for job in self.jobs]}, indent=2),
            encoding="utf-8",
        )
        temporary.replace(self.jobs_path)

    def shutdown(self) -> None:
        if self._worker is not None and self._worker.isRunning():
            self._worker.cancel()
            self._worker.wait(2500)
        if self._active_job and self._active_job.status in {
            JobStatus.RUNNING,
            JobStatus.PAUSED,
        }:
            self._active_job.status = JobStatus.CANCELLED
            self._active_job.ended_at = utc_now()
            self._append_log(self._active_job, "ADAM closed; the active job was stopped.")
            self._save()