#!/usr/bin/env bash
# hf-storage.sh - HuggingFace 公共存储工具
# 用于上传/下载项目文件到 HF Storage Space

set -euo pipefail

SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)"
REPO_ROOT="$(cd -- "$SCRIPT_DIR/.." && pwd)"

STORAGE_REPO="${HF_STORAGE_REPO:-}"
STORAGE_TOKEN="${HF_TOKEN:-}"

usage() {
    cat <<EOF
HuggingFace 公共存储工具

用法: hf-storage.sh <命令> [选项]

命令:
    upload <文件> [目标路径]    上传单个文件
    upload-dir <目录> [前缀]   上传整个目录
    download <文件> [本地路径] 下载文件
    list [路径]               列出存储库中的文件
    sync <本地目录> [远程目录] 同步本地目录到远程

环境变量:
    HF_STORAGE_REPO    存储库名称 (格式: username/repo)
    HF_TOKEN           HuggingFace Token (需要 write 权限)
    HF_API_URL         HF API 地址 (默认: https://huggingface.co)

示例:
    export HF_STORAGE_REPO="myuser/mystorage"
    export HF_TOKEN="hf_xxxx"

    hf-storage.sh upload ./build.tar.gz
    hf-storage.sh upload-dir ./src
    hf-storage.sh download ./build.tar.gz
    hf-storage.sh list
    hf-storage.sh sync ./dist /public

EOF
}

check_config() {
    if [[ -z "$STORAGE_REPO" ]]; then
        echo "错误: HF_STORAGE_REPO 未设置"
        echo "示例: export HF_STORAGE_REPO=\"username/storage\""
        exit 1
    fi
}

check_token() {
    if [[ -z "$STORAGE_TOKEN" ]]; then
        echo "错误: HF_TOKEN 未设置"
        echo "请设置 HF_TOKEN 环境变量"
        exit 1
    fi
}

cmd_upload() {
    local file="$1"
    local dest="${2:-}"

    check_config
    check_token

    if [[ ! -f "$file" ]]; then
        echo "错误: 文件不存在: $file"
        exit 1
    fi

    if [[ -z "$dest" ]]; then
        dest=$(basename "$file")
    fi

    echo "上传文件: $file -> $STORAGE_REPO/$dest"

    python3 -c "
from huggingface_hub import HfApi
api = HfApi()
api.upload_file(
    path_or_fileobj='$file',
    path_in_repo='$dest',
    repo_id='$STORAGE_REPO',
    repo_type='space',
)
" && echo "上传成功!" || echo "上传失败!"
}

cmd_upload_dir() {
    local dir="$1"
    local prefix="${2:-}"

    check_config
    check_token

    if [[ ! -d "$dir" ]]; then
        echo "错误: 目录不存在: $dir"
        exit 1
    fi

    echo "上传目录: $dir -> $STORAGE_REPO/${prefix:-.}"

    find "$dir" -type f | while read -r file; do
        local rel_path="${file#$dir/}"
        local dest
        if [[ -n "$prefix" ]]; then
            dest="$prefix/$rel_path"
        else
            dest="$rel_path"
        fi
        echo "  上传: $rel_path"
        python3 -c "
from huggingface_hub import HfApi
api = HfApi()
api.upload_file(
    path_or_fileobj='$file',
    path_in_repo='$dest',
    repo_id='$STORAGE_REPO',
    repo_type='space',
)
" 2>/dev/null || echo "  失败: $rel_path"
    done

    echo "上传完成!"
}

cmd_download() {
    local file="$1"
    local local_path="${2:-.}"

    check_config
    check_token

    echo "下载文件: $STORAGE_REPO/$file -> $local_path"

    python3 -c "
from huggingface_hub import hf_hub_download
path = hf_hub_download(
    repo_id='$STORAGE_REPO',
    filename='$file',
    local_dir='$local_path',
    token='$STORAGE_TOKEN',
)
print(f'下载到: {path}')
"
}

cmd_list() {
    local path="${1:-}"

    check_config
    check_token

    echo "列出文件: $STORAGE_REPO${path:+/$path}"

    python3 -c "
from huggingface_hub import HfApi
api = HfApi()
files = api.list_repo_files(repo_id='$STORAGE_REPO', repo_type='space')
for f in files:
    if f.startswith('$path'):
        print(f)
"
}

cmd_sync() {
    local local_dir="$1"
    local remote_dir="${2:-}"

    check_config
    check_token

    if [[ ! -d "$local_dir" ]]; then
        echo "错误: 目录不存在: $local_dir"
        exit 1
    fi

    echo "同步目录: $local_dir -> $STORAGE_REPO${remote_dir:+/$remote_dir}"

    # 获取远程文件列表
    python3 -c "
from huggingface_hub import HfApi
api = HfApi()
files = api.list_repo_files(repo_id='$STORAGE_REPO', repo_type='space')
for f in files:
    print(f)
" > /tmp/remote_files.txt

    # 上传本地文件
    find "$local_dir" -type f | while read -r file; do
        local rel_path="${file#$local_dir/}"
        local dest
        if [[ -n "$remote_dir" ]]; then
            dest="$remote_dir/$rel_path"
        else
            dest="$rel_path"
        fi

        # 检查是否需要上传（简单比较文件名）
        if grep -q "^${remote_dir:+$remote_dir/}$rel_path$" /tmp/remote_files.txt 2>/dev/null; then
            echo "  跳过(已存在): $rel_path"
        else
            echo "  上传: $rel_path"
            python3 -c "
from huggingface_hub import HfApi
api = HfApi()
api.upload_file(
    path_or_fileobj='$file',
    path_in_repo='$dest',
    repo_id='$STORAGE_REPO',
    repo_type='space',
)
" 2>/dev/null && echo "    成功!" || echo "    失败!"
        fi
    done

    rm -f /tmp/remote_files.txt
    echo "同步完成!"
}

main() {
    if [[ $# -eq 0 ]]; then
        usage
        exit 1
    fi

    local cmd="$1"
    shift

    case "$cmd" in
        upload)
            cmd_upload "$@"
            ;;
        upload-dir)
            cmd_upload_dir "$@"
            ;;
        download)
            cmd_download "$@"
            ;;
        list)
            cmd_list "$@"
            ;;
        sync)
            cmd_sync "$@"
            ;;
        -h|--help|help)
            usage
            ;;
        *)
            echo "未知命令: $cmd"
            usage
            exit 1
            ;;
    esac
}

main "$@"