#!/bin/bash
set -euo pipefail
# =============================================================================
# VLM (Qwen2.5-VL-3B) 依赖安装脚本
# =============================================================================
# 运行环境：ECS GPU 容器（/workspace 挂载在 EFS 上持久化）
# 前置条件：NVIDIA GPU 可用，CUDA 驱动已安装
# =============================================================================

ROOT="/workspace"
SITE_PACKAGES="${ROOT}/.site-packages"
VLM_MODEL_CACHE="${ROOT}/.cache/huggingface"

RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
NC='\033[0m'

log_info()  { echo -e "${GREEN}[INFO]${NC} $*"; }
log_warn()  { echo -e "${YELLOW}[WARN]${NC} $*"; }
log_error() { echo -e "${RED}[ERROR]${NC} $*"; }

# -----------------------------------------------------------------------------
# 0. 环境检查
# -----------------------------------------------------------------------------
log_info "检查 GPU / CUDA 环境..."

if ! command -v nvidia-smi &>/dev/null; then
    log_error "nvidia-smi 未找到。请确认当前为 GPU 实例且 NVIDIA 驱动已安装。"
    exit 1
fi

nvidia-smi --query-gpu=name,memory.total,driver_version --format=csv,noheader || true

if ! python3 -c "import torch; assert torch.cuda.is_available()" 2>/dev/null; then
    log_warn "PyTorch CUDA 未安装或不可用，将重新安装 PyTorch CUDA 版..."
    NEED_PYTORCH_CUDA=1
else
    python3 -c "import torch; print(f'PyTorch {torch.__version__}, CUDA {torch.version.cuda}, Device: {torch.cuda.get_device_name(0)}')"
    NEED_PYTORCH_CUDA=0
fi

# -----------------------------------------------------------------------------
# 1. 安装 PyTorch CUDA (如需要)
# -----------------------------------------------------------------------------
if [[ "$NEED_PYTORCH_CUDA" == "1" ]]; then
    log_info "安装 PyTorch 2.1.2 + CUDA 12.1 ..."
    pip install --target="$SITE_PACKAGES" --upgrade \
        torch==2.1.2 torchvision==0.16.2 torchaudio==2.1.2 \
        --index-url https://download.pytorch.org/whl/cu121 \
        --no-cache-dir
fi

# -----------------------------------------------------------------------------
# 2. 安装 transformers / accelerate / qwen_vl_utils
# -----------------------------------------------------------------------------
log_info "安装 transformers + accelerate + qwen_vl_utils ..."

pip install --target="$SITE_PACKAGES" --upgrade \
    "transformers>=4.40.0" \
    "accelerate>=0.30.0" \
    "qwen-vl-utils>=0.0.8" \
    "pillow>=10.0.0" \
    "safetensors>=0.4.0" \
    --no-cache-dir

# -----------------------------------------------------------------------------
# 3. 可选：安装 Flash Attention 2 (Ampere+ GPU)
# -----------------------------------------------------------------------------
GPU_ARCH=$(python3 -c "
import sys, os
sys.path.insert(0, '$SITE_PACKAGES')
import torch
major, minor = torch.cuda.get_device_capability()
print(f'{major}{minor}')
" 2>/dev/null || echo "0")

if [[ "$GPU_ARCH" == "8"* ]] || [[ "$GPU_ARCH" == "9"* ]]; then
    log_info "检测到 Ampere/Hopper 架构 (sm_$GPU_ARCH)，尝试安装 Flash Attention 2..."
    pip install --target="$SITE_PACKAGES" --no-build-isolation \
        "flash-attn>=2.5.0" --no-cache-dir || log_warn "Flash Attention 安装失败，将使用 eager attention"
else
    log_warn "GPU 架构 sm_$GPU_ARCH 不支持 Flash Attention 2，跳过安装"
fi

# -----------------------------------------------------------------------------
# 4. 验证安装
# -----------------------------------------------------------------------------
log_info "验证 Python 包..."

python3 -c "
import sys
sys.path.insert(0, '$SITE_PACKAGES')
import torch
print(f'PyTorch: {torch.__version__}')
print(f'CUDA available: {torch.cuda.is_available()}')
if torch.cuda.is_available():
    print(f'Device: {torch.cuda.get_device_name(0)}')
    print(f'Memory: {torch.cuda.get_device_properties(0).total_memory / 1e9:.2f} GB')
"

python3 -c "
import sys
sys.path.insert(0, '$SITE_PACKAGES')
from transformers import Qwen2_5_VLForConditionalGeneration, AutoProcessor
print('transformers + Qwen2.5-VL classes OK')
"

python3 -c "
import sys
sys.path.insert(0, '$SITE_PACKAGES')
from qwen_vl_utils import process_vision_info
print('qwen_vl_utils OK')
"

# -----------------------------------------------------------------------------
# 5. 预下载模型（可选，推荐）
# -----------------------------------------------------------------------------
VLM_MODEL="${VLM_MODEL:-Qwen/Qwen2.5-VL-3B-Instruct}"

log_info "预下载 VLM 模型: $VLM_MODEL ..."
log_info "模型将缓存到: $VLM_MODEL_CACHE"

mkdir -p "$VLM_MODEL_CACHE"

python3 -c "
import sys
import os
sys.path.insert(0, '$SITE_PACKAGES')
os.environ['HF_HOME'] = '$VLM_MODEL_CACHE'

from transformers import Qwen2_5_VLForConditionalGeneration, AutoProcessor
print(f'正在下载模型: $VLM_MODEL ...')
model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
    '$VLM_MODEL',
    torch_dtype='auto',
    device_map='auto',
    trust_remote_code=True,
)
processor = AutoProcessor.from_pretrained(
    '$VLM_MODEL',
    trust_remote_code=True,
)
print('模型下载完成！')
"

# -----------------------------------------------------------------------------
# 6. 完成
# -----------------------------------------------------------------------------
log_info "安装完成！"
log_info ""
log_info "============================================"
log_info "  VLM 安装摘要"
log_info "============================================"
log_info "模型: $VLM_MODEL"
log_info "缓存路径: $VLM_MODEL_CACHE"
log_info "Python 包路径: $SITE_PACKAGES"
log_info ""
log_info "使用方法:"
log_info "  export VISION_BACKEND=vlm"
log_info "  export VLM_INTERVAL=5.0"
log_info "  ros2 run vision_yolo detector"
log_info ""
log_info "如需使用 AWQ 量化版（显存 3-4GB）:"
log_info "  export VLM_MODEL=Qwen/Qwen2.5-VL-3B-Instruct-AWQ"
log_info "============================================"