#!/bin/bash

# ModelScope CLI批量下载脚本
# 使用说明: ./ms_download.sh -f 模型列表.cfg [-F 强制重新下载]
pip install modelscope -i https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple
# 参数解析
CONFIG_FILE=""
FORCE_DOWNLOAD=false
MODELSCOPE_CMD="modelscope download"

while getopts "f:F" opt; do
  case $opt in
    f) CONFIG_FILE="$OPTARG" ;;
    F) FORCE_DOWNLOAD=true ;;
    *) echo "Usage: $0 -f config.cfg [-F]" >&2
       exit 1
  esac
done

# 检查配置文件
if [ ! -f "$CONFIG_FILE" ]; then
  echo "Error: Config file $CONFIG_FILE not found!" >&2
  exit 1
fi

# 检查modelscope是否安装
if ! command -v modelscope &> /dev/null; then
  echo "Error: modelscope CLI not installed. Please install with: pip install modelscope" >&2
  exit 1
fi

# 读取配置文件
TOTAL=0
SUCCESS=0
FAILED=0

echo "=== Starting batch download ==="

while IFS=';' read -r model_id local_dir || [[ -n "$model_id" ]]; do
  # 跳过空行和注释
  [[ -z "$model_id" || "$model_id" =~ ^# ]] && continue
  
  ((TOTAL++))
  
  # 清理变量
  model_id=$(echo "$model_id" | xargs)
  local_dir=$(echo "$local_dir" | xargs)
  
  echo -e "\n[Progress] $TOTAL. Downloading $model_id"
  echo "[Location] $local_dir"

  # 检查目录是否存在
  if [ "$FORCE_DOWNLOAD" = false ] && [ -d "$local_dir" ]; then
    echo "[Status] Skipped (already exists)"
    ((SUCCESS++))
    continue
  fi

  # 创建目录
  mkdir -p "$local_dir" || {
    echo "[Error] Failed to create directory $local_dir" >&2
    ((FAILED++))
    continue
  }

  # 执行下载命令
  if $MODELSCOPE_CMD --model "$model_id" --local_dir "$local_dir"; then
    echo "[Status] Download successful"
    ((SUCCESS++))
  else
    echo "[Error] Download failed" >&2
    ((FAILED++))
    # 删除空目录防止残留
    rmdir "$local_dir" 2>/dev/null
  fi

done < "$CONFIG_FILE"

# 结果统计
echo -e "\n=== Download summary ==="
echo "Total:    $TOTAL"
echo "Success:  $SUCCESS"
echo "Failed:   $FAILED"

# 退出状态
if [ "$FAILED" -gt 0 ]; then
  exit 1
else
  exit 0
fi