#!/bin/bash

echo "正在为FLUX风格塑形API创建Python虚拟环境..."
echo ""

# 检查Python是否已安装
if ! command -v python3 &> /dev/null; then
    echo "错误：未检测到Python安装。请安装Python 3.8或更高版本。"
    exit 1
fi

# 检查是否已存在虚拟环境
if [ -d "venv" ]; then
    echo "虚拟环境已存在。是否重新创建? (y/n)"
    read recreate
    if [[ "$recreate" =~ ^[Yy]$ ]]; then
        echo "正在删除旧的虚拟环境..."
        rm -rf venv
    else
        echo "操作已取消。"
        exit 0
    fi
fi

echo "正在创建新的虚拟环境..."
python3 -m venv venv

if [ $? -ne 0 ]; then
    echo "创建虚拟环境失败。请检查Python版本并确保已安装venv模块。"
    exit 1
fi

echo "正在激活虚拟环境并安装依赖..."
source venv/bin/activate
pip install --upgrade pip

echo ""
echo "现在需要安装PyTorch。请选择您系统上安装的CUDA版本:"
echo "1. CUDA 12.x (最新，适用于RTX 40系列等新显卡)"
echo "2. CUDA 11.8"
echo "3. CUDA 11.7"
echo "4. CUDA 11.6"
echo "5. 无CUDA (CPU版本，不推荐)"
echo "6. 手动选择其他版本"
read -p "请选择 (1-6): " cuda_choice

case $cuda_choice in
    1)
        echo "正在安装PyTorch (CUDA 12.x)..."
        pip install torch==2.2.0+cu121 torchvision==0.17.0+cu121 --index-url https://download.pytorch.org/whl/cu121
        ;;
    2)
        echo "正在安装PyTorch (CUDA 11.8)..."
        pip install torch==2.0.1+cu118 torchvision==0.15.2+cu118 --index-url https://download.pytorch.org/whl/cu118
        ;;
    3)
        echo "正在安装PyTorch (CUDA 11.7)..."
        pip install torch==2.0.1+cu117 torchvision==0.15.2+cu117 --index-url https://download.pytorch.org/whl/cu117
        ;;
    4)
        echo "正在安装PyTorch (CUDA 11.6)..."
        pip install torch==2.0.1+cu116 torchvision==0.15.2+cu116 --index-url https://download.pytorch.org/whl/cu116
        ;;
    5)
        echo "正在安装PyTorch (CPU版本)..."
        echo "警告: CPU版本将无法使用GPU加速，不推荐用于图像生成！"
        pip install torch==2.0.1 torchvision==0.15.2
        ;;
    6)
        echo "请访问 https://pytorch.org/get-started/locally/ 选择合适的PyTorch版本"
        echo "安装完成后，请运行: pip install -r requirements.txt"
        exit 0
        ;;
    *)
        echo "无效的选择，将默认安装CUDA 12.x版本 (适用于最新显卡)"
        pip install torch==2.2.0+cu121 torchvision==0.17.0+cu121 --index-url https://download.pytorch.org/whl/cu121
        ;;
esac

echo "正在安装其他依赖..."
pip install -r requirements.txt

echo ""
echo "验证PyTorch安装是否支持CUDA..."
python -c "import torch; print('CUDA是否可用:', torch.cuda.is_available()); print('CUDA版本:', torch.version.cuda if torch.cuda.is_available() else '不可用')"

echo ""
echo "虚拟环境创建成功！您可以通过以下命令激活它："
echo "source venv/bin/activate"
echo ""
echo "或者直接运行 ./startup.sh 脚本启动API服务。"
echo ""

# 设置脚本为可执行
chmod +x startup.sh
chmod +x client_example.py