StyleRef-styleref-comfyui/server_routes.py

"""
Backend routes for the Load node's search dialog.

The browser cannot call styleref.io directly from the ComfyUI frontend (CORS
aside, an auth token must never reach page JS), so the widget asks the ComfyUI
server, which asks StyleRef using the same client the nodes use.

Everything here is optional. If ComfyUI's server module isn't importable — a
headless test run, an API-only embed, a future frontend rewrite — importing this
fails, `__init__` catches it, and the nodes carry on with the plain text field.
The JS side probes /styleref/search once at startup and hides its button when
these routes are absent, so both halves degrade together.
"""

from __future__ import annotations

import asyncio
import os

from aiohttp import web
from server import PromptServer

import styleref_api as api
from styleref_api import StyleRefError
from styleref_auth import auth_status, clear_credentials, credentials_path, run_login_flow
from styleref_nodes.load import bump_refresh, fetch_style


def _card(entry: dict) -> dict:
    """The picker-row shape. heroImage feeds the row thumbnail."""
    return {
        "slug": entry.get("slug"),
        "name": entry.get("name"),
        "author": entry.get("author"),
        "category": entry.get("category"),
        "tags": (entry.get("tags") or [])[:4],
        "heroImage": entry.get("heroImage"),
        "url": entry.get("url"),
    }


def _int_param(request: web.Request, name: str, default: int) -> int:
    """A query int that tolerates junk — a bad value is a default, not a 400."""
    try:
        return int(request.query.get(name, str(default)))
    except ValueError:
        return default


@PromptServer.instance.routes.get("/styleref/search")
async def styleref_search(request: web.Request) -> web.Response:
    """GET /styleref/search?query=…&category=…&sort=…&limit=…&offset=… → {styles, hasMore}"""
    query = request.query.get("query", "").strip()
    category = request.query.get("category", "").strip() or None
    sort = request.query.get("sort", "").strip() or None
    if sort not in (None, "recent", "popular"):
        sort = None
    limit = _int_param(request, "limit", 8)
    offset = _int_param(request, "offset", 0)

    try:
        payload = api.search_styles(
            query, category=category, sort=sort, limit=limit, offset=offset
        )
    except StyleRefError as err:
        # 200 with an error field: the dialog renders the message inline rather
        # than the frontend surfacing a generic red toast with no detail.
        return web.json_response({"styles": [], "error": err.message})

    return web.json_response(
        {
            "styles": [_card(s) for s in (payload.get("styles") or [])],
            "availableCategories": payload.get("availableCategories") or [],
            # Drives the dialog's Next button without it having to guess from a
            # short page — a full last page is indistinguishable otherwise.
            "hasMore": bool(payload.get("hasMore")),
        }
    )


@PromptServer.instance.routes.get("/styleref/saved")
async def styleref_saved(request: web.Request) -> web.Response:
    """
    The picker's "Saved" sort: gallery styles this user saved. Signed-in only,
    and reported as such rather than as an error — signing out is not a failure.
    """
    signed_in, message = auth_status()
    if not signed_in:
        return web.json_response({"styles": [], "signedIn": False, "message": message})

    query = request.query.get("query", "").strip()
    category = request.query.get("category", "").strip() or None

    try:
        payload = api.list_saved_styles(
            query,
            category=category,
            limit=_int_param(request, "limit", 8),
            offset=_int_param(request, "offset", 0),
        )
    except StyleRefError as err:
        # A 401 here means stored credentials that no longer work (expired, or
        # revoked on styleref.io). That is the signed-out state as far as the UI
        # is concerned, so report it as one and let the dialog offer Sign in —
        # echoing the server's raw 401 told people to set an Authorization header,
        # which is not something anyone can do from inside ComfyUI.
        if err.needs_login:
            return web.json_response({"styles": [], "signedIn": False})
        return web.json_response({"styles": [], "signedIn": True, "error": err.message})

    return web.json_response(
        {
            "signedIn": True,
            "styles": [_card(s) for s in (payload.get("styles") or [])],
            "hasMore": bool(payload.get("hasMore")),
        }
    )


@PromptServer.instance.routes.get("/styleref/my-styles")
async def styleref_my_styles(request: web.Request) -> web.Response:
    """The caller's own library, shown above gallery results when signed in."""
    signed_in, message = auth_status()
    if not signed_in:
        return web.json_response({"styles": [], "signedIn": False, "message": message})

    cursor = request.query.get("cursor", "").strip() or None
    query = request.query.get("query", "").strip()
    try:
        payload = api.list_my_styles(
            limit=_int_param(request, "limit", 8), cursor=cursor, query=query
        )
    except StyleRefError as err:
        # A 401 here means stored credentials that no longer work (expired, or
        # revoked on styleref.io). That is the signed-out state as far as the UI
        # is concerned, so report it as one and let the dialog offer Sign in —
        # echoing the server's raw 401 told people to set an Authorization header,
        # which is not something anyone can do from inside ComfyUI.
        if err.needs_login:
            return web.json_response({"styles": [], "signedIn": False})
        return web.json_response({"styles": [], "signedIn": True, "error": err.message})

    return web.json_response(
        {
            "signedIn": True,
            "styles": [
                {
                    "slug": s.get("ref"),
                    "name": s.get("name"),
                    "author": None,
                    "tags": [],
                    "visibility": s.get("visibility"),
                    # null → never generated on styleref.io, so loading it would
                    # 409. The dialog badges these instead of offering them as
                    # ready to use.
                    "lastGeneratedAt": s.get("lastGeneratedAt"),
                }
                for s in (payload.get("styles") or [])
            ],
            # Passed back as ?cursor= for the dialog's Next button.
            "nextCursor": payload.get("nextCursor"),
        }
    )


@PromptServer.instance.routes.get("/styleref/resolve")
async def styleref_resolve(request: web.Request) -> web.Response:
    """
    GET /styleref/resolve?ref=… → {ok, name, url}

    Backs the name line under the Load node's `style_ref` field. A slug is
    opaque — `72e1zdae-e2d54a18d090` tells you nothing about which style is
    wired in, and until the graph is queued the node cannot say either. This
    answers the question the moment the field changes.

    It goes through the same client (and therefore the same session cache) the
    node uses, so a resolve for a style that is about to be loaded costs one
    request and the queue after it costs none.
    """
    ref = request.query.get("ref", "").strip()
    if not ref:
        return web.json_response({"ok": False})

    try:
        style = fetch_style(ref)
    except StyleRefError as err:
        # 200 with an error field, as with search: an unresolvable ref is
        # ordinary (a half-typed slug, a private style while signed out) and
        # belongs inline on the node, not in a red toast.
        return web.json_response({"ok": False, "error": err.message})

    # No author: the structured lane carries name/sections/referenceImages only,
    # so the node body's attribution line is where a handle would come from.
    return web.json_response(
        {"ok": True, "name": style.get("name"), "url": style.get("url")}
    )


@PromptServer.instance.routes.post("/styleref/refresh")
async def styleref_refresh(request: web.Request) -> web.Response:
    """
    The ↻ Refresh style button's backend: forget the cached style and advance
    the generation counter IS_CHANGED reads, so the next queue re-downloads.
    """
    try:
        payload = await request.json()
    except Exception:  # noqa: BLE001 - malformed body is a client bug, not a crash
        payload = {}
    ref = str(payload.get("ref") or "").strip()
    if not ref:
        return web.json_response({"ok": False, "error": "No style_ref to refresh."})
    generation = bump_refresh(ref)
    return web.json_response({"ok": True, "generation": generation})


# ── auth: the Login node's button backend ────────────────────────────────────
# The Login node's `action` dropdown stays as the queue-driven fallback (it is
# the only path for API-only installs); these routes power the node's buttons.


def _status_payload() -> dict:
    signed_in, message = auth_status()
    return {"signedIn": signed_in, "message": message}


@PromptServer.instance.routes.get("/styleref/auth/status")
async def styleref_auth_status(_request: web.Request) -> web.Response:
    return web.json_response(_status_payload())


@PromptServer.instance.routes.post("/styleref/auth/signin")
async def styleref_auth_signin(_request: web.Request) -> web.Response:
    """
    Run the OAuth loopback flow. Blocking by nature (it waits for the browser
    redirect), so it runs in an executor; the JS button shows a waiting state.
    """
    loop = asyncio.get_event_loop()
    ok, message = await loop.run_in_executor(None, run_login_flow, api.api_base())
    return web.json_response({"ok": ok, "message": message, **_status_payload()})


@PromptServer.instance.routes.post("/styleref/auth/signout")
async def styleref_auth_signout(_request: web.Request) -> web.Response:
    removed = clear_credentials()
    message = (
        f"Signed out — removed {credentials_path()}."
        if removed
        else "No stored credentials to remove."
    )
    if os.environ.get("STYLEREF_TOKEN"):
        message += (
            " Note: STYLEREF_TOKEN is still set in this environment and takes "
            "precedence, so you remain signed in via that token."
        )
    return web.json_response({"ok": True, "message": message, **_status_payload()})