#!/bin/bash

start_comfyui() {
    # 设置代理
    if ! source /etc/waas-script/proxy.sh; then
        echo "未加载代理"
    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 --disable-metadata --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 -n "." # 使用点号表示等待进度
    done
    echo # 换行
    # 验证启动状态
    if [ "$STARTED" -eq 1 ]; then
        echo "ComfyUI 启动成功！ 日志: $LOG_FILE"
    else
        echo "错误: ComfyUI 启动失败!" >&2
        echo "请联系智算云扉官方客服..." >&2
        return 1
    fi
    return 0
}

start_encrypt_services() {
    local python_cmd="/root/miniconda3/bin/python"
    local encrypt_script="/root/comfyui/ComfyUI/main1.py"
    local tools_script="/.tools/main"
    local image_tag="COMFYUI_ENCRYPT"
    local log_dir="/root/comfyui/ComfyUI/logs"
    local encrypt_port=8189
    local tool_port=8089
    cd /root/comfyui/ComfyUI/
    # 检查环境变量
    if [ "$IMAGE_TAG" != "$image_tag" ]; then
        return 0
    fi
    echo "启动加密服务..."
    # 检查必需文件
    required_files=("$encrypt_script" "$tools_script")
    missing_files=()
    for file in "${required_files[@]}"; do
        if [ ! -f "$file" ]; then
            missing_files+=("$file")
        fi
    done
    if [ ${#missing_files[@]} -gt 0 ]; then
        echo "错误：以下必需文件缺失:"
        printf '  %s\n' "${missing_files[@]}"
        return 1
    fi
    # 准备日志目录
    mkdir -p "$log_dir"
    local timestamp=$(date +"%Y%m%d_%H%M%S")
    if ! ss -tln | grep -q ":$encrypt_port "; then
        # 启动加密脚本（使用nohup防止关闭终端时退出）
        nohup $python_cmd "$encrypt_script" >>"$log_dir/encrypt_${timestamp}.log" 2>&1 &
        local encrypt_pid=$!
        echo "启动 $encrypt_script (PID: $encrypt_pid), 日志: $log_dir/encrypt_${timestamp}.log"
    fi
    if ! ss -tln | grep -q ":$tool_port "; then
        # 启动工具脚本
        nohup "$tools_script" >>"$log_dir/tool_${timestamp}.log" 2>&1 &
        local tool_pid=$!
        echo "启动 $tools_script (PID: $tool_pid), 日志: $log_dir/tool_${timestamp}.log"
    fi
    return 0
}

start_ollama() {
    local ollama_cmd="/usr/local/bin/ollama"
    local port=11434

    if [ ! -x "$ollama_cmd" ]; then
        echo "未找到 ollama 可执行文件 ($ollama_cmd)" >&2
        return 0
    fi

    echo "启动 Ollama 服务..."
    # 如果服务已经在运行，跳过
    if ss -tln | grep -q ":$port "; then
        echo "Ollama 已在端口 $port 上运行，跳过启动"
        return 0
    fi
    # 启动 Ollama（无日志、不挂终端）
    nohup "$ollama_cmd" serve >/dev/null 2>&1 &
    # 等待启动完成
    local started=0
    for ((i = 1; i <= 30; i++)); do
        if ss -tln | grep -q ":$port "; then
            started=1
            break
        fi
        sleep 1
    done
    if [ "$started" -eq 1 ]; then
        echo "Ollama 启动成功"
        return 0
    else
        echo "错误: Ollama 启动失败!" >&2
        return 1
    fi
}

start_aigate_app() {
    local python_cmd="/root/miniconda3/bin/python"
    local app_script="/root/aigate-comfyui-app/run.py"
    local portApp1=5006
    local portApp2=16539

    echo "检查 Aigate 应用启动条件..."

    echo "AIGATE_APP 环境变量: $AIGATE_APP"
    echo "应用脚本路径: $app_script"

    if [ "$AIGATE_APP" != "TRUE" ]; then
        echo "AIGATE_APP 环境变量不为 TRUE，跳过启动"
        return 0
    fi

    if [ ! -e "$app_script" ]; then
        echo "应用脚本不存在: $app_script"
        return 0
    fi

    echo "启动 Aigate 应用..."
    # 检查端口是否被占用
    if ss -tln | grep -q ":$portApp1 "; then
        echo "Aigate 应用1已在运行，跳过启动"
        return 0
    fi
    
     # 检查端口是否被占用
    if ss -tln | grep -q ":$portApp2 "; then
        echo "Aigate 应用2已在运行，跳过启动"
        return 0
    fi
    
    # 启动 Aigate 应用
    nohup $python_cmd "$app_script" > /var/log/waas/aigate_app.log 2>&1 &
    local app_pid=$!
    echo "Aigate 应用已启动 (PID: $app_pid), 日志: /var/log/waas/aigate_app.log"
}

# 启动所有服务，记录失败状态但继续执行
echo "开始启动所有服务..."

if ! start_comfyui; then
    echo "警告: ComfyUI 启动失败!" >&2
fi

if ! start_encrypt_services; then
    echo "警告: 加密服务启动失败!" >&2
fi

if ! start_ollama; then
    echo "警告: Ollama 启动失败!" >&2
fi

if ! start_aigate_app; then
    echo "警告: Aigate 应用启动失败!" >&2
fi

echo "所有服务启动流程完成。"
