#!/bin/sh
#
# install_nodes.sh
#
# Clone custom node repositories listed in a node list.
#
# Usage:
#
#   ./install_nodes.sh
#   ./install_nodes.sh custom_nodes.txt
#   ./install_nodes.sh another_nodes.txt
#

set -eu

########################################################################
# Configuration
########################################################################

COMFYUI_DIR="/workspace/runpod-slim/ComfyUI"

DEFAULT_NODE_LIST="custom_nodes.txt"

CONFIG_FILE="config.env"

########################################################################
# Locate script directory
########################################################################

SCRIPT_DIR=$(
    CDPATH= cd -- "$(dirname -- "$0")" &&
    pwd
)

########################################################################
# Load configuration
########################################################################

CONFIG_PATH="$SCRIPT_DIR/$CONFIG_FILE"

if [ -f "$CONFIG_PATH" ]
then
    # shellcheck source=/dev/null
    . "$CONFIG_PATH"
fi

########################################################################
# Select node list
########################################################################

case $# in
    0)
        NODE_LIST="$SCRIPT_DIR/$DEFAULT_NODE_LIST"
        ;;
    1)
        case "$1" in
            /*)
                NODE_LIST="$1"
                ;;
            *)
                NODE_LIST="$SCRIPT_DIR/$1"
                ;;
        esac
        ;;
    *)
        echo "Usage:" >&2
        echo "    $0 [node-list]" >&2
        exit 1
        ;;
esac

########################################################################
# Checks
########################################################################

if [ ! -f "$NODE_LIST" ]
then
    echo "Error: node list not found:" >&2
    echo "    $NODE_LIST" >&2
    exit 1
fi

if [ ! -d "$COMFYUI_DIR" ]
then
    echo "Error: ComfyUI directory not found:" >&2
    echo "    $COMFYUI_DIR" >&2
    exit 1
fi

CUSTOM_NODES_DIR="$COMFYUI_DIR/custom_nodes"

if [ ! -d "$CUSTOM_NODES_DIR" ]
then
    echo "Error: custom_nodes directory not found:" >&2
    echo "    $CUSTOM_NODES_DIR" >&2
    exit 1
fi

if ! command -v git >/dev/null 2>&1
then
    echo "Error: git command not found." >&2
    exit 1
fi

########################################################################
# Clone helper
########################################################################

clone_node()
{
    REPOSITORY_URL=$1
    DIRECTORY_NAME=$2

    if [ -z "$REPOSITORY_URL" ] || [ -z "$DIRECTORY_NAME" ]
    then
        echo "Error: invalid node list entry." >&2
        echo "Each line must be: repository_url|directory_name" >&2
        exit 1
    fi

    TARGET_DIRECTORY="$CUSTOM_NODES_DIR/$DIRECTORY_NAME"

    if [ -d "$TARGET_DIRECTORY" ]
    then
        echo "Skipping $DIRECTORY_NAME: already exists"
        return
    fi

    if [ -e "$TARGET_DIRECTORY" ] || [ -L "$TARGET_DIRECTORY" ]
    then
        echo "Error: target path exists but is not a directory:" >&2
        echo "    $TARGET_DIRECTORY" >&2
        exit 1
    fi

    echo "Cloning $DIRECTORY_NAME..."

    if git clone "$REPOSITORY_URL" "$TARGET_DIRECTORY"
    then
        echo "Successfully cloned $DIRECTORY_NAME"
    else
        rm -rf "$TARGET_DIRECTORY"
        echo "Error: failed to clone $DIRECTORY_NAME from $REPOSITORY_URL" >&2
        exit 1
    fi
}

########################################################################
# Main
########################################################################

CARRIAGE_RETURN=$(printf '\r')

while IFS='|' read -r REPOSITORY_URL DIRECTORY_NAME
do
    REPOSITORY_URL=${REPOSITORY_URL%"$CARRIAGE_RETURN"}
    DIRECTORY_NAME=${DIRECTORY_NAME%"$CARRIAGE_RETURN"}

    #
    # Skip blank lines
    #
    [ -z "$REPOSITORY_URL" ] && [ -z "$DIRECTORY_NAME" ] && continue

    #
    # Skip comments
    #
    case "$REPOSITORY_URL" in
        \#*)
            continue
            ;;
    esac

    clone_node "$REPOSITORY_URL" "$DIRECTORY_NAME"

done < "$NODE_LIST"

echo
echo "=================================================="
echo "All custom nodes installed."
echo "=================================================="