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
import time

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


# Slugs of the signed-in user's own published styles. A style of yours that is
# in the gallery is in *both* lanes of the picker otherwise — the same row
# twice, once under "My styles" and once under "Gallery" — and deduping in the
# frontend only works for the page currently on screen. Walked once and cached:
# the answer changes when you publish something, not between keystrokes.
_OWN_SLUGS: dict[str, object] = {"at": 0.0, "slugs": set()}
_OWN_SLUGS_TTL_S = 300
_OWN_SLUGS_PAGE = 50
_OWN_SLUGS_MAX = 200


def forget_own_slugs() -> None:
    """Sign-in and sign-out both change whose library this is."""
    _OWN_SLUGS["at"] = 0.0
    _OWN_SLUGS["slugs"] = set()


def own_slugs() -> set[str]:
    """
    Every share slug this account owns, or an empty set when signed out.

    Only published styles have a slug, and only a slug can collide with a
    gallery row, so an unpublished library costs one request and returns
    nothing to filter.
    """
    signed_in, _message = auth_status()
    if not signed_in:
        return set()

    now = time.monotonic()
    if now - float(_OWN_SLUGS["at"]) < _OWN_SLUGS_TTL_S:
        return _OWN_SLUGS["slugs"]  # type: ignore[return-value]

    slugs: set[str] = set()
    cursor: str | None = None
    try:
        while len(slugs) < _OWN_SLUGS_MAX:
            payload = api.list_my_styles(limit=_OWN_SLUGS_PAGE, cursor=cursor)
            for entry in payload.get("styles") or []:
                slug = entry.get("slug")
                if slug:
                    slugs.add(slug)
            cursor = payload.get("nextCursor")
            if not cursor:
                break
    except StyleRefError:
        # Filtering is a nicety; a failure here must not break search. An empty
        # set means the gallery lane is unfiltered, which is the old behaviour.
        return set()

    _OWN_SLUGS["at"] = now
    _OWN_SLUGS["slugs"] = slugs
    return slugs


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})

    # The user's own styles already have their own lane above this one. A page
    # can come back one or two rows short as a result — cheaper than showing the
    # same style twice and letting the user wonder which one to click.
    mine = own_slugs()
    rows = [s for s in (payload.get("styles") or []) if s.get("slug") not in mine]

    return web.json_response(
        {
            "styles": [_card(s) for s in rows],
            "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"),
                    # No author line: this is your own library, and the row
                    # shows visibility there instead — the one fact about your
                    # own style you cannot see from its name.
                    "author": None,
                    "category": s.get("category"),
                    "tags": (s.get("tags") or [])[:4],
                    "heroImage": s.get("heroImage"),
                    "url": s.get("url"),
                    "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.
        #
        # `short` is what the widget actually shows. The line under `style_ref`
        # holds one short line, and `message` is a paragraph — rendering it there
        # is how the API's sign-in instructions used to end up stretched across
        # a node. `notGenerated` lets the widget carry the same warning the
        # picker puts on the row, before the queue turns it into a 409.
        return web.json_response(
            {
                "ok": False,
                "error": err.message,
                "short": api.short_error(err),
                "notGenerated": err.status == 409 and err.code == "styleref_not_generated",
            }
        )

    # 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())
    forget_own_slugs()
    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()
    forget_own_slugs()
    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()})