comfy-discord-web/commands/presets.py

"""
commands/presets.py
===================

Named workflow preset commands for the Discord ComfyUI bot.

A preset is a saved snapshot of the current workflow template and runtime
state (prompt, negative_prompt, input_image, seed).  Presets make it easy
to switch between different setups (e.g. "portrait", "landscape", "anime")
with a single command.

All sub-commands are accessed through the single ``ttr!preset`` command:

    ttr!preset save   <name> [description:<text>]  — capture current workflow + state
    ttr!preset load   <name>                        — restore workflow + state
    ttr!preset list                                 — list all saved presets
    ttr!preset view   <name>                        — show preset details
    ttr!preset delete <name>                        — permanently remove a preset
    ttr!preset save-last <name> [description:<text>] — save last generation as preset
"""

from __future__ import annotations

import logging

from discord.ext import commands

from discord_utils import require_comfy_client
from preset_manager import PresetManager


logger = logging.getLogger(__name__)


def _parse_name_and_description(args: str) -> tuple[str, str | None]:
    """
    Split ``<name> [description:<text>]`` into (name, description).

    The name is the first whitespace-delimited token.  Everything after
    ``description:`` (case-insensitive) in the remaining text is the
    description.  Returns (name, None) if no description keyword is found.
    """
    parts = args.strip().split(maxsplit=1)
    name = parts[0] if parts else ""
    description: str | None = None
    if len(parts) > 1:
        rest = parts[1]
        lower = rest.lower()
        idx = lower.find("description:")
        if idx >= 0:
            description = rest[idx + len("description:"):].strip() or None
    return name, description


def setup_preset_commands(bot):
    """
    Register preset commands with the bot.

    Parameters
    ----------
    bot : commands.Bot
        The Discord bot instance.
    """
    preset_manager = PresetManager()

    @bot.command(name="preset", extras={"category": "Presets"})
    @require_comfy_client
    async def preset_command(ctx: commands.Context, *, args: str = "") -> None:
        """
        Save, load, list, view, or delete named workflow presets.

        A preset captures the current workflow template and all runtime
        state changes (prompt, negative_prompt, input_image, seed) under a
        short name.  Load it later to restore everything in one step.

        Usage:
            ttr!preset save   <name> [description:<text>]  — save current workflow + state
            ttr!preset load   <name>                        — restore workflow + state
            ttr!preset list                                 — list all saved presets
            ttr!preset view   <name>                        — show preset details
            ttr!preset delete <name>                        — permanently delete a preset
            ttr!preset save-last <name> [description:<text>] — save last generation as preset

        Names may only contain letters, digits, hyphens, and underscores.

        Examples:
            ttr!preset save portrait description:studio lighting style
            ttr!preset load portrait
            ttr!preset list
            ttr!preset view portrait
            ttr!preset delete portrait
            ttr!preset save-last my-last
        """
        parts = args.strip().split(maxsplit=1)
        subcommand = parts[0].lower() if parts else ""
        rest = parts[1].strip() if len(parts) > 1 else ""

        if subcommand == "save":
            name, description = _parse_name_and_description(rest)
            await _preset_save(ctx, bot, preset_manager, name, description)
        elif subcommand == "load":
            await _preset_load(ctx, bot, preset_manager, rest.split()[0] if rest.split() else "")
        elif subcommand == "list":
            await _preset_list(ctx, preset_manager)
        elif subcommand == "view":
            await _preset_view(ctx, preset_manager, rest.split()[0] if rest.split() else "")
        elif subcommand == "delete":
            await _preset_delete(ctx, preset_manager, rest.split()[0] if rest.split() else "")
        elif subcommand == "save-last":
            name, description = _parse_name_and_description(rest)
            await _preset_save_last(ctx, preset_manager, name, description)
        else:
            await ctx.reply(
                "Usage: `ttr!preset <save|load|list|view|delete|save-last> [name]`\n"
                "Run `ttr!help preset` for full details.",
                mention_author=False,
            )


async def _preset_save(
    ctx: commands.Context, bot, preset_manager: PresetManager, name: str,
    description: str | None = None,
) -> None:
    """Handle ttr!preset save <name> [description:<text>]."""
    if not name:
        await ctx.reply(
            "Please provide a name. Example: `ttr!preset save portrait`",
            mention_author=False,
        )
        return

    if not PresetManager.is_valid_name(name):
        await ctx.reply(
            "Invalid name. Use only letters, digits, hyphens, and underscores (max 64 chars).",
            mention_author=False,
        )
        return

    try:
        workflow_template = bot.comfy.get_workflow_template()
        state = bot.comfy.get_workflow_current_changes()
        preset_manager.save(name, workflow_template, state, description=description)

        # Build a summary of what was saved
        has_workflow = workflow_template is not None
        state_parts = []
        if state.get("prompt"):
            state_parts.append("prompt")
        if state.get("negative_prompt"):
            state_parts.append("negative_prompt")
        if state.get("input_image"):
            state_parts.append("input_image")
        if state.get("seed") is not None:
            state_parts.append(f"seed={state['seed']}")

        summary_parts = []
        if has_workflow:
            summary_parts.append("workflow template")
        summary_parts.extend(state_parts)
        summary = ", ".join(summary_parts) if summary_parts else "empty state"

        desc_note = f"\n> {description}" if description else ""
        await ctx.reply(
            f"Preset **{name}** saved ({summary}).{desc_note}",
            mention_author=False,
        )
    except Exception as exc:
        logger.exception("Failed to save preset '%s'", name)
        await ctx.reply(
            f"Failed to save preset: {type(exc).__name__}: {exc}",
            mention_author=False,
        )


async def _preset_load(
    ctx: commands.Context, bot, preset_manager: PresetManager, name: str
) -> None:
    """Handle ttr!preset load <name>."""
    if not name:
        await ctx.reply(
            "Please provide a name. Example: `ttr!preset load portrait`",
            mention_author=False,
        )
        return

    data = preset_manager.load(name)
    if data is None:
        presets = preset_manager.list_presets()
        hint = f" Available: {', '.join(presets)}" if presets else " No presets saved yet."
        await ctx.reply(
            f"Preset **{name}** not found.{hint}",
            mention_author=False,
        )
        return

    try:
        restored: list[str] = []

        # Restore workflow template if present
        workflow = data.get("workflow")
        if workflow is not None:
            bot.comfy.set_workflow(workflow)
            restored.append("workflow template")

        # Restore state changes
        state = data.get("state", {})
        if state:
            bot.comfy.set_workflow_current_changes(state)
            if state.get("prompt"):
                restored.append("prompt")
            if state.get("negative_prompt"):
                restored.append("negative_prompt")
            if state.get("input_image"):
                restored.append("input_image")
            if state.get("seed") is not None:
                restored.append(f"seed={state['seed']}")

        summary = ", ".join(restored) if restored else "nothing (preset was empty)"
        description = data.get("description")
        desc_note = f"\n> {description}" if description else ""
        await ctx.reply(
            f"Preset **{name}** loaded ({summary}).{desc_note}",
            mention_author=False,
        )
    except Exception as exc:
        logger.exception("Failed to load preset '%s'", name)
        await ctx.reply(
            f"Failed to load preset: {type(exc).__name__}: {exc}",
            mention_author=False,
        )


async def _preset_view(
    ctx: commands.Context, preset_manager: PresetManager, name: str
) -> None:
    """Handle ttr!preset view <name>."""
    if not name:
        await ctx.reply(
            "Please provide a name. Example: `ttr!preset view portrait`",
            mention_author=False,
        )
        return

    data = preset_manager.load(name)
    if data is None:
        await ctx.reply(f"Preset **{name}** not found.", mention_author=False)
        return

    lines = [f"**Preset: {name}**"]
    if data.get("description"):
        lines.append(f"> {data['description']}")
    if data.get("owner"):
        lines.append(f"Owner: {data['owner']}")

    state = data.get("state", {})
    if state.get("prompt"):
        # Truncate long prompts
        p = str(state["prompt"])
        if len(p) > 200:
            p = p[:197] + "…"
        lines.append(f"**Prompt:** {p}")
    if state.get("negative_prompt"):
        np = str(state["negative_prompt"])
        if len(np) > 100:
            np = np[:97] + "…"
        lines.append(f"**Negative:** {np}")
    if state.get("seed") is not None:
        seed_note = " (random)" if state["seed"] == -1 else ""
        lines.append(f"**Seed:** {state['seed']}{seed_note}")

    other = {k: v for k, v in state.items() if k not in ("prompt", "negative_prompt", "seed", "input_image")}
    if other:
        other_str = ", ".join(f"{k}={v}" for k, v in other.items())
        lines.append(f"**Other:** {other_str[:200]}")

    if data.get("workflow") is not None:
        lines.append("_(includes workflow template)_")
    else:
        lines.append("_(no workflow template — load separately)_")

    await ctx.reply("\n".join(lines), mention_author=False)


async def _preset_list(ctx: commands.Context, preset_manager: PresetManager) -> None:
    """Handle ttr!preset list."""
    presets = preset_manager.list_preset_details()
    if not presets:
        await ctx.reply(
            "No presets saved yet. Use `ttr!preset save <name>` to create one.",
            mention_author=False,
        )
        return

    lines = [f"**Saved presets** ({len(presets)})"]
    for p in presets:
        entry = f"  • {p['name']}"
        if p.get("description"):
            entry += f" — {p['description']}"
        lines.append(entry)
    lines.append("\nUse `ttr!preset load <name>` to restore one.")
    await ctx.reply("\n".join(lines), mention_author=False)


async def _preset_delete(
    ctx: commands.Context, preset_manager: PresetManager, name: str
) -> None:
    """Handle ttr!preset delete <name>."""
    if not name:
        await ctx.reply(
            "Please provide a name. Example: `ttr!preset delete portrait`",
            mention_author=False,
        )
        return

    deleted = preset_manager.delete(name)
    if deleted:
        await ctx.reply(f"Preset **{name}** deleted.", mention_author=False)
    else:
        await ctx.reply(
            f"Preset **{name}** not found.",
            mention_author=False,
        )


async def _preset_save_last(
    ctx: commands.Context, preset_manager: PresetManager, name: str,
    description: str | None = None,
) -> None:
    """Handle ttr!preset save-last <name> [description:<text>]."""
    if not name:
        await ctx.reply(
            "Please provide a name. Example: `ttr!preset save-last my-last`",
            mention_author=False,
        )
        return

    if not PresetManager.is_valid_name(name):
        await ctx.reply(
            "Invalid name. Use only letters, digits, hyphens, and underscores (max 64 chars).",
            mention_author=False,
        )
        return

    from generation_db import get_history as db_get_history
    history = db_get_history(limit=1)
    if not history:
        await ctx.reply(
            "No generation history found. Generate something first!",
            mention_author=False,
        )
        return

    last = history[0]
    overrides = last.get("overrides") or {}
    try:
        preset_manager.save(name, None, overrides, description=description)
        desc_note = f"\n> {description}" if description else ""
        await ctx.reply(
            f"Preset **{name}** saved from last generation.{desc_note}\n"
            "Note: workflow template not included — load it separately before generating.",
            mention_author=False,
        )
    except ValueError as exc:
        await ctx.reply(str(exc), mention_author=False)
    except Exception as exc:
        logger.exception("Failed to save preset '%s' from history", name)
        await ctx.reply(
            f"Failed to save preset: {type(exc).__name__}: {exc}",
            mention_author=False,
        )