path_generator_node.py

"""
Path Generator Node for ComfyUI (V3 API)
Generates configurable output paths from dynamic user-managed segments.
© 2026 Created with ❤️ by Alex Munteanu | alexmunteanu.com
"""

import json
import re
from typing_extensions import override

from comfy_api.latest import ComfyExtension, io


DEFAULT_CONFIG = json.dumps([
    {"name": "project", "type": "folder", "value": ""},
    {"name": "shot", "type": "folder", "value": ""},
    {"name": "prefix", "type": "prefix", "value": ""},
    {"name": "suffix", "type": "suffix", "value": ""},
])


def _normalize_path(path):
    """Normalize a path: collapse //, resolve . and .., strip trailing /."""
    if not path:
        return path
    # Normalize slashes
    path = path.replace('\\', '/')
    # Preserve drive letter prefix (e.g. D:/)
    drive = ''
    if len(path) >= 2 and path[1] == ':':
        drive = path[:2]
        path = path[2:]
    # Collapse consecutive slashes
    path = re.sub(r'/+', '/', path)
    # Resolve . and .. segments
    parts = path.split('/')
    resolved = []
    for part in parts:
        if part == '.' or part == '':
            continue
        if part == '..':
            if resolved:
                resolved.pop()
            continue
        resolved.append(part)
    result = '/'.join(resolved)
    if drive:
        result = drive + '/' + result if result else drive + '/'
    return result


def build_path_parts(output_folder, config_json, sequence, frame_padding, delimiter, overrides=None):
    """Build folder path and filename from config segments.

    Args:
        output_folder: Base output folder.
        config_json: JSON string of segment list.
        sequence: If True, filename becomes a subfolder: path/filename/filename.
        frame_padding: If True, append .%04d to filename.
        delimiter: Separator between filename parts.
        overrides: Dict of segment name -> value from connected inputs.

    Returns:
        Tuple of (folder_path, filename).
    """
    try:
        segments = json.loads(config_json)
    except (json.JSONDecodeError, TypeError):
        segments = []

    # Apply overrides by segment name
    if overrides:
        for i, seg in enumerate(segments):
            name = seg.get("name", "")
            if name in overrides:
                segments[i] = dict(seg)
                segments[i]["value"] = overrides[name] if overrides[name] is not None else ""

    # Sanitize inputs
    if output_folder:
        output_folder = output_folder.replace('\\', '/').strip().rstrip('/')
    delimiter = delimiter.replace('/', '').replace('\\', '')
    for s in segments:
        val = s.get("value", "")
        if val:
            s["value"] = val.strip().replace('/', '').replace('\\', '')

    folder_segments = [s for s in segments if s.get("type") == "folder"]
    prefix_segments = [s for s in segments if s.get("type") == "prefix"]
    suffix_segments = [s for s in segments if s.get("type") == "suffix"]

    # Group folder segments into chains (linked segments attach to previous)
    chains = []
    current_chain = []
    for i, seg in enumerate(folder_segments):
        if i > 0 and seg.get("linked") and current_chain:
            current_chain.append(seg)
        else:
            if current_chain:
                chains.append(current_chain)
            current_chain = [seg]
    if current_chain:
        chains.append(current_chain)

    # Build folder hierarchy from chains
    folder_path_parts = []
    for chain in chains:
        if chain[0].get("value"):
            folder_path_parts.append(chain[0]["value"])
        if len(chain) > 1:
            chain_values = [s["value"] for s in chain if s.get("value")]
            if len(chain_values) > 1:
                folder_path_parts.append(delimiter.join(chain_values))

    # Filename: prefixes + ALL folder values + suffixes
    folder_values = [s["value"] for s in folder_segments if s.get("value")]
    prefix_values = [s["value"] for s in prefix_segments if s.get("value")]
    suffix_values = [s["value"] for s in suffix_segments if s.get("value")]
    filename_parts = prefix_values + folder_values + suffix_values
    filename = delimiter.join(filename_parts) if filename_parts else ""

    # Build folder path
    dir_parts = []
    if output_folder:
        dir_parts.append(output_folder)
    dir_parts.extend(folder_path_parts)

    if sequence and filename:
        dir_parts.append(filename)

    # Build final filename
    final_filename = ""
    if filename:
        final_filename = filename
        if frame_padding:
            final_filename += ".%04d"

    folder_path = _normalize_path("/".join(dir_parts))
    return folder_path, final_filename


def build_path(output_folder, config_json, sequence, frame_padding, delimiter, overrides=None):
    """Build the full output path. Convenience wrapper around build_path_parts."""
    folder_path, filename = build_path_parts(
        output_folder, config_json, sequence, frame_padding, delimiter, overrides
    )
    parts = [p for p in [folder_path, filename] if p]
    return _normalize_path("/".join(parts))


class PathGeneratorNode(io.ComfyNode):
    @classmethod
    def define_schema(cls):
        return io.Schema(
            node_id="PathGeneratorNode",
            display_name="Path Generator",
            category="Utils/Path",
            description="Generates configurable output paths from dynamic segments with live preview.",
            accept_all_inputs=True,
            inputs=[
                io.String.Input(
                    "output_folder",
                    default="",
                    tooltip="Base output folder (relative to ComfyUI output directory). Leave empty for default.",
                ),
                io.String.Input(
                    "config",
                    default=DEFAULT_CONFIG,
                    tooltip="Segment configuration (managed by the UI).",
                ),
                io.String.Input(
                    "delimiter",
                    default="_",
                    tooltip="Separator between filename parts.",
                ),
                io.Boolean.Input(
                    "sequence",
                    default=False,
                    tooltip="Create a subfolder with the filename: path/filename/filename",
                ),
                io.Boolean.Input(
                    "frame_padding",
                    default=False,
                    tooltip="Append .%04d frame padding to the filename.",
                ),
            ],
            outputs=[
                io.String.Output(
                    "path",
                    tooltip="Full generated path (folder + filename).",
                ),
                io.String.Output(
                    "folder_path",
                    tooltip="Folder portion of the path.",
                ),
                io.String.Output(
                    "filename",
                    tooltip="Filename portion of the path.",
                ),
            ],
            is_output_node=False,
        )

    @classmethod
    def validate_inputs(cls, output_folder, config, delimiter, sequence, frame_padding, **kwargs) -> bool:
        return True

    @classmethod
    def execute(cls, output_folder, config, delimiter, sequence, frame_padding, **kwargs) -> io.NodeOutput:
        overrides = kwargs
        folder_path, filename = build_path_parts(
            output_folder, config, sequence, frame_padding, delimiter, overrides
        )
        parts = [p for p in [folder_path, filename] if p]
        full_path = _normalize_path("/".join(parts))
        return io.NodeOutput(full_path, folder_path, filename)


class PathGeneratorExtension(ComfyExtension):
    @override
    async def get_node_list(self) -> list[type[io.ComfyNode]]:
        return [PathGeneratorNode]

    @override
    async def on_load(self) -> None:
        pass


async def comfy_entrypoint() -> PathGeneratorExtension:
    return PathGeneratorExtension()