#!/usr/bin/env bash
set -euo pipefail

COMFY_DIR="/opt/dlami/nvme/ComfyUI"
SERVICE_NAME="comfyui"
LOG_DIR="/var/log/comfyui"
ENSURE_SCRIPT="/usr/local/bin/desineuron-ensure-comfyui.sh"

if ! command -v git >/dev/null 2>&1; then
  sudo apt-get update
  sudo apt-get install -y git
fi

sudo tee "${ENSURE_SCRIPT}" >/dev/null <<'EOF'
#!/usr/bin/env bash
set -euo pipefail

COMFY_DIR="/opt/dlami/nvme/ComfyUI"

sudo mkdir -p /opt/dlami/nvme
sudo chown -R ubuntu:ubuntu /opt/dlami/nvme

if [ ! -d "${COMFY_DIR}/.git" ]; then
  rm -rf "${COMFY_DIR}"
  git clone https://github.com/comfyanonymous/ComfyUI.git "${COMFY_DIR}"
else
  git -C "${COMFY_DIR}" fetch --all --prune
  git -C "${COMFY_DIR}" reset --hard origin/master
fi

python3 -m pip install -r "${COMFY_DIR}/requirements.txt"

if [ ! -f "${COMFY_DIR}/main.py" ]; then
  echo "ComfyUI main.py missing after ensure step" >&2
  exit 1
fi
EOF
sudo chmod 0755 "${ENSURE_SCRIPT}"

sudo -u ubuntu "${ENSURE_SCRIPT}"

sudo mkdir -p "${LOG_DIR}"
sudo chown -R ubuntu:ubuntu "${LOG_DIR}"

sudo tee /etc/systemd/system/${SERVICE_NAME}.service >/dev/null <<'EOF'
[Unit]
Description=ComfyUI GPU Service
After=network-online.target
Wants=network-online.target

[Service]
Type=simple
User=ubuntu
Group=ubuntu
WorkingDirectory=/opt/dlami/nvme/ComfyUI
Environment=HOME=/home/ubuntu
Environment=PYTHONUNBUFFERED=1
ExecStartPre=/usr/local/bin/desineuron-ensure-comfyui.sh
ExecStart=/usr/bin/python3 /opt/dlami/nvme/ComfyUI/main.py --listen 0.0.0.0 --port 8188 --disable-auto-launch
Restart=always
RestartSec=5
StandardOutput=append:/var/log/comfyui/service.log
StandardError=append:/var/log/comfyui/service.log

[Install]
WantedBy=multi-user.target
EOF

sudo systemctl daemon-reload
sudo systemctl enable --now "${SERVICE_NAME}.service"
sleep 5
sudo systemctl --no-pager --full status "${SERVICE_NAME}.service"