adam/generations.py

from __future__ import annotations

import json
import re
from html import unescape
from dataclasses import dataclass
from pathlib import Path
from typing import Any

from adam.models import ExecutionPlan, PlanStep
from adam.registry import ToolRegistry, ToolSpec


IMAGE_EXTENSIONS = {".png", ".jpg", ".jpeg", ".webp", ".bmp"}


@dataclass(frozen=True, slots=True)
class ImportedLoRAMetadata:
    """Portable LoRA-generation settings copied from an image or another app."""

    prompt: str
    negative_prompt: str
    seed: int
    steps: int
    cfg_scale: float | None
    base_model_path: str
    lora_path: str
    lora_strength: float | None
    sampler: str
    width: int | None = None
    height: int | None = None
    sampler_note: str = ""


def parse_lora_generation_metadata(text: str) -> ImportedLoRAMetadata:
    """Parse pasted LoRA metadata without requiring an intermediate JSON file.

    Browser copy/paste sometimes HTML-escapes JSON (for example `` `` for a
    space), so decode those entities before loading it.  This intentionally
    accepts only the small, reproducible LoRA schema ADAM understands.
    """
    source = str(text or "").strip()
    # A few sites escape twice when metadata is copied out of a code block.
    for _ in range(2):
        decoded = unescape(source)
        if decoded == source:
            break
        source = decoded
    try:
        payload = json.loads(source)
    except (TypeError, ValueError, json.JSONDecodeError):
        # Some viewers copy a display block rather than strict JSON, omitting
        # braces, commas, or quotes around values. Recover its named fields.
        fields: dict[str, Any] = {}
        for key in ("prompt", "negative_prompt", "seed", "steps", "cfg_scale", "model", "sampler"):
            match = re.search(rf'["\']?{key}["\']?\s*:\s*(?:["\']([^"\']*)["\']|([^\r\n]+))', source, re.I)
            if match:
                fields[key] = (match.group(1) if match.group(1) is not None else match.group(2)).strip().rstrip(",").strip()
        lora_path_match = re.search(r'["\']?path["\']?\s*:\s*["\']?([^"\',\r\n}\]]+)', source, re.I)
        strength_match = re.search(r'["\']?strength["\']?\s*:\s*([^,\r\n}\]]+)', source, re.I)
        if lora_path_match:
            fields["loras"] = [{
                "path": lora_path_match.group(1).strip(),
                "strength": strength_match.group(1).strip() if strength_match else None,
            }]
        payload = fields
    if not isinstance(payload, dict):
        raise ValueError("Metadata must be a JSON object.")

    # Source apps vary between `width`/`height` and `Width`/`Height`.
    payload = {str(key).casefold(): value for key, value in payload.items()}

    def text_value(key: str, *, required: bool = False) -> str:
        value = payload.get(key, "")
        if value is None:
            value = ""
        if not isinstance(value, str):
            raise ValueError(f"{key.replace('_', ' ').title()} must be text.")
        value = value.strip()
        if required and not value:
            raise ValueError(f"Metadata is missing {key.replace('_', ' ')}.")
        return value

    loras = payload.get("loras")
    if not isinstance(loras, list):
        raise ValueError("Metadata needs a 'loras' list (use [] when no LoRA was used).")
    if loras and not isinstance(loras[0], dict):
        raise ValueError("The first LoRA entry must be an object.")
    lora_path = str(loras[0].get("path", "")).strip() if loras else ""
    if loras and not lora_path:
        raise ValueError("The first LoRA entry needs a path.")
    try:
        seed = int(payload.get("seed", 0))
        steps = int(payload.get("steps", 30))
    except (TypeError, ValueError) as exc:
        raise ValueError("Seed and steps must be whole numbers.") from exc
    # Automatic1111-style metadata commonly uses -1 for a fresh random seed.
    # ADAM uses 0 for the same behavior in generation plans.
    if seed == -1:
        seed = 0
    if not 0 <= seed <= 2_147_483_647 or steps < 1:
        raise ValueError("Seed or steps is outside ADAM's supported range.")
    cfg_value = payload.get("cfg_scale")
    try:
        cfg_scale = float(cfg_value) if cfg_value is not None else None
    except (TypeError, ValueError) as exc:
        raise ValueError("CFG scale must be a number.") from exc
    strength_value = loras[0].get("strength") if loras else None
    try:
        lora_strength = float(strength_value) if strength_value is not None else None
    except (TypeError, ValueError) as exc:
        raise ValueError("LoRA strength must be a number.") from exc
    def dimension(key: str) -> int | None:
        value = payload.get(key)
        if value is None or value == "":
            return None
        try:
            result = int(value)
        except (TypeError, ValueError) as exc:
            raise ValueError(f"{key.title()} must be a whole number.") from exc
        if not 256 <= result <= 2048:
            raise ValueError(f"{key.title()} must be between 256 and 2048 pixels.")
        return result

    sampler = text_value("sampler") or "DPM++ 2M"
    supported = {"DPM++ 2M", "DPM++ 2M Karras", "DPM++ 2M SDE", "DPM++ 2M SDE Karras", "DPM++ SDE", "DPM++ SDE Karras", "Euler", "Euler a", "Heun", "LMS", "DDIM"}
    sampler_note = ""
    if sampler not in supported:
        normalized = sampler.casefold()
        replacement = "DPM++ SDE" if "sde" in normalized else "DPM++ 2M" if "2m" in normalized else ""
        if not replacement:
            raise ValueError(f"Sampler '{sampler}' is not supported by ADAM's LoRA generator.")
        sampler_note = f"Sampler '{sampler}' was mapped to '{replacement}' because the connected generator does not support it."
        sampler = replacement
    return ImportedLoRAMetadata(
        prompt=text_value("prompt", required=True),
        negative_prompt=text_value("negative_prompt"),
        seed=seed,
        steps=steps,
        cfg_scale=cfg_scale,
        base_model_path=text_value("model", required=True),
        lora_path=lora_path,
        lora_strength=lora_strength,
        sampler=sampler,
        width=dimension("width"),
        height=dimension("height"),
        sampler_note=sampler_note,
    )


def parse_pasted_lora_metadata_request(text: str) -> "ChatGenerationRequest | None":
    """Recognize copied LoRA metadata as an unambiguous prompt-box command."""
    source = str(text or "")
    if not re.search(r'["\']?prompt["\']?\s*:', source, re.I) or not re.search(
        r'["\']?(?:loras|negative_prompt|sampler)["\']?\s*:', source, re.I
    ):
        return None
    try:
        metadata = parse_lora_generation_metadata(source)
    except ValueError:
        return None
    return ChatGenerationRequest(
        prompt=metadata.prompt,
        subject=Path(metadata.lora_path).stem if metadata.lora_path else "",
        provider_hint="lora",
        model_query=Path(metadata.lora_path).stem if metadata.lora_path else "",
        base_model_query=Path(metadata.base_model_path).stem,
        negative_prompt=metadata.negative_prompt,
        steps=metadata.steps,
        sampler=metadata.sampler,
        seed=metadata.seed,
        cfg_scale=metadata.cfg_scale,
        lora_strength=metadata.lora_strength,
        width=metadata.width,
        height=metadata.height,
        has_positive_prompt=True,
        is_pasted_metadata=True,
        metadata_model_path=metadata.lora_path,
        metadata_base_model_path=metadata.base_model_path,
    )


def parse_plain_generation_metadata(text: str) -> "ChatGenerationRequest | None":
    """Read the common CivitAI/A1111 and PixAI copied-text metadata layouts."""
    source = str(text or "").replace("\r\n", "\n").strip()
    if not source:
        return None
    civitai = re.search(r"\bNegative\s+prompt\s*:", source, re.I)
    pixai = re.search(r"\b(?:Sampling\s+Steps|Original\s+Prompt)\b", source, re.I)
    if not civitai and not pixai:
        return None
    prompt = ""
    negative = ""
    if civitai:
        prompt = source[:civitai.start()].strip(" ,\n")
        tail = source[civitai.end():]
        settings = re.search(r"\b(?:Steps|Size)\s*:", tail, re.I)
        negative = tail[:settings.start()].strip(" ,\n") if settings else tail.strip(" ,\n")
    else:
        original = re.search(r"\bOriginal\s+Prompt\s*\n+(.+?)(?=\n+\s*Size\s*\n)", source, re.I | re.S)
        prompt = (original.group(1) if original else source.split("\n\n", 1)[0]).strip(" ,\n")
        negative_match = re.search(r"\n\s*Negative\s*\n+(.+?)(?=\n\s*(?:Prompt\s+Helper|#|$))", source, re.I | re.S)
        negative = negative_match.group(1).strip(" ,\n") if negative_match else ""

    def number(pattern: str, kind):
        match = re.search(pattern, source, re.I)
        return kind(match.group(1)) if match else None
    steps = number(r"\b(?:Sampling\s+)?Steps\s*:?\s*(\d+)", int)
    cfg = number(r"\bCFG\s*(?:Scale)?\s*:?\s*(\d+(?:\.\d+)?)", float)
    seed = number(r"\bSeed\s*:?\s*(-?\d+)", int)
    if seed == -1:
        seed = 0
    size = re.search(r"\bSize\s*:?\s*(\d+)\s*[x×]\s*(\d+)", source, re.I)
    sampler_match = re.search(r"\b(?:Sampling\s+Method|Sampler)\s*:?\s*([^\n,]+)", source, re.I)
    sampler = sampler_match.group(1).strip() if sampler_match else ""
    if sampler:
        folded = sampler.casefold()
        sampler = next((name for name in ("DPM++ 2M SDE Karras", "DPM++ 2M SDE", "DPM++ 2M Karras", "DPM++ SDE Karras", "DPM++ SDE", "DPM++ 2M") if name.casefold() in folded), sampler)
    loras = re.findall(r"<lora:([^:>]+)(?::([\d.]+))?>", prompt, re.I)
    if loras:
        prompt = re.sub(r"\s*<lora:[^>]+>", "", prompt, flags=re.I).strip(" ,")
    return ChatGenerationRequest(
        prompt=prompt, provider_hint="lora", model_query=loras[0][0].strip() if len(loras) == 1 else "",
        negative_prompt=negative, steps=steps, seed=seed, sampler=sampler, cfg_scale=cfg,
        lora_strength=float(loras[0][1]) if len(loras) == 1 and loras[0][1] else None,
        width=int(size.group(1)) if size else None, height=int(size.group(2)) if size else None,
        has_positive_prompt=True, is_pasted_metadata=True,
    )


@dataclass(frozen=True, slots=True)
class ChatGenerationRequest:
    """Generation settings recognized from a Command Center message."""

    prompt: str
    subject: str = ""
    provider_hint: str = ""
    model_query: str = ""
    base_model_query: str = ""
    negative_prompt: str = ""
    image_count: int | None = None
    steps: int | None = None
    sampler: str = ""
    aspect_ratio: str = ""
    seed: int | None = None
    cfg_scale: float | None = None
    lora_strength: float | None = None
    denoise_strength: float | None = None
    reference_strength: int | None = None
    reference_image: str = ""
    has_positive_prompt: bool = False
    is_pasted_metadata: bool = False
    metadata_model_path: str = ""
    metadata_base_model_path: str = ""
    width: int | None = None
    height: int | None = None


_QUOTED = r'["\u201c\u201d]([^"\u201c\u201d]+)["\u201c\u201d]'


def _clean_chat_value(value: str) -> str:
    return value.strip().strip('"\u201c\u201d').strip(" ,.;")


def generation_model_match_score(query: str, model_name: str) -> int:
    """Score whether conversational subject text clearly names a saved model."""
    def words(value: str) -> list[str]:
        value = re.sub(r"(?<=[a-z0-9])(?=[A-Z])", " ", value)
        value = re.sub(r"\bpixel\s+row\b", " ", value, flags=re.I)
        value = re.sub(r"\binr\s*flow\b", " ", value, flags=re.I)
        ignored = {"a", "an", "the", "of", "image", "picture", "model", "ddpm", "flow", "matching", "lora", "inr"}
        return [word for word in re.findall(r"[a-z0-9]+", value.casefold()) if word not in ignored]

    query_words = words(query)
    model_words = words(model_name)
    if not query_words or not model_words:
        return 0
    query_compact = "".join(query_words)
    model_compact = "".join(model_words)
    if query_compact == model_compact:
        return 120
    if query_compact in model_compact:
        return 100 + min(10, len(query_words))
    shared = len(set(query_words) & set(model_words))
    if shared == len(set(query_words)):
        return 90 + shared
    coverage = shared / len(set(query_words))
    return 60 + shared if shared >= 2 and coverage >= 0.7 else 0


def parse_chat_generation_request(text: str) -> ChatGenerationRequest | None:
    """Recognize a concise natural-language image generation command.

    This intentionally requires both a creation verb and the word image/picture so
    ordinary planning requests continue through the regular Command Center planner.
    """
    pasted_metadata = parse_pasted_lora_metadata_request(text)
    if pasted_metadata is not None:
        return pasted_metadata
    plain_metadata = parse_plain_generation_metadata(text)
    if plain_metadata is not None:
        return plain_metadata
    request = " ".join(text.strip().split())
    # Do not treat any request that happens to contain both words as an image
    # generation command.  Dataset requests commonly say things such as
    # "image mode" and "generate captions"; those must continue to the
    # regular planner (and, in particular, the video dataset collector).
    # Require the creation verb to directly introduce the image noun instead.
    generation_command = re.compile(
        r"\b(?:generate|create|make)\s+"
        r"(?:(?:an?|the|\d+)\s+)?"
        r"(?:[\"\u201c\u201d]?(?:ddpm|ddim|inr\s*flow|flow(?:\s+matching)?|pixel\s*row|lora)[\"\u201c\u201d]?\s+)?"
        r"(?:images?|pictures?)\b",
        re.I,
    )
    if not request or not generation_command.search(request):
        return None

    provider_hint = ""
    provider_match = re.search(
        r"\b(ddpm|ddim|inr\s*flow|flow(?:\s+matching)?|pixel\s*row|lora)\b[\"\u201c\u201d]?(?=\s+(?:image|picture))",
        request,
        re.I,
    )
    if provider_match:
        hint = provider_match.group(1).casefold()
        provider_hint = (
            "ddpm" if hint in {"ddpm", "ddim"}
            else "inrflow" if hint.replace(" ", "") == "inrflow"
            else "flow" if hint.startswith("flow")
            else "pixelrow" if hint.replace(" ", "") == "pixelrow"
            else "lora"
        )
    # Support natural phrasing such as "Generate an image of LoRA OrangeCat".
    lora_subject_match = re.search(
        rf"\b(?:image|picture)s?\s+of\s+(?:a\s+)?LoRA\s+{_QUOTED}",
        request,
        re.I,
    )
    if not lora_subject_match:
        lora_subject_match = re.search(
            r"\b(?:image|picture)s?\s+of\s+(?:a\s+)?LoRA\s+(.+?)(?=\s*(?:,|with\s+base\s+model|base\s+model|positive\s+prompt|negative\s+prompt|\d+\s+steps?|$))",
            request,
            re.I,
        )
    if lora_subject_match:
        provider_hint = "lora"

    subject = ""
    prompt_match = re.search(
        rf"\b(?:image|picture)s?\s+(?:of|showing|depicting)\s+{_QUOTED}",
        request,
        re.I,
    )
    if prompt_match:
        subject = _clean_chat_value(prompt_match.group(1))
    else:
        prompt_match = re.search(
            r"\b(?:image|picture)s?\s+(?:of|showing|depicting)\s+(.+?)(?=\s+(?:for|using|with|on|at)\s+|,|$)",
            request,
            re.I,
        )
        if prompt_match:
            subject = _clean_chat_value(prompt_match.group(1))

    # In promptless commands, a provider suffix is usually part of the saved
    # model name (for example, "Minecraft Flow"), not prompt prose.
    if not provider_hint and subject:
        if re.search(r"\binr\s*flow\s*$", subject, re.I):
            provider_hint = "inrflow"
        elif re.search(r"\bflow(?:\s+match(?:ing)?)?\s*$", subject, re.I):
            provider_hint = "flow"
        elif re.search(r"\bddpm\s*$", subject, re.I):
            provider_hint = "ddpm"
        elif re.search(r"\bpixel\s*row\s*$", subject, re.I):
            provider_hint = "pixelrow"

    positive_match = re.search(
        rf"\bpositive\s+prompt(?:\s+of|\s*=|\s*:)?\s*{_QUOTED}",
        request,
        re.I,
    )
    prompt = _clean_chat_value(positive_match.group(1)) if positive_match else subject
    negative_match = re.search(
        rf"\bnegative\s+prompt(?:\s+of|\s*=|\s*:)?\s*{_QUOTED}",
        request,
        re.I,
    )
    negative_prompt = _clean_chat_value(negative_match.group(1)) if negative_match else ""

    # A quoted value after "using/with model" is treated as a completed model name.
    model_query = ""
    model_match = re.search(
        rf"\b(?:using|with)\s+(?:the\s+)?(?:model\s+)?{_QUOTED}(?:\s+model)?",
        request,
        re.I,
    )
    if model_match:
        model_query = _clean_chat_value(model_match.group(1))
    if lora_subject_match:
        model_query = _clean_chat_value(lora_subject_match.group(1))
    base_model_match = re.search(
        rf"\bbase\s+model(?:\s+of|\s*=|\s*:)?\s*{_QUOTED}",
        request,
        re.I,
    )
    base_model_query = _clean_chat_value(base_model_match.group(1)) if base_model_match else ""

    count_match = re.search(r"\b(?:generate|create|make)\s+[\"\u201c]?([1-9]\d*)[\"\u201d]?\s+(?:images|pictures)\b", request, re.I)
    steps_match = re.search(r"\b[\"\u201c]?(\d{1,4})[\"\u201d]?\s+(?:inference\s+)?steps?\b", request, re.I)
    seed_match = re.search(r"\bseed(?:\s+of|\s*=|\s*:)?\s*[\"\u201c]?(\d{1,10})[\"\u201d]?", request, re.I)
    sampler_match = re.search(
        r"\b(?:on|with|using)\s+[\"\u201c]?(DDIM|DDPM|Heun|Euler(?:\s+a)?|DPM\+\+\s*2M)[\"\u201d]?\s+sampler\b"
        r"|\bsampler(?:\s+of|\s*=|\s*:)?\s+[\"\u201c]?(DDIM|DDPM|Heun|Euler(?:\s+a)?|DPM\+\+\s*2M)",
        request,
        re.I,
    )
    aspect_match = re.search(r"\b(?:aspect\s+ratio(?:\s+of)?|ratio)\s*[\"\u201c]?(\d+\s*:\s*\d+)", request, re.I)
    cfg_match = re.search(r"\bCFG(?:\s+scale)?(?:\s+of|\s*=|\s*:)?\s*[\"\u201c]?(\d+(?:\.\d+)?)", request, re.I)
    lora_strength_match = re.search(r"\bLoRA\s+strength(?:\s+of|\s*=|\s*:)?\s*[\"\u201c]?(\d+(?:\.\d+)?)", request, re.I)
    denoise_match = re.search(r"\bdenoise(?:\s+strength)?(?:\s+of|\s*=|\s*:)?\s*[\"\u201c]?(\d+(?:\.\d+)?)", request, re.I)
    reference_strength_match = re.search(r"\breference\s+strength(?:\s+of|\s*=|\s*:)?\s*[\"\u201c]?(\d{1,3})\s*%?", request, re.I)

    sampler_value = (sampler_match.group(1) or sampler_match.group(2)) if sampler_match else ""
    sampler = sampler_value.upper() if sampler_value else ""
    sampler = {"EULER A": "Euler a", "EULER": "Euler", "HEUN": "Heun", "DPM++ 2M": "DPM++ 2M"}.get(sampler, sampler)
    aspect_ratio = re.sub(r"\s+", "", aspect_match.group(1)) if aspect_match else ""

    # If the only DDPM/DDIM mention labels the sampler, do not let it choose a provider.
    if not provider_match and re.search(r"\b(?:DDPM|DDIM)\b\s+sampler", request, re.I):
        provider_hint = ""
    return ChatGenerationRequest(
        prompt=prompt,
        subject=subject,
        provider_hint=provider_hint,
        model_query=model_query,
        base_model_query=base_model_query,
        negative_prompt=negative_prompt,
        image_count=int(count_match.group(1)) if count_match else None,
        steps=int(steps_match.group(1)) if steps_match else None,
        sampler=sampler,
        aspect_ratio=aspect_ratio,
        seed=int(seed_match.group(1)) if seed_match else None,
        cfg_scale=float(cfg_match.group(1)) if cfg_match else None,
        lora_strength=float(lora_strength_match.group(1)) if lora_strength_match else None,
        denoise_strength=float(denoise_match.group(1)) if denoise_match else None,
        reference_strength=int(reference_strength_match.group(1)) if reference_strength_match else None,
        has_positive_prompt=positive_match is not None,
    )


def generation_output_folder(root: Path, provider_id: str, model_name: str) -> Path:
    """Return the browseable folder shared by all output from one model."""
    def safe(value: str, fallback: str) -> str:
        value = re.sub(r"[<>:\"/\\\\|?*\x00-\x1f]+", " ", value.strip())
        return re.sub(r"\s+", " ", value).strip(" .")[:96] or fallback

    folder = root.resolve() / "data" / "generations" / safe(provider_id, "generator") / safe(model_name, "model")
    folder.mkdir(parents=True, exist_ok=True)
    return folder


def generation_metadata_path(folder: Path, timestamp: str, job_id: str) -> Path:
    return folder / f"generation_{timestamp}_{job_id}.json"


@dataclass(frozen=True, slots=True)
class GenerationRecord:
    metadata_path: Path
    folder: Path
    images: tuple[Path, ...]
    provider_id: str
    provider_name: str
    model_name: str
    model_path: str
    prompt: str
    seed: int
    steps: int
    sampler: str
    aspect_ratio: str
    created_at: str
    smart_generation: dict[str, Any]
    image_evaluations: dict[str, dict[str, Any]]

    @classmethod
    def from_metadata(cls, metadata_path: Path) -> "GenerationRecord | None":
        try:
            payload = json.loads(metadata_path.read_text(encoding="utf-8"))
        except (OSError, ValueError, TypeError, json.JSONDecodeError):
            return None
        folder = metadata_path.parent
        listed_images = payload.get("images", [])
        images = tuple(Path(str(path)) for path in listed_images if Path(str(path)).is_file())
        if not images:  # Support generation folders created by older ADAM versions.
            images = tuple(path for path in sorted(folder.iterdir()) if path.is_file() and path.suffix.casefold() in IMAGE_EXTENSIONS)
        if not images:
            return None
        return cls(
            metadata_path=metadata_path,
            folder=folder,
            images=images,
            provider_id=str(payload.get("provider_id", "")),
            provider_name=str(payload.get("provider_name", "Unknown generator")),
                model_name=str(payload.get("model_name", folder.name)),
            model_path=str(payload.get("model_path", "")),
            prompt=str(payload.get("prompt", "")),
            seed=int(payload.get("seed", 0) or 0),
            steps=int(payload.get("steps", 0) or 0),
            sampler=str(payload.get("sampler", "")),
            aspect_ratio=str(payload.get("aspect_ratio", "")),
            created_at=str(payload.get("created_at", "")),
            smart_generation=dict(payload.get("smart_generation") or {}),
            image_evaluations={
                str(Path(path).resolve()): dict(value)
                for path, value in dict(payload.get("image_evaluations") or {}).items()
                if isinstance(value, dict)
            },
        )


@dataclass(frozen=True, slots=True)
class GenerationModelFolder:
    """A model-centered view over existing generation batches."""

    key: str
    model_name: str
    model_path: str
    provider_id: str
    provider_name: str
    records: tuple[GenerationRecord, ...]
    image_count: int
    cover_image: Path | None
    latest_at: str


@dataclass(frozen=True, slots=True)
class GenerationProviderFolder:
    """A generator-centered view over the on-disk generation folders.

    Output is stored as ``generations/<generator>/<model>/...``.  This view
    intentionally exposes that first directory level in the UI, keeping all
    images made by one generator together without moving any user files.
    """

    key: str
    provider_id: str
    provider_name: str
    records: tuple[GenerationRecord, ...]
    model_count: int
    image_count: int
    latest_at: str


def generation_provider_key(record: GenerationRecord) -> str:
    """Return the stable key for the generator directory containing a batch."""
    provider_id = str(record.provider_id or "").strip()
    if not provider_id:
        # Metadata written by older versions may not have a provider id.  Its
        # parent is still the generator directory in the current file layout.
        provider_id = record.folder.parent.name
    return provider_id.casefold()


def group_generation_providers(
    records: list[GenerationRecord],
) -> list[GenerationProviderFolder]:
    """Build newest-first generator folders from existing generation records."""
    grouped: dict[str, list[GenerationRecord]] = {}
    for record in records:
        grouped.setdefault(generation_provider_key(record), []).append(record)
    folders: list[GenerationProviderFolder] = []
    for key, provider_records in grouped.items():
        newest_first = sorted(
            provider_records,
            key=lambda item: item.created_at or item.folder.name,
            reverse=True,
        )
        latest = newest_first[0]
        provider_id = str(latest.provider_id or latest.folder.parent.name)
        provider_name = str(latest.provider_name or provider_id)
        folders.append(
            GenerationProviderFolder(
                key=key,
                provider_id=provider_id,
                provider_name=provider_name,
                records=tuple(newest_first),
                model_count=len({generation_model_key(record) for record in newest_first}),
                image_count=sum(len(record.images) for record in newest_first),
                latest_at=latest.created_at,
            )
        )
    folders.sort(key=lambda item: (item.latest_at, item.provider_name.casefold()), reverse=True)
    return folders


def generation_model_key(record: GenerationRecord) -> str:
    """Keep renamed or duplicated display names separated by model identity."""
    raw_path = str(record.model_path or "").strip()
    if raw_path:
        try:
            return f"path:{Path(raw_path).expanduser().resolve()}".casefold()
        except OSError:
            return f"path:{raw_path}".casefold()
    return f"name:{record.provider_id}:{record.model_name}".casefold()


def group_generation_records(
    records: list[GenerationRecord],
) -> list[GenerationModelFolder]:
    """Build newest-first automatic model folders without changing files."""
    grouped: dict[str, list[GenerationRecord]] = {}
    for record in records:
        grouped.setdefault(generation_model_key(record), []).append(record)
    folders: list[GenerationModelFolder] = []
    for key, model_records in grouped.items():
        newest_first = sorted(
            model_records,
            key=lambda item: item.created_at or item.folder.name,
            reverse=True,
        )
        latest = newest_first[0]
        cover = next(
            (path for record in newest_first for path in record.images if path.is_file()),
            None,
        )
        folders.append(
            GenerationModelFolder(
                key=key,
                model_name=latest.model_name,
                model_path=latest.model_path,
                provider_id=latest.provider_id,
                provider_name=latest.provider_name,
                records=tuple(newest_first),
                image_count=sum(len(record.images) for record in newest_first),
                cover_image=cover,
                latest_at=latest.created_at,
            )
        )
    folders.sort(key=lambda item: (item.latest_at, item.model_name.casefold()), reverse=True)
    return folders


def generation_tools(registry: ToolRegistry) -> list[ToolSpec]:
    return [
        tool
        for tool in registry.enabled()
        if "image_generation" in tool.capabilities
    ]


def load_generation_history(root: Path, *, limit: int = 200) -> list[GenerationRecord]:
    history_root = root.resolve() / "data" / "generations"
    if not history_root.is_dir():
        return []
    # History can grow into thousands of image batches.  Sort inexpensive file
    # metadata first, then decode only the newest records requested by the UI.
    # This keeps a page refresh responsive without moving or rewriting history.
    try:
        metadata_paths = sorted(
            history_root.rglob("generation*.json"),
            key=lambda path: path.stat().st_mtime,
            reverse=True,
        )
    except OSError:
        metadata_paths = list(history_root.rglob("generation*.json"))
    records = []
    for metadata_path in metadata_paths:
        record = GenerationRecord.from_metadata(metadata_path)
        if record is not None:
            records.append(record)
        if len(records) >= max(1, int(limit)):
            break
    records.sort(key=lambda item: item.created_at or item.folder.name, reverse=True)
    return records


def build_generation_plan(
    tool: ToolSpec,
    *,
    model_name: str,
    model_path: str,
    prompt: str,
    image_count: int,
    steps: int,
    seed: int,
    sampler: str,
    aspect_ratio: str,
    extra_arguments: dict[str, Any] | None = None,
) -> ExecutionPlan:
    if "image_generation" not in tool.capabilities:
        raise ValueError(f"{tool.name} is not registered for image generation.")
    safe_name = model_name.strip() or Path(model_path).name
    arguments: dict[str, Any] = {
        "model_name": safe_name,
        "model_path": model_path,
        "prompt": prompt.strip(),
        "image_count": int(image_count),
        "steps": int(steps),
        "seed": int(seed),
        "sampler": sampler,
        "aspect_ratio": aspect_ratio,
    }
    if extra_arguments:
        arguments.update(extra_arguments)
    return ExecutionPlan(
        request=f"Generate {image_count} image(s) with {safe_name}",
        summary=f"Generate {image_count} image(s) using {tool.name} and {safe_name}.",
        steps=[
            PlanStep(
                tool_id=tool.id,
                title="Generate images",
                description=f"Create a reproducible image batch with {safe_name}.",
                arguments=arguments,
            )
        ],
        requires_confirmation=tool.requires_confirmation,
        confirmation_reason=(
            "This generator is configured to require approval before it runs."
            if tool.requires_confirmation
            else ""
        ),
        project_name=f"{safe_name} generation",
    )


def combine_generation_plans(
    plans: list[ExecutionPlan],
    *,
    display_seconds: int = 5,
    show_labels: bool = True,
    loop: bool = False,
) -> ExecutionPlan:
    """Combine per-model image plans into one sequential presentation cycle."""
    usable = [plan for plan in plans if plan.steps]
    if not usable:
        raise ValueError("A generation cycle needs at least one model.")
    model_names = [
        str(plan.steps[0].arguments.get("model_name", plan.project_name))
        for plan in usable
    ]
    cycle_settings = {
        "display_seconds": max(1, int(display_seconds)),
        "show_labels": bool(show_labels),
        "loop": bool(loop),
        "models": model_names,
    }
    reasons = [plan.confirmation_reason for plan in usable if plan.confirmation_reason]
    return ExecutionPlan(
        request=f"Generate a presentation cycle with {len(usable)} models.",
        summary=(
            f"Generate images sequentially with {len(usable)} models, then play them "
            f"for {cycle_settings['display_seconds']} seconds each"
            + (" with model labels." if show_labels else ".")
        ),
        steps=[step for plan in usable for step in plan.steps],
        requires_confirmation=any(plan.requires_confirmation for plan in usable),
        confirmation_reason="; ".join(dict.fromkeys(reasons)),
        project_name="Generation Cycle",
    )