#!/bin/bash

# 默认设置代理
set_proxy=true

# 解析参数
while [[ $# -gt 0 ]]; do
    case "$1" in
        --proxy)
            if [[ $# -gt 1 && "$2" != --* ]]; then
                if [[ "$2" == "true" ]]; then
                    set_proxy=true
                    shift
                elif [[ "$2" == "false" ]]; then
                    set_proxy=false
                    shift
                else
                    echo "Warning: Invalid value for --proxy: '$2'. Expected 'true' or 'false'."
                    shift
                fi
            else
                set_proxy=true
            fi
            shift
            ;;
        *)
            echo "Warning: Unknown parameter: '$1' (ignored)"
            shift
            ;;
    esac
done

# 设置或取消代理
if [[ "$set_proxy" == "true" ]]; then
    echo "Proxy enabled"
    export https_proxy=http://192.168.20.252:7897
    export http_proxy=http://192.168.20.252:7897
else
    echo "Proxy disabled"
fi

PORT=8188
LOG_FILE="/var/log/waas/comfyui.log"
PYTHON_CMD="/root/miniconda3/bin/python"
COMFYUI_SCRIPT="/root/comfyui/ComfyUI/main.py"
MAX_WAIT=150

# 检查 8188 端口是否被占用
if ss -tlnp | grep -q ":$PORT "; then
    echo "ComfyUI 正在运行，准备关闭..."

    # 提取 PID 并杀死进程
    PID=$(ss -tlnp | grep ":$PORT " | awk '{print $6}' | cut -d',' -f2 | tr -d 'pid=')

    if [ -n "$PID" ]; then
        kill -9 "$PID" 2>/dev/null || true
        sleep 3  # 等待进程完全终止
    fi
    echo "ComfyUI 关闭完成..."
fi

# 启动 ComfyUI（后台运行）
echo "启动 ComfyUI..."
$PYTHON_CMD $COMFYUI_SCRIPT --port $PORT --listen 0.0.0.0 --enable-cors-header >> "$LOG_FILE" 2>&1 &
COMFYUI_PID=$!  # 获取进程 PID

# 等待服务启动完成
echo "等待 ComfyUI 启动完成..."
STARTED=0
for ((i=1; i<=$MAX_WAIT; i++)); do
    if ss -tln | grep -q ":$PORT "; then
        STARTED=1
        break
    fi
    sleep 2
    echo "..."  # 进度提示
done
echo

# 验证启动状态
if [ "$STARTED" -eq 1 ]; then
    echo "ComfyUI 启动成功！ 日志: $LOG_FILE"
else
    echo "错误: ComfyUI 启动失败!" >&2
    echo "请联系智算云扉官方客服..."
    exit 1
fi
