#!/bin/bash

# ComfyUI 环境设置和启动脚本
# 用于创建conda环境、安装依赖并启动ComfyUI

set -e  # 遇到错误立即退出

# 颜色输出
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
RED='\033[0;31m'
NC='\033[0m' # No Color

# 获取脚本所在目录
SCRIPT_DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" && pwd )"
cd "$SCRIPT_DIR"

ENV_NAME="comfyui"
PYTHON_VERSION="3.12"

# 镜像源配置
PYPI_MIRROR="https://mirrors.aliyun.com/pypi/simple/"
PYTORCH_INDEX_URL="https://mirrors.aliyun.com/pypi/simple/"

echo -e "${GREEN}========================================${NC}"
echo -e "${GREEN}ComfyUI 环境设置脚本 (Linux CPU版本)${NC}"
echo -e "${GREEN}pip使用阿里云镜像源加速下载${NC}"
echo -e "${GREEN}========================================${NC}"

# 检查conda是否安装
if ! command -v conda &> /dev/null; then
    echo -e "${RED}错误: 未找到conda，请先安装Anaconda或Miniconda${NC}"
    exit 1
fi

# 检查环境是否已存在
if conda env list | grep -q "^${ENV_NAME} "; then
    echo -e "${YELLOW}环境 ${ENV_NAME} 已存在，是否要重新创建？(y/n)${NC}"
    read -r response
    if [[ "$response" =~ ^([yY][eE][sS]|[yY])$ ]]; then
        echo -e "${YELLOW}正在删除现有环境...${NC}"
        conda env remove -n "$ENV_NAME" -y
    else
        echo -e "${GREEN}使用现有环境${NC}"
        source "$(conda info --base)/etc/profile.d/conda.sh"
        conda activate "$ENV_NAME"
        
        # 检查并安装依赖
        echo -e "${GREEN}检查依赖...${NC}"
        pip install -r requirements.txt -i "$PYPI_MIRROR"
        
        # 启动服务
        echo -e "${GREEN}启动ComfyUI...${NC}"
        python main.py
        exit 0
    fi
fi

# 创建conda环境
echo -e "${GREEN}步骤 1/4: 创建conda环境 ${ENV_NAME} (Python ${PYTHON_VERSION})...${NC}"
echo -e "${YELLOW}注意: conda优先使用默认源（镜像源可能不稳定）...${NC}"

# 使用默认conda源创建环境，添加重试机制
MAX_RETRIES=3
RETRY_COUNT=0
ENV_CREATED=0

# 临时禁用set -e以便处理重试逻辑
set +e

while [ $RETRY_COUNT -lt $MAX_RETRIES ] && [ $ENV_CREATED -eq 0 ]; do
    # 优先尝试使用默认源
    if conda create -n "$ENV_NAME" python="$PYTHON_VERSION" -y -c defaults; then
        echo -e "${GREEN}环境创建成功！${NC}"
        ENV_CREATED=1
    else
        RETRY_COUNT=$((RETRY_COUNT + 1))
        if [ $RETRY_COUNT -lt $MAX_RETRIES ]; then
            echo -e "${YELLOW}创建环境失败，正在重试 ($RETRY_COUNT/$MAX_RETRIES)...${NC}"
            sleep 2
        else
            echo -e "${RED}错误: 创建conda环境失败，请检查网络连接或conda配置${NC}"
            echo -e "${YELLOW}提示: 可以尝试手动运行: conda create -n $ENV_NAME python=$PYTHON_VERSION -y -c defaults${NC}"
            exit 1
        fi
    fi
done

# 重新启用set -e
set -e

if [ $ENV_CREATED -eq 0 ]; then
    echo -e "${RED}错误: 无法创建conda环境${NC}"
    exit 1
fi

# 激活环境
echo -e "${GREEN}步骤 2/4: 激活环境...${NC}"
source "$(conda info --base)/etc/profile.d/conda.sh"
conda activate "$ENV_NAME"

# 配置pip使用阿里云镜像源
echo -e "${GREEN}配置pip使用阿里云镜像源...${NC}"
pip config set global.index-url "$PYPI_MIRROR" 2>/dev/null || true

# 验证pip镜像源配置
if pip config get global.index-url 2>/dev/null | grep -q "aliyun"; then
    echo -e "${GREEN}pip镜像源配置成功${NC}"
else
    echo -e "${YELLOW}警告: pip镜像源配置可能失败，将使用默认源${NC}"
fi

# 安装PyTorch (CPU版本)
echo -e "${GREEN}步骤 3/4: 安装PyTorch和依赖...${NC}"
echo -e "${GREEN}安装PyTorch (CPU only版本)...${NC}"
# 尝试使用阿里云镜像安装PyTorch，如果失败则使用官方源
set +e
pip install torch torchvision torchaudio -i "$PYTORCH_INDEX_URL" > /dev/null 2>&1
PYTORCH_INSTALL_STATUS=$?
set -e

if [ $PYTORCH_INSTALL_STATUS -ne 0 ]; then
    echo -e "${YELLOW}阿里云镜像安装PyTorch失败，尝试使用官方源...${NC}"
    pip install torch torchvision torchaudio
else
    echo -e "${GREEN}PyTorch安装成功（使用阿里云镜像）${NC}"
fi

# 安装requirements.txt中的依赖
echo -e "${GREEN}安装ComfyUI依赖包...${NC}"
pip install -r requirements.txt -i "$PYPI_MIRROR"

echo -e "${GREEN}步骤 4/4: 启动ComfyUI服务 (CPU模式)...${NC}"
echo -e "${GREEN}========================================${NC}"
echo -e "${GREEN}环境设置完成！正在启动ComfyUI...${NC}"
echo -e "${GREEN}服务启动后，请在浏览器中访问显示的地址${NC}"
echo -e "${GREEN}========================================${NC}"

# 启动ComfyUI (CPU模式)
python main.py --cpu

