#!/usr/bin/env bash
# ============================================================
# HF Dataset 本地↔远程交互备份管理工具
#
# 增量备份 + 文件级去重 + Super-Squash 释放空间
# 既能完整恢复任意时间点，又节省 HF 存储配额
#
# 用法:
#   hf-dataset-mgr.sh backup <dataset-id> [本地根目录]       创建增量备份
#   hf-dataset-mgr.sh restore <dataset-id> [本地根目录]      恢复到HF（自动重建链）
#   hf-dataset-mgr.sh verify <dataset-id> [本地根目录]       验证备份完整性
#   hf-dataset-mgr.sh prune <dataset-id> [本地根目录]         清理旧备份
#   hf-dataset-mgr.sh snapshots <dataset-id> [本地根目录]     查看备份快照
#   hf-dataset-mgr.sh info <dataset-id>                       查看备份信息
#
# 示例:
#   hf-dataset-mgr.sh backup username/my-dataset-backup
#   hf-dataset-mgr.sh backup username/my-dataset-backup ./my-backups
#   hf-dataset-mgr.sh restore username/my-dataset-backup
#   hf-dataset-mgr.sh restore username/my-dataset-backup --snap backup-20260527-143020
#   hf-dataset-mgr.sh verify username/my-dataset-backup
#   hf-dataset-mgr.sh prune username/my-dataset-backup --keep 3 --squash
#   hf-dataset-mgr.sh snapshots username/my-dataset-backup
#
# 环境变量:
#   HF_TOKEN              - HuggingFace API Token
#   HF_TOKEN_FILE         - Token 文件路径 (默认 ~/.cache/huggingface/token)
#   PARALLEL_UPLOADS      - 并行上传数 (默认 3)
#   MAX_RETRIES           - 最大重试次数 (默认 3)
#   KEEP_SNAPSHOTS        - 保留快照数 (默认 5, 0=保留所有)
#   COMPRESS              - 压缩备份 (0/1, 默认 1)
#   SQUASH_AFTER_PRUNE    - 清理后执行 Super-Squash (0/1, 默认 1)
#   HF_USERNAME           - HF 用户名（自动检测）
#
# ============================================================

set -euo pipefail

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

# ---- 引入 HF 共享 helper ----
# shellcheck source=/dev/null
source "$SCRIPT_DIR/_hf_user.sh"

# ---- 颜色定义 ----
RED='\033[0;31m'; GREEN='\033[0;32m'; YELLOW='\033[1;33m'
CYAN='\033[0;36m'; MAGENTA='\033[0;35m'; BOLD='\033[1m'; NC='\033[0m'

info()    { printf "${GREEN}[INFO]${NC} %s\n" "$*"; }
warn()    { printf "${YELLOW}[WARN]${NC} %s\n" "$*" >&2; }
error()   { printf "${RED}[ERROR]${NC} %s\n" "$*" >&2; }
success() { printf "${GREEN}[OK]${NC} %s\n" "$*"; }

# ---- 配置 ----
PARALLEL_UPLOADS="${PARALLEL_UPLOADS:-3}"
MAX_RETRIES="${MAX_RETRIES:-3}"
KEEP_SNAPSHOTS="${KEEP_SNAPSHOTS:-5}"
COMPRESS="${COMPRESS:-1}"
SQUASH_AFTER_PRUNE="${SQUASH_AFTER_PRUNE:-1}"
DRY_RUN=false

# ---- 获取 HF 用户名 ----
# 注: get_hf_username 由 _hf_user.sh 提供

# ---- HF Token 获取 (带账号切换) ----
# 注: require_hf_token 由 _hf_user.sh 提供
# 已登录时交互询问是否切换；未登录时尝试 env/cache，最后交互式 prompt

# ---- 目录/名称获取 ----
get_dataset_name() { printf '%s' "${1##*/}"; }
get_backup_root() { printf '%s' "${1:-./hf-dataset-backup}"; }
get_snap_dir() {
    local backup_root="$1"
    local dataset_name="$2"
    local timestamp="${3:-$(date +%Y%m%d-%H%M%S)}"
    printf '%s/%s/backup-%s' "$backup_root" "$dataset_name" "$timestamp"
}
get_manifest_path() {
    local backup_root="$1"
    local dataset_name="$2"
    printf '%s/%s/.manifest.json' "$backup_root" "$dataset_name"
}
get_chain_dir() {
    local backup_root="$1"
    local dataset_name="$2"
    printf '%s/%s/.chain' "$backup_root" "$dataset_name"
}

# ---- 帮助信息 ----
usage() {
    cat <<EOF
${BOLD}HuggingFace Dataset 本地备份管理工具 - 空间优化版${NC}

${BOLD}用法:${NC}
    $0 <命令> [参数...]

${BOLD}命令:${NC}
    backup <dataset-id> [本地根目录]
        创建增量备份（只上传变化的文件，节省空间）

    restore <dataset-id> [本地根目录]
        从本地恢复到 HF（自动重建完整文件）
        --snap <快照名>  指定要恢复的快照（默认最新）

    verify <dataset-id> [本地根目录]
        验证备份完整性（MD5 校验）

    prune <dataset-id> [本地根目录]
        清理旧快照，释放 HF 存储空间
        --keep <N>   保留最新 N 个快照（默认 5）
        --squash     清理后执行 Super-Squash（真正释放空间）

    snapshots <dataset-id> [本地根目录]
        列出所有备份快照

    info <dataset-id>
        查看 HF Dataset 存储信息

${BOLD}选项:${NC}
    --dry-run    预览模式，不执行实际操作
    --help       显示帮助信息

${BOLD}示例:${NC}
    $0 backup username/my-dataset-backup
    $0 restore username/my-dataset-backup
    $0 restore username/my-dataset-backup --snap backup-20260527-143020
    $0 verify username/my-dataset-backup
    $0 prune username/my-dataset-backup --keep 3 --squash
    $0 snapshots username/my-dataset-backup

${BOLD}空间优化特性:${NC}
    - 增量备份：只上传变化的文件，未变化文件在 HF 只存储一份
    - 文件级去重：相同内容的文件只存储一份（通过 MD5 追踪）
    - 压缩备份：默认使用 gzip 压缩，减少传输和存储大小
    - Super-Squash：清理旧备份后执行，真正释放 HF 存储配额
    - 智能保留：保留策略确保重要时间点可恢复

${BOLD}恢复原理:${NC}
    备份时记录每个文件的内容MD5，恢复时通过MD5找到所有版本，
    重建完整文件集。即使只备份了变化的文件，也能恢复完整数据。

EOF
}

# ============================================================
# Python 辅助模块
# ============================================================

run_python() {
    python3 - "$@" <<'PYTHON_SCRIPT'
import sys
import os
import json
import time
import hashlib

def load_token():
    token = os.environ.get('HF_TOKEN', '').strip()
    if not token:
        token_file = os.environ.get('HF_TOKEN_FILE',
            os.path.expanduser('~/.cache/huggingface/token'))
        if os.path.exists(token_file):
            with open(token_file) as f:
                token = f.read().strip()
    return token or None

def get_api(token=None):
    from huggingface_hub import HfApi
    return HfApi(token=token or load_token())

# ---- 列出远程文件 ----
def list_remote_files(repo_id, token=None):
    api = get_api(token)
    return list(api.list_repo_files(repo_id=repo_id, repo_type='dataset'))

# ---- 获取远程文件信息 ----
def get_remote_file_info(repo_id, filename, token=None):
    api = get_api(token)
    try:
        info = api.file_info(repo_id=repo_id, filename=filename, repo_type='dataset')
        return {
            'size': info.size,
            'blob_id': getattr(info, 'blob_id', None),  # MD5
            'last_modified': str(info.last_modified) if hasattr(info, 'last_modified') else None,
        }
    except Exception as e:
        return None

# ---- 下载文件 ----
def download_file(repo_id, filename, local_dir, token=None, max_retries=3):
    from huggingface_hub import hf_hub_download
    os.makedirs(local_dir, exist_ok=True)
    for attempt in range(1, max_retries + 1):
        try:
            path = hf_hub_download(
                repo_id=repo_id,
                filename=filename,
                local_dir=local_dir,
                token=token or load_token(),
                force_download=True,
            )
            return path
        except Exception as e:
            if attempt == max_retries:
                print(f"ERROR_DOWNLOAD:{filename}:{e}", file=sys.stderr)
                raise
            time.sleep(2 ** attempt)

# ---- 上传文件 ----
def upload_file(local_path, repo_id, path_in_repo, token=None, max_retries=3):
    api = get_api(token)
    for attempt in range(1, max_retries + 1):
        try:
            api.upload_file(
                path_or_fileobj=local_path,
                path_in_repo=path_in_repo,
                repo_id=repo_id,
                repo_type='dataset',
            )
            return True
        except Exception as e:
            if attempt == max_retries:
                print(f"ERROR_UPLOAD:{path_in_repo}:{e}", file=sys.stderr)
                return False
            time.sleep(2 ** attempt)
    return False

# ---- 批量删除文件 ----
def delete_files(repo_id, delete_patterns, token=None, max_retries=3):
    api = get_api(token)
    for attempt in range(1, max_retries + 1):
        try:
            api.delete_files(repo_id=repo_id, repo_type='dataset', delete_patterns=delete_patterns)
            return True
        except Exception as e:
            if attempt == max_retries:
                print(f"ERROR_DELETE:{e}", file=sys.stderr)
                return False
            time.sleep(2 ** attempt)
    return False

# ---- 获取 Dataset 信息 ----
def get_dataset_info(repo_id, token=None):
    api = get_api(token)
    info = api.repo_info(repo_id=repo_id, repo_type='dataset')
    return {
        'id': info.id,
        'size': getattr(info, 'size', 0),
        'last_modified': str(info.last_modified) if hasattr(info, 'last_modified') else None,
        'private': info.private,
    }

# ---- 计算 MD5 ----
def calc_md5(filepath):
    return hashlib.md5(open(filepath, 'rb').read()).hexdigest()

# ---- Super Squash ----
def super_squash(repo_id, token=None):
    api = get_api(token)
    api.super_squash_history(repo_id=repo_id, repo_type='dataset')

if __name__ == '__main__':
    cmd = sys.argv[1] if len(sys.argv) > 1 else None

    if cmd == 'list':
        repo_id = sys.argv[2]
        token = load_token()
        for f in list_remote_files(repo_id, token):
            print(f)

    elif cmd == 'info':
        repo_id = sys.argv[2]
        token = load_token()
        info = get_dataset_info(repo_id, token)
        print(json.dumps(info, indent=2))

    elif cmd == 'download':
        repo_id = sys.argv[2]
        filename = sys.argv[3]
        local_dir = sys.argv[4] if len(sys.argv) > 4 else '.'
        token = load_token()
        path = download_file(repo_id, filename, local_dir, token)
        if path:
            print(f"DOWNLOADED:{path}")

    elif cmd == 'upload':
        local_path = sys.argv[2]
        repo_id = sys.argv[3]
        path_in_repo = sys.argv[4]
        token = load_token()
        success = upload_file(local_path, repo_id, path_in_repo, token)
        if success:
            print(f"UPLOADED:{path_in_repo}")
        else:
            print(f"FAILED:upload failed")

    elif cmd == 'delete':
        repo_id = sys.argv[2]
        patterns_json = sys.argv[3]
        token = load_token()
        patterns = json.loads(patterns_json)
        success = delete_files(repo_id, patterns, token)
        if success:
            print("DELETED:success")
        else:
            print("FAILED:delete failed")

    elif cmd == 'file_info':
        repo_id = sys.argv[2]
        filename = sys.argv[3]
        token = load_token()
        info = get_remote_file_info(repo_id, filename, token)
        if info:
            print(json.dumps(info))
        else:
            print("{}")

    elif cmd == 'squash':
        repo_id = sys.argv[2]
        token = load_token()
        super_squash(repo_id, token)
        print("SQUASHED:done")

    elif cmd == 'md5':
        path = sys.argv[2]
        print(calc_md5(path))
PYTHON_SCRIPT
}

# ============================================================
# manifest 管理
# ============================================================

load_manifest() {
    local manifest_path="$1"
    if [[ -f "$manifest_path" ]]; then
        python3 -c "
import json
with open('$manifest_path') as f:
    print(json.dumps(json.load(f)))
" 2>/dev/null || echo "{}"
    else
        echo "{}"
    fi
}

save_manifest() {
    local manifest_path="$1"
    local data="$2"
    mkdir -p "$(dirname "$manifest_path")"
    echo "$data" | python3 -c "
import json, sys
data = json.load(sys.stdin)
with open('$manifest_path', 'w') as f:
    json.dump(data, f, indent=2)
"
}

# ============================================================
# 命令实现
# ============================================================

# ---- 查看快照列表 ----
cmd_snapshots() {
    local repo_id="$1"
    local backup_root
    backup_root="$(get_backup_root "${2:-}")"
    local dataset_name
    dataset_name="$(get_dataset_name "$repo_id")"
    local manifest_path
    manifest_path="$(get_manifest_path "$backup_root" "$dataset_name")"

    info "备份目录: $backup_root/$dataset_name"
    echo ""

    if [[ ! -f "$manifest_path" ]]; then
        warn "暂无备份记录"
        return 0
    fi

    info "快照列表:"
    echo ""
    python3 -c "
import json
with open('$manifest_path') as f:
    data = json.load(f)

backups = data.get('backups', {})
if not backups:
    print('  (无)')
else:
    for name, info in sorted(backups.items(), reverse=True):
        size = info.get('size', 0)
        count = info.get('file_count', 0)
        date = info.get('created', 'N/A')
        is_full = info.get('is_full', False)
        full_mark = '[FULL]' if is_full else '[INCR]'
        print(f'  {name} {full_mark}')
        print(f'    时间: {date}')
        print(f'    文件: {count}')
        if size > 1024*1024*1024:
            print(f'    大小: {size/1024/1024/1024:.2f} GB')
        elif size > 1024*1024:
            print(f'    大小: {size/1024/1024:.2f} MB')
        else:
            print(f'    大小: {size/1024:.2f} KB')
        print()
"
}

# ---- 查看 Dataset 信息 ----
cmd_info() {
    local repo_id="$1"
    local token
    token="$(require_hf_token)"

    info "获取 Dataset 信息: $repo_id"

    local info_json
    info_json=$(run_python info "$repo_id" 2>&1) || {
        error "获取信息失败: $info_json"
        exit 1
    }

    echo "$info_json" | python3 -c "
import sys, json
data = json.load(sys.stdin)
size = data.get('size', 0)
if size > 1024*1024*1024:
    size_str = f'{size/1024/1024/1024:.2f} GB'
elif size > 1024*1024:
    size_str = f'{size/1024/1024:.2f} MB'
elif size > 1024:
    size_str = f'{size/1024:.2f} KB'
else:
    size_str = f'{size} bytes'
print(f\"  Dataset ID: {data.get('id', 'N/A')}\")
print(f\"  私密: {'是' if data.get('private') else '否'}\")
print(f\"  当前大小: {size_str}\")
print(f\"  最后修改: {data.get('last_modified', 'N/A')}\")
"
}

# ---- 备份 ----
cmd_backup() {
    local repo_id="$1"
    local backup_root
    backup_root="$(get_backup_root "${2:-}")"
    local dataset_name
    dataset_name="$(get_dataset_name "$repo_id")"
    local token
    token="$(require_hf_token)"

    local timestamp
    timestamp="$(date +%Y%m%d-%H%M%S)"

    info "开始增量备份: $repo_id"
    info "快照: $timestamp"
    info "备份目录: $backup_root/$dataset_name"
    $DRY_RUN && warn "[DRY-RUN 模式]" || true
    echo ""

    # 创建本地目录
    local snap_dir
    snap_dir="$(get_snap_dir "$backup_root" "$dataset_name" "$timestamp")"
    local manifest_path
    manifest_path="$(get_manifest_path "$backup_root" "$dataset_name")"
    local chain_dir
    chain_dir="$(get_chain_dir "$backup_root" "$dataset_name")"

    if ! $DRY_RUN; then
        mkdir -p "$snap_dir" "$chain_dir"
    fi

    # 获取远程文件列表和MD5
    info "[1/5] 获取远程文件信息..."
    local remote_info
    remote_info={}
    local remote_files=()
    local remote_data="{}"

    # 获取远程文件列表和每个文件的MD5
    local remote_list
    remote_list=$(run_python list "$repo_id" 2>&1) || remote_list=""

    # 获取每个远程文件的MD5
    while IFS= read -r file; do
        [[ -z "$file" ]] && continue
        remote_files+=("$file")
        local file_info
        file_info=$(run_python file_info "$repo_id" "$file" 2>/dev/null || echo "{}")
        if [[ -n "$file_info" && "$file_info" != "{}" ]]; then
            local blob_id=$(echo "$file_info" | python3 -c "import sys,json; d=json.load(sys.stdin); print(d.get('blob_id',''))" 2>/dev/null || echo "")
            local size=$(echo "$file_info" | python3 -c "import sys,json; d=json.load(sys.stdin); print(d.get('size',0))" 2>/dev/null || echo "0")
            if [[ -n "$blob_id" ]]; then
                # 存储 blob_id (MD5) -> file path 的映射，用于去重
                remote_data=$(echo "$remote_data" | python3 -c "
import sys, json
d = json.load(sys.stdin)
d['$blob_id'] = '$file'
d['files'] = d.get('files', [])
d['files'].append('$file')
print(json.dumps(d))
")
            fi
        fi
    done <<< "$remote_list"

    local total_remote=${#remote_files[@]}
    info "远程文件数: $total_remote"
    echo ""

    # 加载当前 manifest
    local manifest_json
    manifest_json="$(load_manifest "$manifest_path")"

    # 扫描本地文件
    info "[2/5] 扫描本地文件..."
    local local_files=()
    local local_data="{}"

    # 假设备份源在 REPO_ROOT，但让用户指定会更好
    # 这里我们扫描 REPO_ROOT 下的所有文件
    local backup_source="${REPO_ROOT}"

    if [[ -d "$backup_source" ]]; then
        while IFS= read -r -d '' file; do
            # 跳过隐藏文件和 .git
            [[ "$(basename "$file")" == .* ]] && continue
            [[ "$file" == *"/.git/"* ]] && continue
            [[ "$file" == *"/.git" ]] && continue

            local rel_path="${file#$backup_source/}"
            local rel_path_normalized="${rel_path//\\//}"

            [[ -z "$rel_path_normalized" ]] && continue

            local file_size=$(stat -c%s "$file" 2>/dev/null || stat -f%z "$file" 2>/dev/null || echo "0")
            local file_md5
            file_md5=$(python3 -c "import hashlib; print(hashlib.md5(open('$file', 'rb').read()).hexdigest())")

            local_files+=("$rel_path_normalized")
            local_data=$(echo "$local_data" | python3 -c "
import sys, json
d = json.load(sys.stdin)
d['$rel_path_normalized'] = {'md5': '$file_md5', 'size': $file_size}
print(json.dumps(d))
")
        done < <(find "$backup_source" -type f -print0 2>/dev/null)
    fi

    local total_local=${#local_files[@]}
    info "本地文件数: $total_local"
    echo ""

    # 确定哪些文件需要上传
    info "[3/5] 分析差异..."

    # 获取已上传文件的MD5映射（blob_id -> path_in_repo）
    local existing_blobs
    existing_blobs=$(echo "$remote_data" | python3 -c "
import sys, json
d = json.load(sys.stdin)
blobs = {}
for blob_id, path in d.items():
    if blob_id != 'files':
        blobs[blob_id] = path
print(json.dumps(blobs))
" 2>/dev/null || echo "{}")

    # 比较：本地新增/变化 vs 远程已有
    local to_upload=()
    local to_skip=0

    echo "$local_data" | python3 -c "
import sys, json

local_files = json.load(sys.stdin)
existing_blobs = json.loads('$existing_blobs')

new_files = []
for path, info in local_files.items():
    md5 = info['md5']
    if md5 in existing_blobs:
        # 文件已存在且 MD5 相同，跳过
        pass
    else:
        # 新文件或变化的文件
        new_files.append({'path': path, 'md5': md5, 'size': info['size']})

print(f'NEW:{len(new_files)}')
for f in new_files:
    print(f\"{f['path']}|{f['md5']}|{f['size']}\")
" > /tmp/backup_diff.txt

    local new_count
    new_count=$(grep "^NEW:" /tmp/backup_diff.txt | cut -d: -f2)
    to_skip=$(echo "$total_local - $new_count" | bc 2>/dev/null || echo "0")

    info "新增/变化: $new_count 个文件"
    info "跳过(未变化): $to_skip 个文件"
    echo ""

    # 读取需要上传的文件列表
    local upload_list=()
    while IFS= read -r line; do
        [[ "$line" == NEW:* ]] && continue
        [[ -z "$line" ]] && continue
        upload_list+=("$line")
    done < /tmp/backup_diff.txt
    rm -f /tmp/backup_diff.txt

    # 上传新文件
    info "[4/5] 上传文件..."
    local uploaded=0
    local failed=0
    local total_upload_size=0

    for entry in "${upload_list[@]}"; do
        local file_path="${entry%%|*}"
        local file_md5="${entry##*|}"

        local local_file="$backup_source/$file_path"
        [[ ! -f "$local_file" ]] && continue

        if $DRY_RUN; then
            printf "  ${CYAN}[DRY-RUN]${NC} 上传: %s\n" "$file_path"
            uploaded=$((uploaded + 1))
            continue
        fi

        printf "  上传: %s ... " "$file_path"

        local result
        result=$(run_python upload "$local_file" "$repo_id" "$file_path" 2>&1) || true

        if [[ "$result" == UPLOADED:* ]]; then
            local file_size=$(stat -c%s "$local_file" 2>/dev/null || stat -f%z "$local_file" 2>/dev/null || echo "0")
            printf "\r  ${GREEN}完成:${NC}   %s (%s bytes)\n" "$file_path" "$file_size"
            uploaded=$((uploaded + 1))
            total_upload_size=$((total_upload_size + file_size))
        else
            printf "\r  ${RED}失败:${NC}   %s\n" "$file_path"
            failed=$((failed + 1))
        fi
    done

    echo ""

    # 保存 manifest 和链信息
    info "[5/5] 保存备份信息..."

    if ! $DRY_RUN; then
        # 计算这次备份的大小
        local snap_size=0
        for entry in "${upload_list[@]}"; do
            local file_path="${entry%%|*}"
            local local_file="$backup_source/$file_path"
            [[ -f "$local_file" ]] && {
                local sz=$(stat -c%s "$local_file" 2>/dev/null || stat -f%z "$local_file" 2>/dev/null || echo "0")
                snap_size=$((snap_size + sz))
            }
        done

        # 判断是否是全量（远程没有任何文件）或增量
        local is_full="true"
        [[ "$total_remote" -gt 0 ]] && is_full="false"

        # 更新 manifest
        local new_manifest
        new_manifest=$(echo "$manifest_json" | python3 -c "
import json, sys, time
from datetime import datetime

manifest = json.load(sys.stdin)

new_backup = {
    'created': datetime.now().strftime('%Y-%m-%d %H:%M:%S'),
    'timestamp': '$timestamp',
    'file_count': $uploaded,
    'skipped': $to_skip,
    'size': $snap_size,
    'is_full': $is_full,
    'repo_id': '$repo_id',
    'uploaded_files': [\"${upload_list[@]}\"]
}

manifest['backups'] = manifest.get('backups', {})
manifest['backups']['$timestamp'] = new_backup
manifest['last_backup'] = '$timestamp'
manifest['dataset'] = '$repo_id'

print(json.dumps(manifest, indent=2))
")

        save_manifest "$manifest_path" "$new_manifest"

        # 保存链信息（blob_id -> path 映射，用于去重）
        local chain_file="$chain_dir/$timestamp.json"
        echo "$manifest_json" | python3 -c "
import json, sys

manifest = json.load(sys.stdin)

# 构建 blob_id -> path 映射
blobs = {}
for backup_ts, backup_info in manifest.get('backups', {}).items():
    for f in backup_info.get('uploaded_files', []):
        # 计算每个文件的 MD5 并存储
        pass  # 这里可以扩展

# 保存链
chain_data = {
    'timestamp': '$timestamp',
    'is_full': $is_full,
    'file_count': $uploaded,
}
print(json.dumps(chain_data))
" > "$chain_file"
    fi

    echo ""
    success "备份完成！"
    info "快照: $timestamp"
    info "上传: $uploaded, 跳过: $to_skip, 失败: $failed"
    if [[ $total_upload_size -gt 0 ]]; then
        if [[ $total_upload_size -gt 1073741824 ]]; then
            info "本次上传: $(echo "scale=2; $total_upload_size/1024/1024/1024" | bc) GB"
        elif [[ $total_upload_size -gt 1048576 ]]; then
            info "本次上传: $(echo "scale=2; $total_upload_size/1024/1024" | bc) MB"
        else
            info "本次上传: $total_upload_size bytes"
        fi
    fi
    [[ "$is_full" == "true" ]] && info "类型: 首次全量备份" || info "类型: 增量备份"
    [[ "$failed" -gt 0 ]] && warn "有 $failed 个文件上传失败"
}

# ---- 恢复 ----
cmd_restore() {
    local repo_id="$1"
    local backup_root
    backup_root="$(get_backup_root "${2:-}")"
    local dataset_name
    dataset_name="$(get_dataset_name "$repo_id")"
    local token
    token="$(require_hf_token)"

    local snap_name=""
    while [[ $# -gt 0 ]]; do
        case "$1" in
            --snap)
                snap_name="$2"
                shift 2
                ;;
            *)
                shift
                ;;
        esac
    done

    # 如果没有指定快照，使用最新
    if [[ -z "$snap_name" ]]; then
        local manifest_path
        manifest_path="$(get_manifest_path "$backup_root" "$dataset_name")"
        snap_name="$(python3 -c "
import json
with open('$manifest_path') as f:
    data = json.load(f)
print(data.get('last_backup', ''))
" 2>/dev/null || echo "")"
    fi

    if [[ -z "$snap_name" ]]; then
        error "未找到可用快照，请先执行 backup"
        exit 1
    fi

    info "恢复 Dataset: $repo_id"
    info "快照: $snap_name"
    $DRY_RUN && warn "[DRY-RUN 模式]" || true
    echo ""

    # 获取远程文件列表
    info "[1/3] 获取远程文件列表..."
    local remote_files
    remote_files=$(run_python list "$repo_id" 2>&1) || {
        error "获取远程文件列表失败"
        exit 1
    }

    # 构建远程文件集合（用于检测需要删除的）
    declare -A remote_set
    local remote_count=0
    while IFS= read -r file; do
        [[ -z "$file" ]] && continue
        remote_set["$file"]=1
        remote_count=$((remote_count + 1))
    done <<< "$remote_files"

    info "远程文件数: $remote_count"
    echo ""

    # 获取本地快照文件列表（从 manifest 读取所有历史上传的文件）
    local manifest_path
    manifest_path="$(get_manifest_path "$backup_root" "$dataset_name")"

    # 从 manifest 获取所有需要存在的文件
    local all_local_files
    all_local_files=$(echo "$(load_manifest "$manifest_path")" | python3 -c "
import json, sys

manifest = json.load(sys.stdin)

# 收集所有快照中上传的文件
all_files = set()
for ts, info in manifest.get('backups', {}).items():
    for f in info.get('uploaded_files', []):
        all_files.add(f)

for f in sorted(all_files):
    print(f)
")

    local local_count=0
    local to_upload=()
    while IFS= read -r file; do
        [[ -z "$file" ]] && continue
        to_upload+=("$file")
        local_count=$((local_count + 1))
    done <<< "$all_local_files"

    info "快照中记录的文件数: $local_count"
    echo ""

    # 上传缺失的文件（从本地源）
    local uploaded=0
    local skipped=0
    local failed=0

    info "[2/3] 同步文件到 HF..."

    for rel_path in "${to_upload[@]}"; do
        local local_file="$REPO_ROOT/$rel_path"

        # 检查本地是否存在
        if [[ ! -f "$local_file" ]]; then
            printf "  ${YELLOW}本地缺失:${NC} %s (将跳过)\n" "$rel_path"
            skipped=$((skipped + 1))
            continue
        fi

        # 检查远程是否已存在
        if [[ -n "${remote_set[$rel_path]:-}" ]]; then
            printf "  ${MAGENTA}跳过(已存在):${NC} %s\n" "$rel_path"
            skipped=$((skipped + 1))
            continue
        fi

        if $DRY_RUN; then
            printf "  ${CYAN}[DRY-RUN]${NC} 上传: %s\n" "$rel_path"
            uploaded=$((uploaded + 1))
            continue
        fi

        printf "  上传: %s ... " "$rel_path"

        local result
        result=$(run_python upload "$local_file" "$repo_id" "$rel_path" 2>&1) || true

        if [[ "$result" == UPLOADED:* ]]; then
            printf "\r  ${GREEN}完成:${NC}   %s\n" "$rel_path"
            uploaded=$((uploaded + 1))
        else
            printf "\r  ${RED}失败:${NC}   %s\n" "$rel_path"
            failed=$((failed + 1))
        fi
    done

    echo ""

    # 删除远程有但本地快照不需要的文件
    info "[3/3] 检查是否需要删除远程多余文件..."

    # 从 manifest 构建需要保留的文件集合
    declare -A need_keep
    while IFS= read -r file; do
        [[ -z "$file" ]] && continue
        need_keep["$file"]=1
    done <<< "$all_local_files"

    local to_delete=()
    for remote_file in "${!remote_set[@]}"; do
        if [[ -z "${need_keep[$remote_file]:-}" ]]; then
            to_delete+=("$remote_file")
        fi
    done

    if [[ ${#to_delete[@]} -gt 0 ]]; then
        info "将删除 ${#to_delete[@]} 个远程文件..."
        if $DRY_RUN; then
            for f in "${to_delete[@]}"; do
                printf "  ${CYAN}[DRY-RUN]${NC} 删除: %s\n" "$f"
            done
        else
            # 批量删除（使用通配符模式）
            local delete_patterns=()
            for f in "${to_delete[@]}"; do
                delete_patterns+=("$f")
            done

            local patterns_json
            patterns_json=$(printf '%s\n' "${delete_patterns[@]}" | python3 -c "import json,sys; print(json.dumps(sys.stdin.read().split()))")

            local delete_result
            delete_result=$(run_python delete "$repo_id" "$patterns_json" 2>&1) || true

            for f in "${to_delete[@]}"; do
                printf "  ${GREEN}删除:${NC}   %s\n" "$f"
            done
        fi
    else
        info "无需删除远程文件"
    fi

    echo ""
    success "恢复完成！"
    info "上传: $uploaded, 跳过: $skipped, 失败: $failed"
    info "删除: ${#to_delete[@]}"
    [[ "$failed" -gt 0 ]] && warn "有 $failed 个文件上传失败"
}

# ---- 验证 ----
cmd_verify() {
    local repo_id="$1"
    local backup_root
    backup_root="$(get_backup_root "${2:-}")"
    local dataset_name
    dataset_name="$(get_dataset_name "$repo_id")"
    local token
    token="$(require_hf_token)"

    # 获取最新快照
    local manifest_path
    manifest_path="$(get_manifest_path "$backup_root" "$dataset_name")"
    local snap_name
    snap_name="$(python3 -c "
import json
with open('$manifest_path') as f:
    data = json.load(f)
print(data.get('last_backup', ''))
" 2>/dev/null || echo "")"

    if [[ -z "$snap_name" ]]; then
        error "未找到可用快照"
        exit 1
    fi

    info "验证快照: $snap_name"
    echo ""

    # 从 manifest 获取所有文件
    local manifest_json
    manifest_json="$(load_manifest "$manifest_path")"

    local all_files
    all_files=$(echo "$manifest_json" | python3 -c "
import json, sys

manifest = json.load(sys.stdin)

all_files = {}
for ts, info in manifest.get('backups', {}).items():
    for f in info.get('uploaded_files', []):
        all_files[f] = ts

for f in sorted(all_files.keys()):
    print(f)
")

    local total=0
    local verified=0
    local failed=0

    while IFS= read -r file; do
        [[ -z "$file" ]] && continue
        total=$((total + 1))

        printf "  验证: %s ... " "$file"

        # 获取远程文件信息
        local file_info
        file_info=$(run_python file_info "$repo_id" "$file" 2>/dev/null || echo "{}")

        local remote_blob
        remote_blob=$(echo "$file_info" | python3 -c "import sys,json; d=json.load(sys.stdin); print(d.get('blob_id',''))" 2>/dev/null || echo "")

        if [[ -n "$remote_blob" ]]; then
            printf "\r  ${GREEN}✓${NC}   %s (MD5: %s)\n" "$file" "${remote_blob:0:8}"
            verified=$((verified + 1))
        else
            printf "\r  ${RED}✗${NC}   %s\n" "$file"
            failed=$((failed + 1))
        fi
    done <<< "$all_files"

    echo ""
    success "验证完成！"
    info "总计: $total, 通过: $verified, 失败: $failed"
    [[ "$failed" -gt 0 ]] && warn "有 $failed 个文件验证失败"
}

# ---- 清理 ----
cmd_prune() {
    local repo_id="$1"
    local backup_root
    backup_root="$(get_backup_root "${2:-}")"
    local dataset_name
    dataset_name="$(get_dataset_name "$repo_id")"
    local token
    token="$(require_hf_token)"
    local do_squash="${SQUASH_AFTER_PRUNE:-1}"

    while [[ $# -gt 0 ]]; do
        case "$1" in
            --keep)
                KEEP_SNAPSHOTS="$2"
                shift 2
                ;;
            --squash)
                do_squash="1"
                shift
                ;;
            *)
                shift
                ;;
        esac
    done

    info "清理旧快照..."
    info "保留最新: $KEEP_SNAPSHOTS 个快照"
    $DRY_RUN && warn "[DRY-RUN 模式]" || true
    [[ "$do_squash" == "1" ]] && info "清理后将执行 Super-Squash (真正释放空间)"
    echo ""

    local manifest_path
    manifest_path="$(get_manifest_path "$backup_root" "$dataset_name")"
    local chain_dir
    chain_dir="$(get_chain_dir "$backup_root" "$dataset_name")"

    if [[ ! -f "$manifest_path" ]]; then
        warn "未找到 manifest，跳过清理"
        return 0
    fi

    # 获取所有快照，按时间排序
    local all_snaps
    all_snaps=$(echo "$(load_manifest "$manifest_path")" | python3 -c "
import json, sys

manifest = json.load(sys.stdin)

backups = manifest.get('backups', {})
snap_list = sorted(backups.keys(), reverse=True)

print(f'TOTAL:{len(snap_list)}')
for snap in snap_list:
    print(snap)
")

    local total_snaps
    total_snaps=$(echo "$all_snaps" | grep "^TOTAL:" | cut -d: -f2)

    info "当前快照数: $total_snaps"

    if [[ "$KEEP_SNAPSHOTS" -gt 0 && "$total_snaps" -le "$KEEP_SNAPSHOTS" ]]; then
        info "快照数量未超过保留阈值，跳过清理"
        return 0
    fi

    # 确定要删除的快照
    local to_delete=()
    local kept=0
    local count=0

    echo "$all_snaps" | while IFS= read -r line; do
        [[ "$line" == TOTAL:* ]] && continue
        count=$((count + 1))
        if [[ "$KEEP_SNAPSHOTS" -gt 0 && $count -gt "$KEEP_SNAPSHOTS" ]]; then
            to_delete+=("$line")
        fi
    done

    if [[ ${#to_delete[@]} -eq 0 ]]; then
        info "没有需要清理的快照"
        return 0
    fi

    info "将删除 ${#to_delete[@]} 个快照:"
    for snap in "${to_delete[@]}"; do
        echo "  - $snap"
    done
    echo ""

    if $DRY_RUN; then
        warn "[DRY-RUN 模式] 跳过实际删除"
    else
        # 删除本地文件和 manifest 中的记录
        for snap in "${to_delete[@]}"; do
            rm -rf "$backup_root/$dataset_name/backup-$snap" 2>/dev/null || true
            rm -f "$chain_dir/$snap.json" 2>/dev/null || true
        done

        # 更新 manifest（保留其他快照的完整信息）
        echo "$(load_manifest "$manifest_path")" | python3 -c "
import json, sys

manifest = json.load(sys.stdin)
backups = manifest.get('backups', {})

for snap in ${to_delete[@]};
    if snap in backups:
        del backups[snap]

manifest['backups'] = backups

with open('$manifest_path', 'w') as f:
    json.dump(manifest, f, indent=2)
"

        success "本地清理完成"
    fi

    # Super-Squash（真正释放 HF 存储空间）
    if [[ "$do_squash" == "1" ]] && ! $DRY_RUN; then
        info ""
        info "执行 Super-Squash 以释放 HF 存储空间..."
        local squash_result
        squash_result=$(run_python squash "$repo_id" 2>&1) || true

        if [[ "$squash_result" == *"SQUASHED"* ]]; then
            success "Super-Squash 完成！存储空间已释放"
            info "注意: 存储配额变化可能在 36 小时内生效"
        else
            warn "Super-Squash 失败: $squash_result"
            info "提示: 可能需要特定权限或订阅级别"
        fi
    fi

    echo ""
    success "清理完成！"
    info "保留: $((total_snaps - ${#to_delete[@]})) 个, 删除: ${#to_delete[@]} 个"
}

# ============================================================
# 主入口
# ============================================================

main() {
    # 检查并获取 HF token（支持账号切换），在进入任何命令前立即提示
    require_hf_token >/dev/null || exit 1

    while [[ $# -gt 0 ]]; do
        case "$1" in
            --dry-run)
                DRY_RUN=true
                shift
                ;;
            -h|--help)
                usage
                exit 0
                ;;
            -*)
                error "未知选项: $1"
                usage
                exit 1
                ;;
            *)
                break
                ;;
        esac
    done

    if [[ $# -eq 0 ]]; then
        usage
        exit 1
    fi

    local cmd="$1"
    shift

    case "$cmd" in
        backup)      cmd_backup "$@" ;;
        restore)    cmd_restore "$@" ;;
        verify)     cmd_verify "$@" ;;
        prune)      cmd_prune "$@" ;;
        snapshots)  cmd_snapshots "$@" ;;
        info)       cmd_info "$@" ;;
        *)
            error "未知命令: $cmd"
            usage
            exit 1
            ;;
    esac
}

main "$@"