comfy-discord-web/web/routers/presets_router.py

"""CRUD for workflow presets via /api/presets"""
from __future__ import annotations

from typing import Optional

from fastapi import APIRouter, Depends, HTTPException
from pydantic import BaseModel

from web.auth import require_auth
from web.deps import get_comfy, get_user_registry

router = APIRouter()


class SavePresetRequest(BaseModel):
    name: str
    description: Optional[str] = None


class SaveFromHistoryRequest(BaseModel):
    name: str
    description: Optional[str] = None


def _get_pm():
    from web.deps import get_bot
    bot = get_bot()
    pm = getattr(bot, "preset_manager", None) if bot else None
    if pm is None:
        from preset_manager import PresetManager
        pm = PresetManager()
    return pm


@router.get("")
async def list_presets(_: dict = Depends(require_auth)):
    pm = _get_pm()
    return {"presets": pm.list_preset_details()}


@router.post("")
async def save_preset(body: SavePresetRequest, user: dict = Depends(require_auth)):
    """Capture the user's overrides + workflow template as a named preset."""
    user_label: str = user["sub"]
    registry = get_user_registry()
    pm = _get_pm()

    if registry:
        workflow_template = registry.get_workflow_template(user_label)
        overrides = registry.get_state_manager(user_label).get_overrides()
    else:
        comfy = get_comfy()
        if comfy is None:
            raise HTTPException(503, "ComfyUI not available")
        workflow_template = comfy.get_workflow_template()
        overrides = comfy.state_manager.get_overrides()

    try:
        pm.save(body.name, workflow_template, overrides, owner=user_label, description=body.description)
    except ValueError as exc:
        raise HTTPException(400, str(exc))
    return {"ok": True, "name": body.name}


@router.get("/{name}")
async def get_preset(name: str, _: dict = Depends(require_auth)):
    pm = _get_pm()
    data = pm.load(name)
    if data is None:
        raise HTTPException(404, "Preset not found")
    return data


@router.post("/{name}/load")
async def load_preset(name: str, user: dict = Depends(require_auth)):
    """Restore overrides (and optionally workflow template) from a preset into the user's state."""
    pm = _get_pm()
    data = pm.load(name)
    if data is None:
        raise HTTPException(404, "Preset not found")

    user_label: str = user["sub"]
    registry = get_user_registry()

    if registry:
        wf = data.get("workflow")
        if wf:
            registry.set_workflow(user_label, wf, name)
        else:
            # No workflow in preset — just clear overrides and restore state
            registry.get_state_manager(user_label).clear_overrides()
        state = data.get("state", {})
        sm = registry.get_state_manager(user_label)
        for k, v in state.items():
            if v is not None:
                sm.set_override(k, v)
    else:
        comfy = get_comfy()
        if comfy is None:
            raise HTTPException(503, "ComfyUI not available")
        comfy.state_manager.clear_overrides()
        state = data.get("state", {})
        for k, v in state.items():
            if v is not None:
                comfy.state_manager.set_override(k, v)
        wf = data.get("workflow")
        if wf:
            comfy.workflow_manager.set_workflow_template(wf)

    return {"ok": True, "name": name, "overrides_restored": list(data.get("state", {}).keys())}


@router.delete("/{name}")
async def delete_preset(name: str, user: dict = Depends(require_auth)):
    pm = _get_pm()
    data = pm.load(name)
    if data is None:
        raise HTTPException(404, "Preset not found")

    user_label: str = user["sub"]
    is_admin = user.get("admin") is True
    owner = data.get("owner")
    if owner is not None and owner != user_label and not is_admin:
        raise HTTPException(403, "You do not have permission to delete this preset")

    pm.delete(name)
    return {"ok": True}


@router.post("/from-history/{prompt_id}")
async def save_preset_from_history(
    prompt_id: str,
    body: SaveFromHistoryRequest,
    user: dict = Depends(require_auth),
):
    """Create a preset from a past generation's overrides."""
    from generation_db import get_generation_full

    gen = get_generation_full(prompt_id)
    if gen is None:
        raise HTTPException(404, "Generation not found")

    user_label: str = user["sub"]
    is_admin = user.get("admin") is True
    if not is_admin and gen.get("user_label") != user_label:
        raise HTTPException(404, "Generation not found")

    pm = _get_pm()
    try:
        pm.save(body.name, None, gen["overrides"], owner=user_label, description=body.description)
    except ValueError as exc:
        raise HTTPException(400, str(exc))
    return {"ok": True, "name": body.name}