comfy-discord-web/preset_manager.py

"""
preset_manager.py
=================

Preset management for the Discord ComfyUI bot.

A preset is a named snapshot of the current workflow template and runtime
state (prompt, negative_prompt, input_image, seed).  Presets are stored as
individual JSON files inside a ``presets/`` directory so they can be
inspected, backed up, or shared manually.

Usage via bot commands:
    ttr!preset save   <name>  — capture current workflow + state
    ttr!preset load   <name>  — restore workflow + state from snapshot
    ttr!preset list           — list all saved presets
    ttr!preset delete <name>  — permanently remove a preset
"""

from __future__ import annotations

import json
import logging
import re
from pathlib import Path
from typing import Any, Optional


logger = logging.getLogger(__name__)

# Only allow alphanumeric characters, hyphens, and underscores in preset names.
_SAFE_NAME_RE = re.compile(r"^[a-zA-Z0-9_-]{1,64}$")


class PresetManager:
    """
    Manages named workflow presets on disk.

    Each preset is stored as ``<presets_dir>/<name>.json`` and contains:
    - ``name``:     the preset name
    - ``workflow``: the full ComfyUI workflow template dict (may be null)
    - ``state``:    the runtime state changes (prompt, negative_prompt,
                    input_image, seed)

    Parameters
    ----------
    presets_dir : str
        Directory where preset files are stored.  Created automatically if
        it does not exist.
    """

    def __init__(self, presets_dir: str = "presets") -> None:
        self.presets_dir = Path(presets_dir)
        self.presets_dir.mkdir(parents=True, exist_ok=True)

    # ------------------------------------------------------------------
    # Internal helpers
    # ------------------------------------------------------------------

    @staticmethod
    def is_valid_name(name: str) -> bool:
        """Return True if the name only contains safe characters."""
        return bool(_SAFE_NAME_RE.match(name))

    def _path(self, name: str) -> Path:
        """Return the file path for a preset by name (no validation)."""
        return self.presets_dir / f"{name}.json"

    # ------------------------------------------------------------------
    # Public API
    # ------------------------------------------------------------------

    def save(
        self,
        name: str,
        workflow_template: Optional[dict[str, Any]],
        state: dict[str, Any],
        owner: Optional[str] = None,
        description: Optional[str] = None,
    ) -> None:
        """
        Save a preset to disk.

        Parameters
        ----------
        name : str
            The preset name (alphanumeric, hyphens, underscores only).
        workflow_template : Optional[dict]
            The current workflow template, or None if none is loaded.
        state : dict
            The current runtime state from WorkflowStateManager.get_changes().
        owner : Optional[str]
            The user label of the preset creator.  Stored for access control.
        description : Optional[str]
            A human-readable description of what this preset does.

        Raises
        ------
        ValueError
            If the name contains invalid characters.
        OSError
            If the file cannot be written.
        """
        if not self.is_valid_name(name):
            raise ValueError(
                f"Invalid preset name '{name}'. "
                "Use only letters, digits, hyphens, and underscores (max 64 chars)."
            )
        data: dict[str, Any] = {"name": name, "workflow": workflow_template, "state": state}
        if owner is not None:
            data["owner"] = owner
        if description is not None:
            data["description"] = description
        path = self._path(name)
        with open(path, "w", encoding="utf-8") as f:
            json.dump(data, f, indent=2)
        logger.info("Saved preset '%s' to %s", name, path)

    def load(self, name: str) -> Optional[dict[str, Any]]:
        """
        Load a preset from disk.

        Parameters
        ----------
        name : str
            The preset name.

        Returns
        -------
        Optional[dict]
            The preset data dict, or None if the preset does not exist.
        """
        if not self.is_valid_name(name):
            return None
        path = self._path(name)
        if not path.exists():
            return None
        try:
            with open(path, "r", encoding="utf-8") as f:
                return json.load(f)
        except Exception as exc:
            logger.warning("Failed to load preset '%s': %s", name, exc)
            return None

    def delete(self, name: str) -> bool:
        """
        Delete a preset file.

        Parameters
        ----------
        name : str
            The preset name.

        Returns
        -------
        bool
            True if the preset existed and was deleted, False otherwise.
        """
        if not self.is_valid_name(name):
            return False
        path = self._path(name)
        if path.exists():
            path.unlink()
            logger.info("Deleted preset '%s'", name)
            return True
        return False

    def list_presets(self) -> list[str]:
        """
        List all saved preset names, sorted alphabetically.

        Returns
        -------
        list[str]
            Sorted list of preset names (without the .json extension).
        """
        return sorted(p.stem for p in self.presets_dir.glob("*.json"))

    def list_preset_details(self) -> list[dict[str, Any]]:
        """
        List all presets with their metadata, sorted alphabetically by name.

        Returns
        -------
        list[dict]
            Each entry has ``"name"``, ``"owner"`` (may be None), and
            ``"description"`` (may be None).
        """
        result = []
        for p in sorted(self.presets_dir.glob("*.json"), key=lambda x: x.stem):
            owner = None
            description = None
            try:
                with open(p, "r", encoding="utf-8") as f:
                    data = json.load(f)
                owner = data.get("owner")
                description = data.get("description")
            except Exception:
                pass
            result.append({"name": p.stem, "owner": owner, "description": description})
        return result

    def exists(self, name: str) -> bool:
        """Return True if a preset with this name exists on disk."""
        if not self.is_valid_name(name):
            return False
        return self._path(name).exists()