#!/bin/bash

# 虚拟环境路径 - 使用当前目录下的venv
VENV_PATH="$(pwd)/venv"

# 配置参数
BASE_PORT=8001                # 起始端口号
GPU_COUNT=3                   # 默认使用的GPU数量
SCRIPT_PATH="flux_style_shaper_api/flux_style_shaper_api.py"  # 服务脚本路径
GPU_IDS=(1 2 3)                 # 指定要使用的GPU ID

# 处理命令行参数
while [[ $# -gt 0 ]]; do
  case $1 in
    --gpus)
      GPU_COUNT="$2"
      shift 2
      ;;
    --base-port)
      BASE_PORT="$2"
      shift 2
      ;;
    *)
      echo "未知参数: $1"
      echo "用法: $0 [--gpus 数量] [--base-port 起始端口]"
      exit 1
      ;;
  esac
done

# 激活虚拟环境
echo "正在激活Python虚拟环境..."
source "$VENV_PATH/bin/activate"

# 检查虚拟环境是否成功激活
if [ $? -ne 0 ]; then
    echo "错误: 无法激活虚拟环境，请检查路径: $VENV_PATH"
    exit 1
fi

echo "已激活虚拟环境: $(which python3)"

# 创建日志目录
mkdir -p logs

# 启动多个实例
declare -a INSTANCE_PIDS
for ((i=0; i<$GPU_COUNT; i++)); do
    GPU_ID=${GPU_IDS[$i]}
    CURRENT_PORT=$((BASE_PORT + i))
    echo "启动 GPU $GPU_ID 上的服务实例，端口$CURRENT_PORT..."
    CUDA_VISIBLE_DEVICES=$GPU_ID nohup python3 $SCRIPT_PATH --port $CURRENT_PORT > logs/flux_api_gpu$GPU_ID.log 2>&1 &
    INSTANCE_PIDS[$i]=$!
    echo "GPU $GPU_ID 服务已启动，进程ID: ${INSTANCE_PIDS[$i]}"
done

echo "所有服务实例已启动!"
echo "可以通过以下地址访问服务："

# 获取本机IP
HOST_IP=$(hostname -I | awk '{print $1}')

# 显示所有服务地址
for ((i=0; i<$GPU_COUNT; i++)); do
    GPU_ID=${GPU_IDS[$i]}
    CURRENT_PORT=$((BASE_PORT + i))
    echo "  - 实例$((i+1)): http://$HOST_IP:$CURRENT_PORT (GPU $GPU_ID)"
done

echo "查看日志："
for ((i=0; i<$GPU_COUNT; i++)); do
    GPU_ID=${GPU_IDS[$i]}
    echo "  tail -f logs/flux_api_gpu$GPU_ID.log"
done

# 保持虚拟环境激活状态
# deactivate