#!/bin/sh
SERVICE="comfyui.service"
DIR="${XDG_CONFIG_HOME:-$HOME/.config}/systemd/user"
FILE="$DIR/$SERVICE"

detect_comfyui_dir() {
	for d in \
		"$HOME/AI/ComfyUI" \
		"$HOME/ComfyUI" \
		"/media/$USER/"*/AI/ComfyUI \
		"/media/$USER/"*/ComfyUI \
		"/mnt/"*/AI/ComfyUI \
		"/mnt/"*/ComfyUI
	do
		if [ -f "$d/main.py" ] && [ -d "$d/comfy" ]; then
			echo "$d"
			return 0
		fi
	done
	return 1
}

install() {
	MNT="${1:-$(detect_comfyui_dir)}"
	[ -n "$MNT" ] || { echo "ComfyUI directory not found. Specify it as an argument: $0 install /path/to/ComfyUI" >&2; exit 1; }

	mkdir -p "$DIR"
	cat > "$FILE" <<- EOF
	[Unit]
	Description=ComfyUI
	After=graphical-session.target
	Requires=graphical-session.target

	[Service]
	ExecStart=/bin/sh -c '\
	  MNT="$MNT" && \
	  cd "\$MNT" && \
	  OPTIONS="--enable-manager" && \
	  OPTIONS="\$OPTIONS --preview-method auto" && \
	  OPTIONS="\$OPTIONS --listen" && \
	  git pull --rebase && \
	  if [ -d venv-linux ]; then python3 -m venv --upgrade venv-linux; else python3 -m venv venv-linux; fi && \
	  . venv-linux/bin/activate && \
	  pip install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu130 && \
	  pip install -r requirements.txt && \
	  pip install -r manager_requirements.txt && \
	  pip install matrix-nio && \
	  find custom_nodes -mindepth 1 -maxdepth 1 -type d | while read n; do \
	    [ -d "\$n" ] && (cd "\$n" && { [ -d .git ] && git pull --rebase; }; [ -r requirements.txt ] && pip install -r requirements.txt) || true; \
	  done && \
	  exec python main.py \$OPTIONS'
	Restart=on-failure
	RestartSec=10s

	[Install]
	WantedBy=default.target
	EOF
	systemctl --user daemon-reload
	systemctl --user enable --now "$SERVICE"
	echo "installed and started $SERVICE"
}

uninstall() {
	systemctl --user disable --now "$SERVICE" 2>/dev/null || true
	rm -f "$FILE"
	systemctl --user daemon-reload
	echo "removed $SERVICE"
}

usage() {
	echo "usage: ${0##*/} [-h] [-u] [-l] [-p PATH]"
	echo "  -h  show this help"
	echo "  -u  disable and remove the systemd service"
	echo "  -l  show the service log"
	echo "  -p  path to ComfyUI directory (auto-detected if omitted)"
}

PATH_ARG=""
while getopts hulp: opt; do
	case $opt in
		h) usage; exit ;;
		u) uninstall; exit ;;
		l) journalctl --user -u "$SERVICE" -f; exit ;;
		p) PATH_ARG="$OPTARG" ;;
	esac
done
install "$PATH_ARG"