"""
parser.py — ComfyUI workflow asset extractor

Supports three workflow formats:
  • UI export  — top-level ``"nodes"`` array  (what you get when you save from the browser)
  • API format — top-level ``"prompt"`` dict   (what the /prompt endpoint accepts)
  • Prompt     — top-level keys are node IDs   (bare prompt dict)

Type inference runs through six layers in priority order:

  1. Exact node-type lookup       (``NODE_TYPE_TO_ASSET_TYPE``)
  2. Exact field-name lookup      (``FIELD_TO_ASSET_TYPE``)
  3. Output-slot type of the node (ComfyUI's own type system: MODEL, VAE, CLIP …)
  4. PascalCase / token decomposition of the node type string
  5. Keyword extraction from ``_meta.title`` or node ``title``
  6. Filename heuristics          (last resort)

Layers 3-5 are the new additions that make inference work for arbitrary custom
nodes without requiring a lookup-table entry.
"""

from __future__ import annotations

import json
import logging
import re
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any

logger = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# Constants
# ---------------------------------------------------------------------------

#: File extensions we treat as model assets.
ASSET_EXTENSIONS: frozenset[str] = frozenset(
    [".safetensors", ".ckpt", ".pt", ".bin", ".pth", ".gguf", ".sft"]
)

#: Field names that always indicate an asset reference regardless of node type.
KNOWN_ASSET_FIELDS: frozenset[str] = frozenset([
    "ckpt_name",
    "lora_name",
    "vae_name",
    "model_name",
    "unet_name",
    "clip_name",
    "clip_name1",
    "clip_name2",
    "control_net_name",
    "upscale_model",
    "ipadapter",
    "gligen_model",
    "style_model",
    "hypernetwork",
    "weight_dtype",   # sometimes used in custom loaders
])

# ---------------------------------------------------------------------------
# Layer 1 — exact node-type → asset type
# ---------------------------------------------------------------------------

NODE_TYPE_TO_ASSET_TYPE: dict[str, str] = {
    # ── Checkpoints ─────────────────────────────────────────────────────────
    "CheckpointLoaderSimple":       "checkpoints",
    "CheckpointLoader":             "checkpoints",
    "unCLIPCheckpointLoader":       "checkpoints",
    # ── UNETs / diffusion models ─────────────────────────────────────────────
    "UNETLoader":                   "unet",
    "DiffusionModelLoader":         "unet",
    # ── LoRA ────────────────────────────────────────────────────────────────
    "LoraLoader":                   "loras",
    "LoraLoaderModelOnly":          "loras",
    "Power Lora Loader (rgthree)":  "loras",
    "LoraLoaderBlockWeight":        "loras",
    "LoraSave":                     "loras",
    # ── VAE ─────────────────────────────────────────────────────────────────
    "VAELoader":                    "vae",
    # ── Text encoders / CLIP ─────────────────────────────────────────────────
    "CLIPLoader":                   "text_encoders",
    "DualCLIPLoader":               "text_encoders",
    "TripleCLIPLoader":             "text_encoders",
    # ── ControlNet ──────────────────────────────────────────────────────────
    "ControlNetLoader":             "controlnet",
    "ControlNetLoaderAdvanced":     "controlnet",
    "DiffControlNetLoader":         "controlnet",
    # ── Upscale ─────────────────────────────────────────────────────────────
    "UpscaleModelLoader":           "upscale_models",
    # ── IP-Adapter ──────────────────────────────────────────────────────────
    "IPAdapterModelLoader":         "ipadapter",
    "IPAdapterUnifiedLoader":       "ipadapter",
    "IPAdapter":                    "ipadapter",
    # ── GLIGEN ──────────────────────────────────────────────────────────────
    "GLIGENLoader":                 "gligen",
    # ── Style models ────────────────────────────────────────────────────────
    "StyleModelLoader":             "style_models",
    # ── Hypernetworks ───────────────────────────────────────────────────────
    "HypernetworkLoader":           "hypernetworks",
    # ── Diffusers ───────────────────────────────────────────────────────────
    "DiffusersLoader":              "diffusers",
    # ── GGUF diffusion models (ComfyUI-GGUF node pack) ───────────────────────
    "UnetLoaderGGUF":               "unet",
    "UnetLoaderGGUFAdvanced":       "unet",
    # ── GGUF text encoders (ComfyUI-GGUF node pack) ──────────────────────────
    "CLIPLoaderGGUF":               "text_encoders",
    "DualCLIPLoaderGGUF":           "text_encoders",
    "TripleCLIPLoaderGGUF":         "text_encoders",
}

# ---------------------------------------------------------------------------
# Layer 2 — exact field-name → asset type
# ---------------------------------------------------------------------------

FIELD_TO_ASSET_TYPE: dict[str, str] = {
    "ckpt_name":        "checkpoints",
    "lora_name":        "loras",
    "vae_name":         "vae",
    "model_name":       "checkpoints",
    "unet_name":        "unet",
    "clip_name":        "text_encoders",
    "clip_name1":       "text_encoders",
    "clip_name2":       "text_encoders",
    "control_net_name": "controlnet",
    "upscale_model":    "upscale_models",
    "ipadapter":        "ipadapter",
    "gligen_model":     "gligen",
    "style_model":      "style_models",
    "hypernetwork":     "hypernetworks",
}

# ---------------------------------------------------------------------------
# Layer 3 — ComfyUI output-slot type → asset type
#
# The ComfyUI type system is standardised even for custom nodes.
# When a node's first output slot is typed e.g. "VAE" we know exactly what
# kind of model it loads.
# ---------------------------------------------------------------------------

OUTPUT_TYPE_TO_ASSET_TYPE: dict[str, str] = {
    "MODEL":          "checkpoints",
    "UNET":           "unet",
    "VAE":            "vae",
    "CLIP":           "text_encoders",
    "CONTROL_NET":    "controlnet",
    "UPSCALE_MODEL":  "upscale_models",
    "IPADAPTER":      "ipadapter",
    "STYLE_MODEL":    "style_models",
    "GLIGEN":         "gligen",
    "HYPERNETWORK":   "hypernetworks",
    "LORA":           "loras",
}

# ---------------------------------------------------------------------------
# Layer 4 — keyword tokens extracted from PascalCase node type / title
# ---------------------------------------------------------------------------

#: Each entry is ``(token_set, asset_type)``.  The *first* match wins so more
#: specific entries should come first.
TOKEN_RULES: list[tuple[frozenset[str], str]] = [
    (frozenset({"lora", "lycoris"}),                        "loras"),
    (frozenset({"controlnet", "control_net", "control"}),   "controlnet"),
    (frozenset({"ipadapter", "ip_adapter", "ip-adapter"}),  "ipadapter"),
    (frozenset({"upscale", "esrgan", "realesrgan"}),        "upscale_models"),
    (frozenset({"hypernetwork", "hyperNet"}),               "hypernetworks"),
    (frozenset({"clip", "text_encoder", "textencode"}),     "text_encoders"),
    (frozenset({"vae"}),                                    "vae"),
    (frozenset({"unet", "diffusion_model", "gguf"}),        "unet"),
    (frozenset({"checkpoint", "ckpt"}),                     "checkpoints"),
    (frozenset({"embedding", "textual_inversion"}),         "embeddings"),
    (frozenset({"style_model", "stylemodel"}),              "style_models"),
    (frozenset({"gligen"}),                                 "gligen"),
    (frozenset({"diffuser"}),                               "diffusers"),
]

# ---------------------------------------------------------------------------
# Data class
# ---------------------------------------------------------------------------

@dataclass
class AssetReference:
    """
    A single model file referenced inside a ComfyUI workflow.

    Attributes:
        raw_name:       Original filename string as it appears in the workflow.
        asset_type:     Inferred category (e.g. ``"checkpoints"``, ``"loras"``).
        node_id:        ID of the node that references the file.
        node_type:      ``class_type`` / ``type`` of that node.
        inference_path: Human-readable description of which inference layer
                        resolved the type — useful for debugging / logging.
    """
    raw_name:       str
    asset_type:     str
    node_id:        str
    node_type:      str
    inference_path: str = field(default="unknown", compare=False)


# ---------------------------------------------------------------------------
# Public API
# ---------------------------------------------------------------------------

def parse_workflow(path: Path) -> list[AssetReference]:
    """
    Parse a ComfyUI workflow JSON file and return all asset references.

    Handles UI-export, API, and bare-prompt formats automatically.

    Args:
        path: Path to the workflow ``.json`` file.

    Returns:
        Deduplicated list of :class:`AssetReference` objects.

    Raises:
        json.JSONDecodeError: If the file is not valid JSON.
        OSError:              If the file cannot be read.
    """
    with open(path, "r", encoding="utf-8") as fh:
        workflow: dict[str, Any] = json.load(fh)

    assets = extract_assets(workflow)
    logger.debug(
        "Parsed %d asset(s) from %s", len(assets), path.name
    )
    return assets


def extract_assets(workflow: dict[str, Any]) -> list[AssetReference]:
    """
    Extract asset references from an already-loaded workflow dictionary.

    Args:
        workflow: Loaded workflow dict in any supported ComfyUI format.

    Returns:
        Deduplicated list of :class:`AssetReference` objects.
    """
    if "nodes" in workflow and isinstance(workflow["nodes"], list):
        return _extract_ui_format(workflow)
    else:
        return _extract_api_format(workflow)


def infer_asset_type(
    field_name: str,
    node_type: str,
    filename: str,
    output_types: list[str] | None = None,
    title: str = "",
) -> tuple[str, str]:
    """
    Infer the asset category for a model file.

    Runs six inference layers in priority order and returns on the first hit.

    Args:
        field_name:   Name of the input field containing the file reference.
        node_type:    Node ``class_type`` / ``type`` string.
        filename:     The model filename (used for heuristic fallback).
        output_types: List of ComfyUI type strings from the node's output
                      slots (e.g. ``["MODEL", "CLIP", "VAE"]``).
        title:        Human-readable node title from ``_meta.title`` or
                      the ``"title"`` key in UI-format nodes.

    Returns:
        ``(asset_type, inference_path)`` — a two-tuple so callers can log
        or display how the decision was made.
    """
    # ── Layer 1: exact node-type lookup ────────────────────────────────────
    if node_type in NODE_TYPE_TO_ASSET_TYPE:
        return NODE_TYPE_TO_ASSET_TYPE[node_type], f"node_type:{node_type}"

    # ── Layer 2: exact field-name lookup ───────────────────────────────────
    if field_name in FIELD_TO_ASSET_TYPE:
        return FIELD_TO_ASSET_TYPE[field_name], f"field_name:{field_name}"

    # ── Layer 3: output-slot type ──────────────────────────────────────────
    if output_types:
        for slot_type in output_types:
            normalised = slot_type.upper().replace(" ", "_")
            if normalised in OUTPUT_TYPE_TO_ASSET_TYPE:
                return (
                    OUTPUT_TYPE_TO_ASSET_TYPE[normalised],
                    f"output_slot:{slot_type}",
                )

    # ── Layer 4: token decomposition of node type ──────────────────────────
    node_tokens = _tokenise(node_type)
    result = _match_tokens(node_tokens)
    if result:
        return result, f"node_tokens:{node_type}"

    # ── Layer 5: keyword extraction from title ─────────────────────────────
    if title:
        title_tokens = _tokenise(title)
        result = _match_tokens(title_tokens)
        if result:
            return result, f"title:{title}"

    # ── Layer 6: filename heuristics ───────────────────────────────────────
    fn       = filename.lower()
    is_gguf  = fn.endswith(".gguf")

    heuristic_rules: list[tuple[list[str], str]] = [
        # ── Must come first — most specific ──────────────────────────────────
        (["lora", "lycoris"],                               "loras"),
        (["controlnet", "control_net", "control-net"],      "controlnet"),
        (["ipadapter", "ip_adapter", "ip-adapter"],         "ipadapter"),
        (["upscale", "esrgan", "realesrgan"],               "upscale_models"),
        (["embedding", "textual_inversion"],                "embeddings"),
        (["hypernetwork"],                                  "hypernetworks"),
        # ── Text encoders — explicit signals ─────────────────────────────────
        # mmproj = multimodal projection layer, always a text-encoder companion
        (["mmproj"],                                        "text_encoders"),
        # Embeddings connectors used by LTX2, Wan, etc.
        (["embeddings_connector", "embed_connector",
          "text_proj", "text_encoder"],                    "text_encoders"),
        # Well-known text encoder architectures
        (["clip_l", "clip_g", "clip_h",
          "t5xxl", "t5-xxl", "t5_xxl",
          "t5xl", "t5-xl", "t5_xl"],                       "text_encoders"),
        # LLM-family models used as text encoders in modern pipelines
        # (Gemma, LLaMA, Mistral, Phi used as prompt encoders)
        (["gemma", "llama", "mistral", "phi"],              "text_encoders"),
        # Qwen VL models used as *text encoders* carry VL/VLM/instruct tags
        # Qwen *diffusion* models (image editors) don't — handled below
        (["qwen2.5-vl", "qwen2-vl", "qwen_vl",
          "qwen-vl", "qwenvl", "instruct"],                "text_encoders"),
        # ── VAE ───────────────────────────────────────────────────────────────
        (["vae"],                                           "vae"),
        # ── UNet / diffusion model ────────────────────────────────────────────
        (["unet", "diffusion_model"],                       "unet"),
    ]

    for keywords, asset_type in heuristic_rules:
        if any(kw in fn for kw in keywords):
            return asset_type, "filename_heuristic"

    # ── GGUF default: models/unet (not checkpoints) ──────────────────────
    # ComfyUI-GGUF always loads diffusion GGUFs from models/unet.
    # Text-encoder GGUFs are caught above; anything else is a diffusion model.
    if is_gguf:
        return "unet", "gguf_default"

    return "other", "fallback"


# ---------------------------------------------------------------------------
# Format-specific extractors
# ---------------------------------------------------------------------------

def _extract_ui_format(workflow: dict[str, Any]) -> list[AssetReference]:
    """
    Extract assets from the UI export format.

    In this format nodes are an array and widget values (including file names)
    are stored in ``widgets_values``.  The node's output slot types are stored
    in the ``outputs`` array and are the best inference signal for custom nodes.

    Args:
        workflow: Loaded workflow dict containing a ``"nodes"`` list.

    Returns:
        Deduplicated list of :class:`AssetReference` objects.
    """
    assets:     list[AssetReference] = []
    seen_names: set[str]             = set()

    for node_data in workflow["nodes"]:
        if not isinstance(node_data, dict):
            continue

        node_id   = str(node_data.get("id", "unknown"))
        node_type = node_data.get("type", "Unknown")
        title     = node_data.get("title", "")

        # Collect output slot types — Layer 3 inference signal
        output_types = _collect_output_types(node_data.get("outputs", []))

        # widgets_values holds the actual field values in UI format
        widgets_values = node_data.get("widgets_values", [])
        if isinstance(widgets_values, list):
            for value in widgets_values:
                if isinstance(value, str) and _is_asset_filename(value):
                    _add_asset(
                        assets, seen_names,
                        filename=value,
                        field_name="",
                        node_id=node_id,
                        node_type=node_type,
                        output_types=output_types,
                        title=title,
                    )
                elif isinstance(value, dict):
                    # Nested structure e.g. Power Lora Loader
                    _extract_from_dict(
                        value, node_id, node_type, output_types, title,
                        assets, seen_names,
                    )

        # Also scan the inputs array — some UI-format nodes store field names
        # here which helps Layer 2 inference
        inputs = node_data.get("inputs", [])
        if isinstance(inputs, list):
            for inp in inputs:
                if not isinstance(inp, dict):
                    continue
                field_name = inp.get("name", "")
                # widget-linked inputs sometimes carry the value directly
                value = inp.get("widget", {})
                if isinstance(value, dict):
                    v = value.get("value", "")
                    if isinstance(v, str) and _is_asset_filename(v):
                        _add_asset(
                            assets, seen_names,
                            filename=v,
                            field_name=field_name,
                            node_id=node_id,
                            node_type=node_type,
                            output_types=output_types,
                            title=title,
                        )

    return assets


def _extract_api_format(workflow: dict[str, Any]) -> list[AssetReference]:
    """
    Extract assets from the API / prompt format.

    In this format nodes are a dict keyed by node ID.  Input values are stored
    directly inside ``"inputs"``.  We also read ``_meta.title`` for Layer 5
    inference.

    Args:
        workflow: Loaded workflow dict.  May have a top-level ``"prompt"`` key
                  wrapping the node dict (API format), or the node dict may be
                  at the top level (bare prompt format).

    Returns:
        Deduplicated list of :class:`AssetReference` objects.
    """
    assets:     list[AssetReference] = []
    seen_names: set[str]             = set()

    nodes: dict[str, Any] = workflow.get("prompt", workflow)

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

        node_type = node_data.get("class_type", "Unknown")
        title     = node_data.get("_meta", {}).get("title", "")
        inputs    = node_data.get("inputs", {})

        if not isinstance(inputs, dict):
            continue

        # API format doesn't have explicit output slot declarations in the
        # node dict, so Layer 3 is not available here.  Layers 1, 2, 4, 5, 6
        # still apply.
        for field_name, field_value in inputs.items():
            if not isinstance(field_value, str):
                continue
            if field_name in KNOWN_ASSET_FIELDS or _is_asset_filename(field_value):
                _add_asset(
                    assets, seen_names,
                    filename=field_value,
                    field_name=field_name,
                    node_id=str(node_id),
                    node_type=node_type,
                    output_types=None,
                    title=title,
                )

    return assets


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def _add_asset(
    assets:       list[AssetReference],
    seen_names:   set[str],
    filename:     str,
    field_name:   str,
    node_id:      str,
    node_type:    str,
    output_types: list[str] | None,
    title:        str,
) -> None:
    """
    Add an asset to *assets* if it has not been seen before.

    Deduplication is by filename only — the first occurrence wins.

    Args:
        assets:       Accumulator list.
        seen_names:   Set of filenames already added.
        filename:     The model filename.
        field_name:   Input field name (may be empty for UI format).
        node_id:      Node ID string.
        node_type:    Node type / class_type string.
        output_types: Output slot type strings (UI format only).
        title:        Human-readable node title.
    """
    if filename in seen_names:
        return
    if not _is_asset_filename(filename):
        return

    asset_type, inference_path = infer_asset_type(
        field_name=field_name,
        node_type=node_type,
        filename=filename,
        output_types=output_types,
        title=title,
    )

    assets.append(AssetReference(
        raw_name=filename,
        asset_type=asset_type,
        node_id=node_id,
        node_type=node_type,
        inference_path=inference_path,
    ))
    seen_names.add(filename)
    logger.debug(
        "Found asset %-50s  type=%-16s  via=%s",
        filename, asset_type, inference_path,
    )


def _extract_from_dict(
    data:         dict[str, Any],
    node_id:      str,
    node_type:    str,
    output_types: list[str] | None,
    title:        str,
    assets:       list[AssetReference],
    seen_names:   set[str],
) -> None:
    """
    Recursively scan a nested dict/list for asset filenames.

    Used for custom nodes that embed file references inside structured widget
    values (e.g. Power Lora Loader).

    Args:
        data:         The dict to scan.
        node_id:      Parent node ID.
        node_type:    Parent node type.
        output_types: Parent node output slot types.
        title:        Parent node title.
        assets:       Accumulator list.
        seen_names:   Deduplication set.
    """
    for key, value in data.items():
        if isinstance(value, str) and _is_asset_filename(value):
            _add_asset(
                assets, seen_names,
                filename=value,
                field_name=key,
                node_id=node_id,
                node_type=node_type,
                output_types=output_types,
                title=title,
            )
        elif isinstance(value, dict):
            _extract_from_dict(
                value, node_id, node_type, output_types, title, assets, seen_names
            )
        elif isinstance(value, list):
            for item in value:
                if isinstance(item, dict):
                    _extract_from_dict(
                        item, node_id, node_type, output_types, title, assets, seen_names
                    )


def _collect_output_types(outputs: Any) -> list[str]:
    """
    Extract the ComfyUI type strings from a node's ``outputs`` array.

    Args:
        outputs: Value of the ``"outputs"`` key in a UI-format node dict.

    Returns:
        List of type strings, e.g. ``["MODEL", "CLIP", "VAE"]``.
        Empty list if ``outputs`` is missing or malformed.
    """
    if not isinstance(outputs, list):
        return []
    result: list[str] = []
    for slot in outputs:
        if isinstance(slot, dict):
            t = slot.get("type", "")
            if isinstance(t, str) and t:
                result.append(t)
    return result


def _is_asset_filename(value: str) -> bool:
    """
    Return True if *value* looks like a model filename.

    Checks the file extension against :data:`ASSET_EXTENSIONS` and
    additionally filters out obvious non-file strings (e.g. node variable
    names used in SetNode/GetNode).

    Args:
        value: String to test.

    Returns:
        True if the string ends with a known asset extension.
    """
    if not value:
        return False
    # Quick extension check
    suffix = Path(value).suffix.lower()
    return suffix in ASSET_EXTENSIONS


def _tokenise(text: str) -> set[str]:
    """
    Split a CamelCase / snake_case / space-separated / emoji-prefixed string
    into lowercase tokens for keyword matching.

    Examples::

        "DualCLIPLoader"            → {"dual", "clip", "loader"}
        "LTXVAudioVAELoader"        → {"ltxv", "audio", "vae", "loader"}
        "🅛🅣🅧 LTXV Audio VAE Loader" → {"ltxv", "audio", "vae", "loader"}
        "control_net_name"          → {"control", "net", "name"}

    Args:
        text: Raw string to tokenise.

    Returns:
        Set of lowercase token strings.
    """
    # Strip emoji and other non-ASCII characters
    ascii_text = text.encode("ascii", errors="ignore").decode()

    # Insert a space before every uppercase letter that follows a lowercase
    # letter or digit (handles PascalCase and abbreviations like "VAE")
    spaced = re.sub(r"([a-z0-9])([A-Z])", r"\1 \2", ascii_text)
    # Also split sequences of uppercase letters followed by a lowercase letter
    # e.g. "CLIPLoader" → "CLIP Loader"
    spaced = re.sub(r"([A-Z]+)([A-Z][a-z])", r"\1 \2", spaced)

    # Split on any non-alphanumeric character
    raw_tokens = re.split(r"[^a-zA-Z0-9]+", spaced)
    return {t.lower() for t in raw_tokens if len(t) >= 2}


def _match_tokens(tokens: set[str]) -> str | None:
    """
    Test *tokens* against :data:`TOKEN_RULES` and return the first match.

    Args:
        tokens: Set of lowercase token strings (from :func:`_tokenise`).

    Returns:
        Asset type string on match, ``None`` if no rule matched.
    """
    for keyword_set, asset_type in TOKEN_RULES:
        if tokens & keyword_set:   # non-empty intersection → match
            return asset_type
    return None