#!/bin/bash
# 一键下载数据集、预训练模型、创建软链接、下载信息文件
# 调用方式：bash HyperDrive/tools/one_step_data.sh /data/datasets/nuscenes /data/datasets/FlashOCC /workspace/FlashOCC

# 设置默认参数
NUSCENES_DATA_PATH=${1:-/data/datasets/nuscenes}
FLASHOCC_DATA_PATH=${2:-/data/datasets/FlashOCC}
WORKSPACE_PATH=${3:-/workspace/FlashOCC}
CACHE_PATH=${4:-null}  # 可选缓存路径，如/cfs/datasets/hyperdrive，下载时会优先在此目录下查找文件。

if [[ "${CACHE_PATH}" != "null" && ! -d "${CACHE_PATH}" ]]; then
    echo "⚠️ **注意：缓存路径不存在**: ${CACHE_PATH}"
    sleep 1
fi

echo "**=== 参数配置 ===**"
echo "**FlashOCC数据路径**: $FLASHOCC_DATA_PATH"
echo "**工作空间路径**: $WORKSPACE_PATH"
echo ""

# TEST_MODE=false
# 如果最后一个参数是 --test，则启用测试模式，否则启用完整模式
# last_arg="${@: -1}"
# if [ "$last_arg" = "--test" ]; then
#     TEST_MODE=true
#     echo "🧪 **测试模式已启用，将下载 mini 数据集**"
# fi

# 检查必要的依赖
check_dependencies() {
    local missing=()
    if ! command -v aria2c &>/dev/null; then
        missing+=("aria2")
    fi
    if ! command -v zstd &>/dev/null; then
        missing+=("zstd")
    fi
    if [[ ${#missing[@]} -gt 0 ]]; then
        echo "❌ **错误: 缺少必要的依赖: ${missing[*]}**"
        echo "请先安装: apt-get install -y ${missing[*]}"
        exit 1
    fi
}
check_dependencies

# 根据地域选择 COS 域名
get_domain_name() {
    local zone region
    zone=$(curl -s --connect-timeout 2 http://169.254.10.10/latest/meta-data/placement/zone 2>/dev/null)
    if [[ -n "$zone" ]]; then
        # 按 "-" 分割，取前两部分组成地域字符串
        region=$(echo "$zone" | awk -F'-' '{print $1"-"$2}')
    fi

    case "$region" in
        ap-shanghai)
            echo "https://haihub-model-1251001002.cos.ap-shanghai.myqcloud.com"
            ;;
        ap-beijing)
            echo "https://haihub-model-bj-1251001002.cos.ap-beijing.myqcloud.com"
            ;;
        ap-guangzhou)
            echo "https://haihub-model-gz-1251001002.cos.ap-guangzhou.myqcloud.com"
            ;;
        ap-nanjing)
            echo "https://haihub-model-nj-1251001002.cos.ap-nanjing.myqcloud.com"
            ;;
        *)
            # default: shanghai
            echo "https://haihub-model-bj-1251001002.cos.ap-beijing.myqcloud.com"
            ;;
    esac
}

BASE_URL="$(get_domain_name)/HyperDrive"
echo "=== 当前使用的 BASE_URL: $BASE_URL ==="

check_directory() {
    [ -d "$1" ] && [ "$(ls -A "$1" 2>/dev/null)" ]
}

check_file() {
    local target_md5="${2:-}"
    if [[ -f "$1" ]]; then
        md5=$(md5sum "$1" | awk '{ print $1 }')
        if [[ -n "$target_md5" && "$md5" != "$target_md5" ]]; then
            echo "❌ MD5 mismatch: $md5 != $target_md5"
            return 1
        fi
        return 0
    else
        return 1
    fi
}

safe_download_with_cache() {
    local url="$1" filename="$2"
    # 优先从缓存路径复制文件
    local return_code=0
    if [[ "$CACHE_PATH" != "null" && -d "$CACHE_PATH" ]]; then
        local cached_file="$CACHE_PATH/$(basename "$url")"
        cached_file="$(realpath -m "$cached_file")"  # 规范化路径，顺便避免双斜杠

        if [[ -f "$cached_file" ]]; then
            echo "**从缓存复制文件: $cached_file -> $filename **"
            rsync -a "$cached_file" "$filename"
            return_code=$?
        else
            return_code=1
            if safe_download "$url" "$filename"; then
                echo "**将下载的文件复制到缓存: $filename -> $cached_file **"
                rsync -a "$filename" "$cached_file"
                return_code=$?
            fi
        fi
    else
        safe_download "$url" "$filename"
        return_code=$?
    fi
    return $return_code
}

safe_download() {
    local url="$1" filename="$2" max_retries=3 retry_count=0

    while [[ $retry_count -lt $max_retries ]]; do
        echo "**下载尝试 $((retry_count + 1))/$max_retries**: $filename"
        rm -f "$filename"

        echo "**使用 aria2c 进行多线程下载**"
        if aria2c --enable-color=false --show-console-readout=false -x 8 -s 8 -k 1M -d "$(dirname "$filename")" -o "$(basename "$filename")" "$url" && [[ -f "$filename" ]]; then
            echo "✅ **下载完成**: $filename"
            return 0
        fi

        retry_count=$((retry_count + 1))
    done

    echo "❌ **下载失败**: $filename"
    return 1
}

setup_flashocc_data() {
    echo "**=== 设置FlashOCC数据 ===**"

    # 下载 ckpt 文件到工作区 ckpts 目录
    mkdir -p "$WORKSPACE_PATH/ckpts"
    if check_file "$WORKSPACE_PATH/ckpts/bevdet-r50-cbgs.pth" "996ff0a2ba830fd0049f5b61748cf60f"; then
        echo "✅ **bevdet-r50-cbgs ckpt 已存在，跳过下载解压**"
    else
        echo "**下载 ckpts 到: $WORKSPACE_PATH/ckpts**"
        safe_download_with_cache "$BASE_URL/ckpt/bevdet-r50-cbgs.pth" "$WORKSPACE_PATH/ckpts/bevdet-r50-cbgs.pth" || { echo "❌ **bevdet-r50-cbgs ckpt 下载失败**"; return 1; }

    echo "✅ **FlashOCC数据设置完成**"
    fi

    local target_dir="$1"

    mkdir -p "$target_dir"

    # 下载infos pkl gt 压缩包
    local pkl_train_url="$BASE_URL/flashocc_archive.tar"

    if check_directory "$target_dir/gts"; then
        num_files="$(ls -lgG "$target_dir/gts" | grep "^d" | wc -l)"
        if [[ "${num_files}" -eq 850 ]]; then
            echo "✅ **完整的数据已存在，跳过下载解压**"
            echo "如果需要重新解压，请删除目录: $target_dir/gts"
            return 0
        fi
    fi

    if check_file "$target_dir/flashocc_archive.tar" "84207b6f15c0038ad20dc8fbc4cde267"; then
        echo "**检测到已存在完整的 flashocc 压缩包，跳过下载**"
        tar -xf "$target_dir/flashocc_archive.tar" -C "$target_dir"
    else
        echo "**下载并解压 flashocc 数据到: $target_dir**"
        safe_download_with_cache "$pkl_train_url" "$target_dir/flashocc_archive.tar" || { echo "❌ **train pkl 下载失败**"; false; } && { tar -xf "$target_dir/flashocc_archive.tar" -C "$target_dir"; }
    fi
}

# 创建软链接
create_symlink() {
    local source="$1" target="$2"

    if [ -L "$target" ]; then
        [ "$(readlink "$target")" = "$source" ] && return 0
        rm -f "$target"
    elif [ -e "$target" ]; then
        rm -rf "$target"
    fi

    ln -s "$source" "$target"
    echo "✅ **软链接创建**: $(basename "$target") -> $source"
}

# # 创建NuScenes数据软链接
# create_nuscenes_symlinks() {
#     echo "**=== 创建NuScenes数据软链接 ===**"
#     mkdir -p "$2" && cd "$2"

#     # 处理可能的嵌套目录
#     local source_path="$1"
#     [ -d "$1/nuscenes" ] && source_path="$1/nuscenes"

#     # 创建核心目录软链接
#     for item in samples sweeps maps v1.0-trainval; do
#         [ -e "$source_path/$item" ] && create_symlink "$source_path/$item" "$2/$item"
#     done

#     echo "✅ **NuScenes数据软链接创建完成**"
# }

# # **第一步：检查和下载NuScenes原始数据**
# echo "**=== 第一步：检查NuScenes原始数据 ===**"
# if check_directory "$NUSCENES_DATA_PATH"; then
#     echo "✅ **NuScenes数据集已存在**"
# else
#     download_nuscenes "$NUSCENES_DATA_PATH" || { echo "❌ **NuScenes下载失败**"; exit 1; }
# fi

# **第二步：设置FlashOCC数据**
echo "**=== 第二步：设置FlashOCC数据 ===**"
setup_flashocc_data "$FLASHOCC_DATA_PATH" || { echo "❌ **FlashOCC数据设置失败**"; exit 1; }

# **第四步：在工作空间创建软链接**
echo ""
echo "**=== 第四步：在工作空间创建软链接 ===**"
[ ! -d "$WORKSPACE_PATH" ] && { echo "❌ **工作空间不存在**: $WORKSPACE_PATH"; exit 1; }

mkdir -p "$WORKSPACE_PATH/data"
create_symlink "$FLASHOCC_DATA_PATH" "$WORKSPACE_PATH/data/nuscenes"

echo ""
echo "🎉 **================================================================**"
echo "🎉 **所有步骤完成! FlashOCC环境设置成功!**"
echo "🎉 **================================================================**"
echo ""
echo "📊 **数据架构总览**:"
echo "├── 📁 **FlashOCC数据管理**: $FLASHOCC_DATA_PATH"
echo "│   ├── 🔗 samples/"
echo "│   ├── 📁 gts/"
echo "│   └── 📄 bevdetv2-nuscenes_infos_train.pkl"
echo "└── 💼 **工作空间**: $WORKSPACE_PATH"
echo "    └── 🔗 data/nuscenes -> $FLASHOCC_DATA_PATH"
