#!/bin/bash
set -e

echo "Creating directories for models..."
MODEL_DIRECTORIES=(
    "checkpoints"
    "clip"
    "clip_vision"
    "configs"
    "controlnet"
    "diffusers"
    "diffusion_models"
    "embeddings"
    "gligen"
    "hypernetworks"
    "loras"
    "photomaker"
    "style_models"
    "text_encoders"
    "unet"
    "upscale_models"
    "vae"
    "vae_approx"
)
for MODEL_DIRECTORY in "${MODEL_DIRECTORIES[@]}"; do
    mkdir -p "/opt/comfyui/models/$MODEL_DIRECTORY"
done

echo "Preparing ComfyUI Manager..."
export PYTHONPATH="/opt/comfyui-manager${PYTHONPATH:+:$PYTHONPATH}"
rm --force /opt/comfyui/custom_nodes/ComfyUI-Manager
ln -s /opt/comfyui-manager /opt/comfyui/custom_nodes/ComfyUI-Manager
mkdir -p /opt/comfyui/user/__manager
python - <<'PY'
from configparser import ConfigParser
from pathlib import Path

config_path = Path("/opt/comfyui/user/__manager/config.ini")
config = ConfigParser(strict=False)
config.read(config_path)
if "default" not in config:
    config["default"] = {}

defaults = {
    "git_exe": "",
    "use_uv": "True",
    "channel_url": "https://raw.githubusercontent.com/ltdrdata/ComfyUI-Manager/main",
    "share_option": "all",
    "bypass_ssl": "False",
    "file_logging": "True",
    "update_policy": "stable-comfyui",
    "windows_selector_event_loop_policy": "False",
    "model_download_by_agent": "False",
    "downgrade_blacklist": "",
    "security_level": "normal",
    "always_lazy_install": "False",
    "verbose": "False",
}
for key, value in defaults.items():
    config["default"].setdefault(key, value)

config["default"]["network_mode"] = "offline"
config["default"]["db_mode"] = "local"

with config_path.open("w", encoding="utf-8") as handle:
    config.write(handle)
PY

echo "Installing requirements for custom nodes..."
for CUSTOM_NODE_DIRECTORY in /opt/comfyui/custom_nodes/*; do
    if [ "$CUSTOM_NODE_DIRECTORY" != "/opt/comfyui/custom_nodes/ComfyUI-Manager" ]; then
        if [ -f "$CUSTOM_NODE_DIRECTORY/requirements.txt" ]; then
            CUSTOM_NODE_NAME=${CUSTOM_NODE_DIRECTORY##*/}
            CUSTOM_NODE_NAME=${CUSTOM_NODE_NAME//[-_]/ }
            echo "Installing requirements for $CUSTOM_NODE_NAME..."
            pip install --requirement "$CUSTOM_NODE_DIRECTORY/requirements.txt"
        fi
    fi
done

if [ -z "$USER_ID" ] || [ -z "$GROUP_ID" ]; then
    echo "Running container as $USER..."
    exec /opt/conda/bin/python main.py \
        --port 8188 \
        --listen 0.0.0.0 \
        --disable-auto-launch \
        "$@"
else
    echo "Creating non-root user..."
    getent group "$GROUP_ID" > /dev/null 2>&1 || groupadd --gid "$GROUP_ID" comfyui-user
    id -u "$USER_ID" > /dev/null 2>&1 || useradd --uid "$USER_ID" --gid "$GROUP_ID" --create-home comfyui-user
    chown --recursive "$USER_ID:$GROUP_ID" /opt/comfyui
    chown --recursive "$USER_ID:$GROUP_ID" /opt/comfyui-manager
    export PATH=$PATH:/home/comfyui-user/.local/bin

    echo "Running container as comfyui-user ($USER_ID:$GROUP_ID)..."
    sudo --set-home --preserve-env=PATH,PYTHONPATH --user "#$USER_ID" \
        /opt/conda/bin/python main.py \
            --port 8188 \
            --listen 0.0.0.0 \
            --disable-auto-launch \
            "$@"
fi