#!/bin/bash

# Ollama 精简镜像下载脚本
# 默认只下载 amd64 架构文件，跳过已存在的文件
# 用法: 
#   ./ollama-mirror.sh          # 下载最新版本amd64文件
#   ./ollama-mirror.sh v0.6.2   # 下载指定版本amd64文件
#   ./ollama-mirror.sh --all    # 下载所有版本amd64文件

MIRROR_DIR="ollama_releases"
BASE_URL="https://github.com/ollama/ollama/releases"

# 创建镜像目录
mkdir -p "$MIRROR_DIR"

# 获取最新版本号
get_latest_version() {
    echo "获取最新 Ollama 版本..."
    latest_version=$(curl -s "https://api.github.com/repos/ollama/ollama/releases/latest" | grep -oP '"tag_name": "\K[^"]+')
    echo "最新版本: $latest_version"
    echo "$latest_version"
}

# 下载指定版本的文件（如果不存在）
download_file() {
    local version=$1
    local file=$2
    local url="$BASE_URL/download/$version/$file"
    local save_path="$MIRROR_DIR/$version/$file"
    
    # 如果文件已存在则跳过
    if [ -f "$save_path" ]; then
        echo "文件已存在，跳过: $save_path"
        return
    fi
    
    # 创建版本目录
    mkdir -p "$(dirname "$save_path")"
    
    echo "下载: $url => $save_path"
    if ! wget -q -O "$save_path" "$url"; then
        echo "下载失败: $file"
        # 删除可能下载不完整的文件
        rm -f "$save_path"
    fi
}

# 下载指定版本的amd64文件
download_amd64_release() {
    local version=$1
    
    echo "正在下载版本 $version (amd64)..."
    
    # 只下载amd64架构文件
    local files=(
        "ollama-linux-amd64.tgz"
        "ollama-darwin-amd64.tgz"
        "ollama-windows-amd64.zip"
        "checksums.txt"  # 保留校验文件
    )
    
    for file in "${files[@]}"; do
        download_file "$version" "$file"
    done
}

# 获取所有版本
fetch_all_versions() {
    echo "获取所有 Ollama 版本..."
    curl -s "https://api.github.com/repos/ollama/ollama/releases" | \
    grep -oP '"tag_name": "\K[^"]+' | \
    while read -r version; do
        download_amd64_release "$version"
    done
}

# 主逻辑
if [ "$1" == "--all" ]; then
    # 下载所有版本的amd64文件
    fetch_all_versions
elif [ $# -eq 1 ]; then
    # 下载指定版本的amd64文件
    download_amd64_release "$1"
else
    # 下载最新版本的amd64文件
    version=$(get_latest_version)
    download_amd64_release "$version"
fi

echo "操作完成！文件保存在 $MIRROR_DIR 目录中"