#!/bin/bash
# DreamZero 模型权重一键下载脚本
#
# 用法:
#   bash scripts/data/download_checkpoints.sh
#   bash scripts/data/download_checkpoints.sh --mirror   # 使用 HF 镜像（中国用户）
#   bash scripts/data/download_checkpoints.sh --offline /path/to/cache  # 从本地缓存加载
#
# 下载内容:
#   - Wan2.2-TI2V-5B:     Wan2.2 视频 DiT 骨干 (~32G)
#   - umt5-xxl:           T5 tokenizer (~49G)
#   - CLIP encoder:       图像编码器 (~2.4G)

set -euo pipefail

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

MIRROR=false
OFFLINE_CACHE=""

# Parse arguments
while [[ $# -gt 0 ]]; do
    case $1 in
        --mirror)
            MIRROR=true
            shift
            ;;
        --offline)
            OFFLINE_CACHE="$2"
            shift 2
            ;;
        *)
            echo "Unknown option: $1"
            echo "Usage: $0 [--mirror] [--offline /path/to/cache]"
            exit 1
            ;;
    esac
done

# ============ 网络配置 ============
if [ "$MIRROR" = true ]; then
    export HF_ENDPOINT=https://hf-mirror.com
    echo ">>> 使用 HF 镜像: hf-mirror.com"
fi

# ============ 离线模式 ============
if [ -n "$OFFLINE_CACHE" ]; then
    if [ ! -d "$OFFLINE_CACHE" ]; then
        echo "ERROR: 离线缓存目录不存在: $OFFLINE_CACHE"
        exit 1
    fi
    echo ">>> 离线模式，从 $OFFLINE_CACHE 复制..."
    cp -r "$OFFLINE_CACHE"/* "$CHECKPOINT_DIR/"
    echo ">>> 复制完成！"
    exit 0
fi

# ============ 下载 Wan2.2-TI2V-5B ============
WAN22_DIR="$CHECKPOINT_DIR/Wan2.2-TI2V-5B"
if [ -d "$WAN22_DIR" ] && [ -n "$(ls -A "$WAN22_DIR" 2>/dev/null)" ]; then
    echo "[SKIP] Wan2.2-TI2V-5B 已存在: $WAN22_DIR"
else
    echo "[DOWNLOAD] Wan2.2-TI2V-5B → $WAN22_DIR"
    mkdir -p "$WAN22_DIR"
    huggingface-cli download Wan-AI/Wan2.2-TI2V-5B --local-dir "$WAN22_DIR"
    echo "[OK] Wan2.2-TI2V-5B 下载完成"
fi

# ============ 下载 umt5-xxl ============
TOKENIZER_DIR="$CHECKPOINT_DIR/umt5-xxl"
if [ -d "$TOKENIZER_DIR" ] && [ -n "$(ls -A "$TOKENIZER_DIR" 2>/dev/null)" ]; then
    echo "[SKIP] umt5-xxl 已存在: $TOKENIZER_DIR"
else
    echo "[DOWNLOAD] umt5-xxl → $TOKENIZER_DIR"
    mkdir -p "$TOKENIZER_DIR"
    huggingface-cli download google/umt5-xxl --local-dir "$TOKENIZER_DIR"
    echo "[OK] umt5-xxl 下载完成"
fi

# ============ 下载 CLIP encoder ============
CLIP_DIR="$CHECKPOINT_DIR/clip-encoder"
CLIP_FILE="$CLIP_DIR/models_clip_open-clip-xlm-roberta-large-vit-huge-14.pth"
if [ -f "$CLIP_FILE" ]; then
    echo "[SKIP] CLIP encoder 已存在: $CLIP_FILE"
else
    echo "[DOWNLOAD] CLIP encoder → $CLIP_DIR"
    mkdir -p "$CLIP_DIR"

    # 尝试从 dreamzero 独立仓库下载
    if huggingface-cli download poet70/clip-encoder \
        --local-dir "$CLIP_DIR" 2>/dev/null; then
        echo "[OK] CLIP encoder 下载完成"
    else
        echo "[WARN] 独立 CLIP 仓库暂不可用，尝试从 Wan2.1 提取..."
        # 后备：下载 Wan2.1-I2V-14B-480P（包含 CLIP 权重）
        WAN21_DIR="$CHECKPOINT_DIR/Wan2.1-I2V-14B-480P"
        if [ ! -d "$WAN21_DIR" ]; then
            huggingface-cli download Wan-AI/Wan2.1-I2V-14B-480P --local-dir "$WAN21_DIR"
        fi
        # 复制 CLIP 权重
        if [ -f "$WAN21_DIR/models_clip_open-clip-xlm-roberta-large-vit-huge-14.pth" ]; then
            cp "$WAN21_DIR/models_clip_open-clip-xlm-roberta-large-vit-huge-14.pth" "$CLIP_DIR/"
            echo "[OK] CLIP encoder 从 Wan2.1 提取完成"
        else
            echo "[ERROR] 无法获取 CLIP encoder"
            echo "请手动从 DreamZero-DROID 检查点提取:"
            echo "  cp /path/to/DreamZero-DROID/models_clip_*.pth $CLIP_DIR/"
            exit 1
        fi
    fi
fi

echo ""
echo "============================================"
echo " 所有权重下载完成！"
echo "============================================"
echo "Wan2.2:      $WAN22_DIR"
echo "umt5-xxl:    $TOKENIZER_DIR"
echo "CLIP:        $CLIP_DIR"
echo "============================================"