ComfyUI-OpenClaw/services/preflight.py

"""
Preflight Diagnostics Service (R42).

Provides logic to validate a workflow against the local ComfyUI environment,
checking for missing node classes and models.
"""

import logging
import time
from typing import Any, Dict, List, Set, Tuple

logger = logging.getLogger("ComfyUI-OpenClaw.services.preflight")

# IMPORTANT (ComfyUI runtime wiring):
# This module is imported both:
# - inside a real ComfyUI runtime (where `nodes` and `folder_paths` exist), and
# - in unit tests / tooling contexts (where they may not).
# Keep these imports optional and keep references guarded.
try:  # pragma: no cover (best-effort ComfyUI imports)
    import nodes  # type: ignore
except Exception:  # pragma: no cover
    nodes = None  # type: ignore

try:  # pragma: no cover (best-effort ComfyUI imports)
    import folder_paths  # type: ignore
except Exception:  # pragma: no cover
    folder_paths = None  # type: ignore

_CACHE = {}
_CACHE_TTL = 60  # seconds

# Heuristic mapping: input_key -> folder_paths type
_INPUT_KEY_MAP = {
    "ckpt_name": "checkpoints",
    "checkpoint": "checkpoints",
    "lora_name": "loras",
    "vae_name": "vae",
    "control_net_name": "controlnet",
    "upscale_model_name": "upscale_models",
    "style_model_name": "style_models",
    "clip_name": "clip",
    "unet_name": "unet",
    # Add more as discovered
}


def _get_node_class_mappings() -> Dict[str, Any]:
    """Safely retrieve the global NODE_CLASS_MAPPINGS."""
    if nodes and hasattr(nodes, "NODE_CLASS_MAPPINGS"):
        return nodes.NODE_CLASS_MAPPINGS
    return {}


def _get_model_inventory() -> Dict[str, List[str]]:
    """
    Retrieve snapshot of available models using folder_paths.
    Returns a dict mapping folder name (e.g., 'checkpoints') to list of filenames.
    Cached for 60s to prevent IO spam.
    """
    global _CACHE
    now = time.time()

    cached = _CACHE.get("inventory")
    if cached:
        timestamp, data = cached
        if now - timestamp < _CACHE_TTL:
            return data

    inventory = {}
    if not folder_paths:
        return inventory

    # Common model types to check
    # We use the keys from folder_paths.folder_names_and_paths if available,
    # or a hardcoded list of common ones.
    model_types = [
        "checkpoints",
        "loras",
        "vae",
        "embeddings",
        "controlnet",
        "upscale_models",
        "clip",
        "unet",
        "clip_vision",
        "style_models",
        "diffusers",
        "vae_approx",
        "photomaker",
    ]

    # Add any dynamic ones
    if hasattr(folder_paths, "folder_names_and_paths"):
        for k in folder_paths.folder_names_and_paths.keys():
            if k not in model_types:
                model_types.append(k)

    for mtype in model_types:
        try:
            files = folder_paths.get_filename_list(mtype)
            if files:
                inventory[mtype] = list(files)
        except Exception:
            # Some folders might not exist or raise error
            continue

    _CACHE["inventory"] = (now, inventory)
    return inventory


def run_preflight_check(workflow: Dict[str, Any]) -> Dict[str, Any]:
    """
    Analyze a workflow (API format) and return a diagnostic report.

    Args:
        workflow: The ComfyUI workflow JSON (node ID -> node data).

    Returns:
        Dict containing validation results (missing_nodes, missing_models, etc.)
    """
    report = {
        "ok": True,
        "summary": {"missing_nodes": 0, "missing_models": 0, "invalid_inputs": 0},
        "missing_nodes": [],
        "missing_models": [],
        "invalid_inputs": [],
        "notes": [],
    }

    if not isinstance(workflow, dict):
        report["ok"] = False
        report["notes"].append("Workflow must be a JSON object (API format).")
        return report

    # 1. Check Nodes
    available_nodes = _get_node_class_mappings()
    missing_node_counts: Dict[str, int] = {}

    # 2. Check Models (Heuristic)
    inventory = _get_model_inventory()
    missing_models_counts: Dict[str, Dict[str, Any]] = {}

    for node_id, node_data in workflow.items():
        if not isinstance(node_data, dict):
            continue

        # Check Node Class
        class_type = node_data.get("class_type")
        if not class_type:
            continue

        if available_nodes and class_type not in available_nodes:
            missing_node_counts[class_type] = missing_node_counts.get(class_type, 0) + 1

        # Check Inputs for Models
        inputs = node_data.get("inputs")
        if isinstance(inputs, dict):
            _check_inputs_for_models(inputs, inventory, missing_models_counts)

    # Format Results
    for cls, count in missing_node_counts.items():
        report["missing_nodes"].append({"class_type": cls, "count": count})

    for key, info in missing_models_counts.items():
        report["missing_models"].append(
            {"type": info["type"], "name": info["name"], "count": info["count"]}
        )

    # Summarize
    report["summary"]["missing_nodes"] = len(report["missing_nodes"])
    report["summary"]["missing_models"] = len(report["missing_models"])

    if (
        report["summary"]["missing_nodes"] > 0
        or report["summary"]["missing_models"] > 0
    ):
        report["ok"] = False

    if not nodes:
        report["notes"].append("Node inventory unavailable (backend import failed).")
    if not folder_paths:
        report["notes"].append("Model inventory unavailable (backend import failed).")

    # F49: Inject Guidance Banners
    # We serialize them so they are ready for JSON response
    banners = generate_preflight_banners(report)
    report["banners"] = [b.to_dict() for b in banners]

    return report


def _check_inputs_for_models(
    inputs: Dict[str, Any],
    inventory: Dict[str, List[str]],
    missing_counts: Dict[str, Dict[str, Any]],
):
    """
    Heuristic to detect missing models in node inputs.
    We look for keys that hint at model types (e.g. 'ckpt_name', 'lora_name').
    """
    # Mapping heuristic: input_key -> folder_paths type
    key_map = _INPUT_KEY_MAP

    for key, value in inputs.items():
        if not isinstance(value, str):
            continue

        target_type = key_map.get(key)
        if target_type:
            # Check if exists
            available = inventory.get(target_type, [])
            if value not in available:
                # Also try normalizing separators just in case (e.g. windows vs linux paths)
                # But typically ComfyUI expects exact match or relative match.
                # Use simple exact match for now.

                unique_key = f"{target_type}:{value}"

                if unique_key not in missing_counts:
                    missing_counts[unique_key] = {
                        "type": target_type,
                        "name": value,
                        "count": 0,
                    }
                missing_counts[unique_key]["count"] += 1


# F49: Banner Generation Support
def generate_preflight_banners(report: Dict[str, Any]) -> List["OperatorBanner"]:
    """
    Generate actionable guidance banners from a preflight report.
    Returns list of OperatorBanner objects.
    """
    # CRITICAL: keep package-relative import first.
    # Direct `services.*` imports can fail when loaded as a ComfyUI package module.
    if __package__ and "." in __package__:
        from .operator_guidance import BannerSeverity, OperatorAction, OperatorBanner
    else:  # pragma: no cover (standalone/test import mode)
        from services.operator_guidance import (  # type: ignore
            BannerSeverity,
            OperatorAction,
            OperatorBanner,
        )

    banners = []

    if report.get("ok"):
        return banners

    # 1. Missing Nodes
    missing_nodes = report.get("missing_nodes", [])
    # Sort for determinism
    missing_nodes.sort(key=lambda x: x["class_type"])

    if missing_nodes:
        node_names = [n["class_type"] for n in missing_nodes]
        count = len(node_names)
        preview = ", ".join(node_names[:3])
        if count > 3:
            preview += f" and {count - 3} more"

        banners.append(
            OperatorBanner(
                id="missing_nodes",
                severity=BannerSeverity.ERROR,
                message=f"Workflow requires missing custom nodes: {preview}",
                source="Preflight",
                action=OperatorAction(
                    label="Manager",
                    type="tab",
                    payload="manager",  # Future: deep link to manager
                ).to_dict(),
            )
        )

    # 2. Missing Models
    missing_models = report.get("missing_models", [])
    # Sort for determinism
    missing_models.sort(key=lambda x: (x["type"], x["name"]))

    if missing_models:
        model_names = [f"{m['name']} ({m['type']})" for m in missing_models]
        count = len(model_names)
        preview = ", ".join(model_names[:3])
        if count > 3:
            preview += f" and {count - 3} more"

        banners.append(
            OperatorBanner(
                id="missing_models",
                severity=BannerSeverity.WARNING,
                message=f"Workflow refers to missing models: {preview}",
                source="Preflight",
                # No specific action for models yet, maybe just docs or upload.
            )
        )

    # 3. Notes/Errors
    notes = report.get("notes", [])
    for i, note in enumerate(notes):
        banners.append(
            OperatorBanner(
                id=f"preflight_note_{i}",
                severity=BannerSeverity.WARNING,
                message=note,
                source="Preflight",
            )
        )

    return banners