#!/bin/bash
# Phantom HuggingFace Spaces 安装脚本
# 仅 Inference 模式 - 跳过 training 相关依赖

set -e

PHANTOM_DIR="/home/user/app/phantom"
LOG_FILE="/tmp/phantom_setup.log"

log() {
    echo "[$(date +'%H:%M:%S')] $1" | tee -a "$LOG_FILE"
}

log "🚀 开始配置 Phantom 环境 (Inference Only)"

# 检查 phantom 目录
if [ ! -d "$PHANTOM_DIR" ]; then
    log "❌ Phantom 目录不存在"
    exit 1
fi

cd "$PHANTOM_DIR"

# ========== 安装 Inference 必需依赖 ==========

# 1. 安装 PyTorch (如果尚未安装)
if ! python -c "import torch" 2>/dev/null; then
    log "📦 安装 PyTorch..."
    pip install -q torch==2.1.0 torchvision==0.16.0 --index-url https://download.pytorch.org/whl/cu121
fi

# 2. SAM2 (分割模型)
if [ ! -f "/tmp/.sam2_installed" ]; then
    log "📦 安装 SAM2..."
    cd "$PHANTOM_DIR/submodules/sam2"
    pip install -q -e . 2>&1 | tee -a "$LOG_FILE" || log "⚠️ SAM2 警告"
    touch /tmp/.sam2_installed
    log "✅ SAM2 完成"
fi

# 3. HaMeR (手部姿态估计)
if [ ! -f "/tmp/.hamer_installed" ]; then
    log "📦 安装 HaMeR..."
    cd "$PHANTOM_DIR/submodules/phantom-hamer"
    pip install -q -e .[all] 2>&1 | tee -a "$LOG_FILE" || log "⚠️ HaMeR 警告"

    # 安装 ViTPose
    if [ -d "third-party/ViTPose" ]; then
        log "📦 安装 ViTPose..."
        pip install -q -e third-party/ViTPose 2>&1 | tee -a "$LOG_FILE" || true
    fi

    touch /tmp/.hamer_installed
    log "✅ HaMeR 完成"
fi

# 4. 下载 HaMeR demo 数据
if [ ! -d "$PHANTOM_DIR/submodules/phantom-hamer/_DATA/hamer_demo_data" ]; then
    log "📥 下载 HaMeR demo 数据..."
    cd "$PHANTOM_DIR/submodules/phantom-hamer"
    mkdir -p _DATA && cd _DATA
    if [ ! -f "hamer_demo_data.tar.gz" ]; then
        wget -q https://www.cs.utexas.edu/~pavlakos/hamer/data/hamer_demo_data.tar.gz || log "⚠️ HaMeR 数据下载失败"
    fi
    if [ -f "hamer_demo_data.tar.gz" ]; then
        tar --warning=no-unknown-keyword -xzf hamer_demo_data.tar.gz 2>&1 | tee -a "$LOG_FILE" || true
        rm -f hamer_demo_data.tar.gz
        log "✅ HaMeR 数据完成"
    fi
fi

# 5. MMCV (仅基础版本，inference 够用)
if [ ! -f "/tmp/.mmcv_installed" ]; then
    log "📦 安装 MMCV..."
    pip install -q mmcv==1.3.9 2>&1 | tee -a "$LOG_FILE" || true
    # 尝试安装 mmcv-full，失败也没关系
    pip install -q mmcv-full -f https://download.openmmlab.com/mmcv/dist/cu121/torch2.1/index.html 2>&1 | tee -a "$LOG_FILE" || log "⚠️ MMCV-full 跳过，使用基础版本"
    touch /tmp/.mmcv_installed
    log "✅ MMCV 完成"
fi

# 6. E2FGVI (视频修复)
E2FGVI_DIR="$PHANTOM_DIR/submodules/phantom-E2FGVI/E2FGVI/release_model"
if [ ! -f "$E2FGVI_DIR/E2FGVI-HQ.pth" ]; then
    log "📥 下载 E2FGVI 权重..."
    mkdir -p "$E2FGVI_DIR"
    cd "$E2FGVI_DIR"
    pip install -q gdown
    gdown --fuzzy "https://drive.google.com/file/d/10wGdKSUOie0XmCr8SQ2A2FeDe-mfn5w3/view?usp=sharing" 2>&1 | tee -a "$LOG_FILE" || log "⚠️ E2FGVI 权重下载失败"
    log "✅ E2FGVI 权重完成"
fi

if [ ! -f "/tmp/.e2fgvi_installed" ]; then
    log "📦 安装 E2FGVI..."
    cd "$PHANTOM_DIR/submodules/phantom-E2FGVI"
    pip install -q -e . 2>&1 | tee -a "$LOG_FILE"
    touch /tmp/.e2fgvi_installed
    log "✅ E2FGVI 完成"
fi

# ========== 跳过 Training 依赖 ==========
# 以下包仅用于训练，inference 不需要:
# - phantom-robosuite (机器人仿真)
# - phantom-robomimic (机器人学习)
log "⏭️ 跳过 Training 依赖 (robosuite, robomimic)"

# 7. 其他 inference 依赖
log "📦 安装其他依赖..."
pip install -q joblib mediapy 2>&1 | tee -a "$LOG_FILE" || true
pip install -q transformers==4.42.4 2>&1 | tee -a "$LOG_FILE" || true
pip install -q PyOpenGL==3.1.4 Rtree protobuf==3.20.0 2>&1 | tee -a "$LOG_FILE" || true
pip install -q hydra-core==1.3.2 omegaconf==2.3.0 2>&1 | tee -a "$LOG_FILE" || true
pip install -q numpy==1.26.4 2>&1 | tee -a "$LOG_FILE" || true
# open3d 体积大，尝试安装但不强求
pip install -q open3d 2>&1 | tee -a "$LOG_FILE" || log "⚠️ open3d 跳过"

# 8. Phantom 主包
if [ ! -f "/tmp/.phantom_pkg_installed" ]; then
    log "📦 安装 Phantom 主包..."
    cd "$PHANTOM_DIR"
    pip install -q -e . 2>&1 | tee -a "$LOG_FILE"
    touch /tmp/.phantom_pkg_installed
    log "✅ Phantom 主包完成"
fi

# 9. 下载示例数据（可选）
SAMPLE_DATA_DIR="$PHANTOM_DIR/data/raw"
if [ ! -d "$SAMPLE_DATA_DIR/pick_and_place" ]; then
    log "📥 下载示例数据..."
    mkdir -p "$SAMPLE_DATA_DIR"
    cd "$SAMPLE_DATA_DIR"
    wget -q https://download.cs.stanford.edu/juno/phantom/pick_and_place.zip || log "⚠️ 示例数据下载失败"
    if [ -f "pick_and_place.zip" ]; then
        unzip -q pick_and_place.zip
        rm -f pick_and_place.zip
        log "✅ 示例数据完成"
    fi
fi

# 10. 检查 MANO 模型
MANO_DIR="$PHANTOM_DIR/submodules/phantom-hamer/_DATA/data/mano"
mkdir -p "$MANO_DIR"

# 检查是否已存在（可能用户已经放在仓库里了）
if [ -f "$MANO_DIR/MANO_LEFT.pkl" ] && [ -f "$MANO_DIR/MANO_RIGHT.pkl" ]; then
    log "✅ MANO 模型已就绪"
else
    log "⚠️ MANO 模型缺失！"
    log "   请将文件放到: $MANO_DIR"
fi

# 标记完成
touch /tmp/.phantom_ready

log "🎉 Phantom 环境配置完成 (Inference Only)"
log "📝 日志文件: $LOG_FILE"