adam/generations.py

from __future__ import annotations

import json
import re
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 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


_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)
        ignored = {"a", "an", "the", "of", "image", "picture", "model", "ddpm", "flow", "matching", "lora"}
        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.
    """
    request = " ".join(text.strip().split())
    if not request or not re.search(r"\b(generate|create|make)\b", request, re.I):
        return None
    if not re.search(r"\b(image|images|picture|pictures)\b", request, re.I):
        return None

    provider_hint = ""
    provider_match = re.search(
        r"\b(ddpm|ddim|flow(?:\s+matching)?|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 "flow" if hint.startswith("flow") 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"\bflow(?:\s+match(?:ing)?)?\s*$", subject, re.I):
            provider_hint = "flow"
        elif re.search(r"\bddpm\s*$", subject, re.I):
            provider_hint = "ddpm"

    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:
    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

    @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(
            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", "")),
        )


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 []
    records = [
        record
        for metadata_path in history_root.rglob("generation*.json")
        for record in [GenerationRecord.from_metadata(metadata_path)]
        if record is not None
    ]
    records.sort(key=lambda item: item.created_at or item.folder.name, reverse=True)
    return records[: max(1, int(limit))]


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",
    )