ComfyUI-LM-Remote/proxy.py

"""
Reverse-proxy middleware that forwards LoRA Manager requests to the remote instance.

Registered as an aiohttp middleware on PromptServer.instance.app.  It intercepts
requests matching known LoRA Manager URL prefixes and proxies them to the remote
Docker instance.  Non-matching requests fall through to the regular ComfyUI router.

Routes that use ``send_sync`` are handled locally so that events are broadcast
to the local ComfyUI frontend (the remote instance has no connected browsers).
"""
from __future__ import annotations

import asyncio
import json
import logging
from pathlib import Path

import aiohttp
from aiohttp import web, WSMsgType

try:
    from .config import remote_config
except ImportError:
    from config import remote_config  # type: ignore[no-redef]

try:
    from .remote_client import RemoteLoraClient
except ImportError:
    from remote_client import RemoteLoraClient  # type: ignore[no-redef]

logger = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# URL prefixes that should be forwarded to the remote LoRA Manager
# ---------------------------------------------------------------------------
_PROXY_PREFIXES = (
    "/api/lm/",
    "/loras_static/",
    "/locales/",
    "/example_images_static/",
    "/extensions/ComfyUI-Lora-Manager/",
)

# Page routes served by the standalone LoRA Manager web UI
_PROXY_PAGE_ROUTES = {
    "/loras",
    "/checkpoints",
    "/embeddings",
    "/loras/recipes",
    "/statistics",
}

# WebSocket endpoints to proxy
_WS_ROUTES = {
    "/ws/fetch-progress",
    "/ws/download-progress",
    "/ws/init-progress",
}

# ---------------------------------------------------------------------------
# Local handlers for routes that need send_sync (event broadcasting)
# ---------------------------------------------------------------------------
# These routes are NOT proxied.  They are handled locally so that events
# reach the local ComfyUI frontend via PromptServer.send_sync().


def _get_prompt_server():
    """Lazily import PromptServer to avoid circular imports at module level."""
    from server import PromptServer  # type: ignore
    return PromptServer.instance


def _parse_node_id(entry):
    """Parse a node ID entry that can be int, string, or dict.

    Returns (parsed_id, graph_id_or_None).
    """
    node_identifier = entry
    graph_identifier = None
    if isinstance(entry, dict):
        node_identifier = entry.get("node_id")
        graph_identifier = entry.get("graph_id")

    try:
        parsed_id = int(node_identifier)
    except (TypeError, ValueError):
        parsed_id = node_identifier

    return parsed_id, graph_identifier


async def _handle_get_trigger_words(request: web.Request) -> web.Response:
    """Fetch trigger words from remote and broadcast via send_sync."""
    try:
        data = await request.json()
        lora_names = data.get("lora_names", [])
        node_ids = data.get("node_ids", [])

        client = RemoteLoraClient.get_instance()
        server = _get_prompt_server()

        # Collect trigger words for ALL loras into a single combined list,
        # then broadcast the same combined text to ALL node_ids.
        all_trigger_words = []
        for lora_name in lora_names:
            _, trigger_words = await client.get_lora_info(lora_name)
            all_trigger_words.extend(trigger_words)

        trigger_words_text = ",, ".join(all_trigger_words) if all_trigger_words else ""

        for entry in node_ids:
            parsed_id, graph_id = _parse_node_id(entry)
            payload = {"id": parsed_id, "message": trigger_words_text}
            if graph_id is not None:
                payload["graph_id"] = str(graph_id)
            server.send_sync("trigger_word_update", payload)

        return web.json_response({"success": True})
    except Exception as exc:
        logger.error("[LM-Remote] Error getting trigger words: %s", exc)
        return web.json_response(
            {"success": False, "error": str(exc)}, status=500
        )


async def _handle_update_lora_code(request: web.Request) -> web.Response:
    """Parse lora code update and broadcast via send_sync."""
    data = await request.json()
    node_ids = data.get("node_ids")
    lora_code = data.get("lora_code", "")
    mode = data.get("mode", "append")

    server = _get_prompt_server()

    if node_ids is None:
        # Broadcast to all nodes
        server.send_sync(
            "lora_code_update",
            {"id": -1, "lora_code": lora_code, "mode": mode},
        )
    else:
        for entry in node_ids:
            parsed_id, graph_id = _parse_node_id(entry)
            payload = {"id": parsed_id, "lora_code": lora_code, "mode": mode}
            if graph_id is not None:
                payload["graph_id"] = str(graph_id)
            server.send_sync("lora_code_update", payload)

    return web.json_response({"success": True})


async def _handle_update_node_widget(request: web.Request) -> web.Response:
    """Parse widget update and broadcast via send_sync."""
    data = await request.json()
    widget_name = data.get("widget_name")
    value = data.get("value")
    node_ids = data.get("node_ids")

    if not widget_name or value is None or not node_ids:
        return web.json_response(
            {"error": "widget_name, value, and node_ids are required"},
            status=400,
        )

    server = _get_prompt_server()

    for entry in node_ids:
        parsed_id, graph_id = _parse_node_id(entry)
        payload = {"id": parsed_id, "widget_name": widget_name, "value": value}
        if graph_id is not None:
            payload["graph_id"] = str(graph_id)
        server.send_sync("lm_widget_update", payload)

    return web.json_response({"success": True})


async def _handle_register_nodes(request: web.Request) -> web.Response:
    """No-op handler — node registration is not needed in remote mode."""
    return web.json_response({"success": True, "message": "No-op in remote mode"})


# Dispatch table for send_sync routes
_SEND_SYNC_HANDLERS = {
    "/api/lm/loras/get_trigger_words": _handle_get_trigger_words,
    "/api/lm/update-lora-code": _handle_update_lora_code,
    "/api/lm/update-node-widget": _handle_update_node_widget,
    "/api/lm/register-nodes": _handle_register_nodes,
}

# ---------------------------------------------------------------------------
# Stats enrichment
# ---------------------------------------------------------------------------

_STATS_JS_PATH = Path(__file__).resolve().parent / "static" / "lm_stats_ui.js"

_stats_db = None


def _get_stats_db():
    """Lazy-initialize the stats database."""
    global _stats_db
    if _stats_db is None:
        try:
            from .stats_db import StatsDB as _StatsDBClass
        except ImportError:
            from stats_db import StatsDB as _StatsDBClass
        _stats_db = _StatsDBClass()
        _stats_db.init()
    return _stats_db


def _enrich_list_response(response_data: dict, db=None) -> dict:
    """Merge stats from the local DB into a list API response."""
    items = response_data.get("items", [])
    if not items:
        return response_data

    if db is None:
        db = _get_stats_db()

    hashes = [item["sha256"] for item in items if item.get("sha256")]
    if not hashes:
        return response_data

    stats_map = db.get_by_hashes(hashes)

    for item in items:
        sha = item.get("sha256")
        if sha and sha in stats_map:
            s = stats_map[sha]
            item["download_count"] = s["download_count"]
            item["rating"] = s["rating"]
            item["rating_count"] = s["rating_count"]
            item["thumbs_up_count"] = s["thumbs_up_count"]

    return response_data


# ---------------------------------------------------------------------------
# Stats-based sorting
# ---------------------------------------------------------------------------

_STATS_SORT_KEYS = {"downloads", "rating", "thumbsup"}
_STATS_SORT_COLUMN = {
    "downloads": "download_count",
    "rating": "rating",
    "thumbsup": "thumbs_up_count",
}


def _is_stats_sort(sort_key: str) -> bool:
    """Return True if sort_key is a stats-based sort."""
    return sort_key in _STATS_SORT_KEYS


def _sort_items_by_stats(items, sort_key, order, db=None):
    """Sort items by a stats column. Items without stats go last."""
    if db is None:
        db = _get_stats_db()

    column = _STATS_SORT_COLUMN.get(sort_key, "download_count")
    hashes = [item["sha256"] for item in items if item.get("sha256")]
    stats_map = db.get_by_hashes(hashes)
    reverse = order == "desc"

    def get_value(item):
        sha = item.get("sha256")
        if sha and sha in stats_map:
            return stats_map[sha].get(column, 0)
        return None

    # Split items into those with stats and those without, then sort only
    # the former so that items without stats always end up last.
    with_stats = [i for i in items if get_value(i) is not None]
    without_stats = [i for i in items if get_value(i) is None]

    with_stats.sort(key=get_value, reverse=reverse)
    return with_stats + without_stats


async def _proxy_list_with_stats_sort(
    request: web.Request, sort_by: str
) -> web.Response:
    """Fetch all items from remote, sort by stats locally, paginate."""
    from urllib.parse import urlencode

    page = max(1, int(request.query.get("page", "1")))
    page_size = max(1, int(request.query.get("page_size", "20")))
    sort_key, _, order = sort_by.partition(":")
    order = order or "desc"

    # Build query to fetch all items from remote with default sort
    query = dict(request.query)
    query["page"] = "1"
    query["page_size"] = "9999"
    query["sort_by"] = "name:asc"
    remote_url = f"{remote_config.remote_url}{request.path}?{urlencode(query)}"

    # Forward relevant headers
    headers = {}
    skip = {"host", "transfer-encoding", "connection", "keep-alive", "upgrade"}
    for k, v in request.headers.items():
        if k.lower() not in skip:
            headers[k] = v

    session = await _get_proxy_session()
    try:
        async with session.get(remote_url, headers=headers) as resp:
            if resp.status != 200:
                return await _proxy_http(request)
            data = await resp.json()
    except Exception:
        return await _proxy_http(request)

    items = data.get("items", [])
    db = _get_stats_db()

    # Enrich all items with stats
    _enrich_list_response(data, db)

    # Sort by stats
    sorted_items = _sort_items_by_stats(items, sort_key, order, db)

    # Paginate locally
    total = len(sorted_items)
    start = (page - 1) * page_size
    end = start + page_size
    page_items = sorted_items[start:end]

    return web.json_response({
        "items": page_items,
        "total": total,
        "page": page,
        "page_size": page_size,
        "total_pages": (total + page_size - 1) // page_size,
    })


# ---------------------------------------------------------------------------
# Extra API endpoints (stats)
# ---------------------------------------------------------------------------

async def _handle_fetch_stats(request: web.Request) -> web.Response:
    """Bulk fetch CivitAI stats for all models."""
    if request.method != "POST":
        return web.json_response({"error": "POST required"}, status=405)
    try:
        from .stats_service import StatsFetchService
    except ImportError:
        from stats_service import StatsFetchService

    db = _get_stats_db()

    # Get the lora list from remote
    session = await _get_proxy_session()
    remote_url = f"{remote_config.remote_url}/api/lm/loras/list?page=1&page_size=9999"
    try:
        async with session.get(remote_url) as resp:
            if resp.status != 200:
                return web.json_response(
                    {"success": False, "error": "Failed to fetch lora list"},
                    status=500,
                )
            data = await resp.json()
    except Exception as exc:
        return web.json_response(
            {"success": False, "error": f"Failed to fetch lora list: {exc}"},
            status=500,
        )

    items = data.get("items", [])
    models = []
    for item in items:
        civitai = item.get("civitai") or {}
        model_id = civitai.get("modelId")
        sha256 = item.get("sha256")
        if model_id and sha256:
            models.append({"sha256": sha256, "civitai_model_id": model_id})

    if not models:
        return web.json_response({"success": True, "updated": 0,
                                   "message": "No models with CivitAI IDs found"})

    service = StatsFetchService(db, api_key=remote_config.civitai_api_key or None)
    try:
        try:
            server = _get_prompt_server()

            async def progress(current, total):
                server.send_sync("lm_stats_progress", {
                    "current": current, "total": total,
                })

            updated = await service.fetch_stats_for_models(models, progress_callback=progress)
        except Exception:
            updated = await service.fetch_stats_for_models(models)
    finally:
        await service.close()

    return web.json_response({"success": True, "updated": updated, "total": len(models)})


async def _handle_stats_status(request: web.Request) -> web.Response:
    """Check if stats DB is populated."""
    db = _get_stats_db()
    return web.json_response({"success": True, "count": db.count()})


async def _handle_stats_ui_js(request: web.Request) -> web.Response:
    """Serve the stats UI JavaScript file."""
    try:
        content = _STATS_JS_PATH.read_text(encoding="utf-8")
        return web.Response(
            text=content,
            content_type="application/javascript",
        )
    except Exception as exc:
        return web.Response(status=404, text=str(exc))


_EXTRA_API_HANDLERS = {
    "/api/lm-extra/fetch-stats": _handle_fetch_stats,
    "/api/lm-extra/stats-status": _handle_stats_status,
    "/lm-extra/stats-ui.js": _handle_stats_ui_js,
}


# Shared HTTP session for proxied requests (connection pooling)
_proxy_session: aiohttp.ClientSession | None = None


async def _get_proxy_session() -> aiohttp.ClientSession:
    """Return a shared aiohttp session for HTTP proxy requests."""
    global _proxy_session
    if _proxy_session is None or _proxy_session.closed:
        timeout = aiohttp.ClientTimeout(total=remote_config.timeout)
        _proxy_session = aiohttp.ClientSession(timeout=timeout)
    return _proxy_session


def _should_proxy(path: str) -> bool:
    """Return True if *path* should be proxied to the remote instance."""
    if any(path.startswith(p) for p in _PROXY_PREFIXES):
        return True
    if path in _PROXY_PAGE_ROUTES or path.rstrip("/") in _PROXY_PAGE_ROUTES:
        return True
    return False


def _is_ws_route(path: str) -> bool:
    return path in _WS_ROUTES


async def _proxy_ws(request: web.Request) -> web.WebSocketResponse:
    """Proxy a WebSocket connection to the remote LoRA Manager."""
    remote_url = remote_config.remote_url.replace("http://", "ws://").replace("https://", "wss://")
    remote_ws_url = f"{remote_url}{request.path}"
    if request.query_string:
        remote_ws_url += f"?{request.query_string}"

    local_ws = web.WebSocketResponse()
    await local_ws.prepare(request)

    timeout = aiohttp.ClientTimeout(total=None)
    session = aiohttp.ClientSession(timeout=timeout)
    try:
        async with session.ws_connect(remote_ws_url) as remote_ws:

            async def forward_local_to_remote():
                async for msg in local_ws:
                    if msg.type == WSMsgType.TEXT:
                        await remote_ws.send_str(msg.data)
                    elif msg.type == WSMsgType.BINARY:
                        await remote_ws.send_bytes(msg.data)
                    elif msg.type in (WSMsgType.CLOSE, WSMsgType.CLOSING, WSMsgType.CLOSED):
                        return

            async def forward_remote_to_local():
                async for msg in remote_ws:
                    if msg.type == WSMsgType.TEXT:
                        await local_ws.send_str(msg.data)
                    elif msg.type == WSMsgType.BINARY:
                        await local_ws.send_bytes(msg.data)
                    elif msg.type in (WSMsgType.CLOSE, WSMsgType.CLOSING, WSMsgType.CLOSED):
                        return

            # Run both directions concurrently.  When either side closes,
            # cancel the other to prevent hanging.
            task_l2r = asyncio.create_task(forward_local_to_remote())
            task_r2l = asyncio.create_task(forward_remote_to_local())
            try:
                done, pending = await asyncio.wait(
                    {task_l2r, task_r2l}, return_when=asyncio.FIRST_COMPLETED
                )
                for task in pending:
                    task.cancel()
                    try:
                        await task
                    except asyncio.CancelledError:
                        pass
            finally:
                # Ensure both sides are closed
                if not remote_ws.closed:
                    await remote_ws.close()
                if not local_ws.closed:
                    await local_ws.close()

    except Exception as exc:
        logger.warning("[LM-Remote] WebSocket proxy error for %s: %s", request.path, exc)
    finally:
        await session.close()

    return local_ws


async def _proxy_http(request: web.Request) -> web.Response:
    """Forward an HTTP request to the remote LoRA Manager and return its response."""
    remote_url = f"{remote_config.remote_url}{request.path}"
    if request.query_string:
        remote_url += f"?{request.query_string}"

    # Read the request body (if any)
    body = await request.read() if request.can_read_body else None

    # Filter hop-by-hop headers
    headers = {}
    skip = {"host", "transfer-encoding", "connection", "keep-alive", "upgrade"}
    for k, v in request.headers.items():
        if k.lower() not in skip:
            headers[k] = v

    session = await _get_proxy_session()
    try:
        async with session.request(
            method=request.method,
            url=remote_url,
            headers=headers,
            data=body,
        ) as resp:
            resp_body = await resp.read()
            resp_headers = {}
            for k, v in resp.headers.items():
                if k.lower() not in ("transfer-encoding", "content-encoding", "content-length"):
                    resp_headers[k] = v
            return web.Response(
                status=resp.status,
                body=resp_body,
                headers=resp_headers,
            )
    except Exception as exc:
        logger.error("[LM-Remote] Proxy error for %s %s: %s", request.method, request.path, exc)
        return web.json_response(
            {"error": f"Remote LoRA Manager unavailable: {exc}"},
            status=502,
        )


# ---------------------------------------------------------------------------
# Middleware factory
# ---------------------------------------------------------------------------

@web.middleware
async def lm_remote_proxy_middleware(request: web.Request, handler):
    """aiohttp middleware that intercepts LoRA Manager requests."""
    if not remote_config.is_configured:
        return await handler(request)

    path = request.path

    # Routes that need send_sync are handled locally so events reach
    # the local browser (the remote instance has no connected browsers).
    local_handler = _SEND_SYNC_HANDLERS.get(path)
    if local_handler is not None:
        return await local_handler(request)

    # WebSocket routes
    if _is_ws_route(path):
        return await _proxy_ws(request)

    # Extra API endpoints (stats, etc.)
    extra_handler = _EXTRA_API_HANDLERS.get(path)
    if extra_handler is not None:
        return await extra_handler(request)

    # Stats-based sorting: intercept list requests, fetch all from remote,
    # sort locally, paginate locally
    if (_should_proxy(path) and path.endswith("/list")
            and path.startswith("/api/lm/") and request.method == "GET"):
        sort_by = request.query.get("sort_by", "name")
        sort_key = sort_by.split(":")[0] if ":" in sort_by else sort_by
        if _is_stats_sort(sort_key):
            return await _proxy_list_with_stats_sort(request, sort_by)

    # Regular proxy routes
    if _should_proxy(path):
        response = await _proxy_http(request)
        # Enrich model list responses with stats from local DB
        if (path.endswith("/list") and path.startswith("/api/lm/")
                and response.status == 200):
            try:
                body = json.loads(response.body)
                enriched = _enrich_list_response(body)
                return web.json_response(enriched)
            except Exception as exc:
                logger.debug("[LM-Remote] Stats enrichment failed: %s", exc)
        # Inject stats UI script into HTML pages
        if (path in _PROXY_PAGE_ROUTES or path.rstrip("/") in _PROXY_PAGE_ROUTES):
            if response.status == 200:
                content_type = response.headers.get("Content-Type", "")
                if "text/html" in content_type:
                    try:
                        html = response.body.decode("utf-8")
                        inject = '<script src="/lm-extra/stats-ui.js"></script>'
                        html = html.replace("</body>", f"{inject}\n</body>")
                        return web.Response(
                            text=html,
                            content_type="text/html",
                            charset="utf-8",
                            status=response.status,
                        )
                    except Exception as exc:
                        logger.debug("[LM-Remote] HTML injection failed: %s", exc)
        return response

    # Not a LoRA Manager route — fall through
    return await handler(request)


async def _cleanup_proxy_session(app) -> None:
    """Shutdown hook to close the shared proxy session and stats DB."""
    global _proxy_session, _stats_db
    if _proxy_session and not _proxy_session.closed:
        await _proxy_session.close()
        _proxy_session = None
    if _stats_db:
        _stats_db.close()
        _stats_db = None


def register_proxy(app) -> None:
    """Insert the proxy middleware into the aiohttp app."""
    if not remote_config.is_configured:
        logger.warning("[LM-Remote] No remote_url configured — proxy disabled")
        return

    # Insert at position 0 so we run before the original LoRA Manager routes
    app.middlewares.insert(0, lm_remote_proxy_middleware)
    app.on_shutdown.append(_cleanup_proxy_session)
    logger.info("[LM-Remote] Proxy routes registered -> %s", remote_config.remote_url)