#!/bin/bash
set -euo pipefail
# =============================================================================
# VLM (Qwen2.5-VL-3B) 依赖安装脚本 — 国内镜像优化版
# =============================================================================
# 运行环境：ECS GPU 容器（/workspace 挂载在 EFS 上持久化）
# 前置条件：NVIDIA GPU 可用，CUDA 驱动已安装
#
# 国内镜像策略：
#   - pip 包：阿里云镜像
#   - HuggingFace 模型：hf-mirror.com（默认）或魔搭 ModelScope
#   - PyTorch：优先阿里云，若缺失则回退官方源
#
# 如需使用官方源（境外环境）：
#   USE_CHINA_MIRROR=false bash scripts/install_vlm.sh
# =============================================================================

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

# 国内镜像开关（默认启用）
USE_CHINA_MIRROR="${USE_CHINA_MIRROR:-true}"

# pip 镜像
PIP_MIRROR_ALIYUN="https://mirrors.aliyun.com/pypi/simple/"
PIP_MIRROR_TSINGHUA="https://pypi.tuna.tsinghua.edu.cn/simple"
PIP_MIRROR_OFFICIAL="https://pypi.org/simple"

# HuggingFace 镜像
HF_MIRROR="https://hf-mirror.com"

# 魔搭 ModelScope（阿里云出品，国内速度极快）
MODELSCOPE_CACHE="${ROOT}/.cache/modelscope"

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

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

# -----------------------------------------------------------------------------
# 配置镜像
# -----------------------------------------------------------------------------
if [[ "$USE_CHINA_MIRROR" == "true" ]]; then
    log_info "国内镜像模式已启用（USE_CHINA_MIRROR=true）"
    PIP_INDEX="--index-url $PIP_MIRROR_ALIYUN"
    HF_ENDPOINT="$HF_MIRROR"
    export HF_ENDPOINT
    export PIP_INDEX_URL="$PIP_MIRROR_ALIYUN"
    # 让 huggingface_hub 使用国内 endpoint
    export HF_HUB_ENABLE_HF_TRANSFER=0
else
    log_info "官方源模式（USE_CHINA_MIRROR=false）"
    PIP_INDEX=""
    HF_ENDPOINT="https://huggingface.co"
fi

# -----------------------------------------------------------------------------
# 0. 环境检查
# -----------------------------------------------------------------------------
log_step "0. 检查 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_step "1. 安装 PyTorch 2.1.2 + CUDA 12.1 ..."

    if [[ "$USE_CHINA_MIRROR" == "true" ]]; then
        # 国内环境：先尝试阿里云 pip 镜像安装 PyTorch
        # 注意：阿里云上 PyTorch CUDA wheel 可能滞后，若失败则回退官方源
        log_info "尝试从阿里云镜像安装 PyTorch..."
        pip install --target="$SITE_PACKAGES" --upgrade \
            torch==2.1.2 torchvision==0.16.2 torchaudio==2.1.2 \
            --index-url "$PIP_MIRROR_ALIYUN" \
            --no-cache-dir || {
            log_warn "阿里云镜像安装 PyTorch 失败，回退到官方源..."
            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
        }
    else
        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
fi

# -----------------------------------------------------------------------------
# 2. 安装 transformers / accelerate / qwen_vl_utils 等
# -----------------------------------------------------------------------------
log_step "2. 安装 Python 依赖包..."

# 先升级 pip 本身（国内镜像）
pip install $PIP_INDEX --upgrade pip setuptools wheel --no-cache-dir || true

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

# -----------------------------------------------------------------------------
# 3. 可选：安装 Flash Attention 2 (Ampere+ GPU)
# -----------------------------------------------------------------------------
log_step "3. 检查 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" $PIP_INDEX --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. 验证 Python 包
# -----------------------------------------------------------------------------
log_step "4. 验证 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_step "5. 预下载 VLM 模型: $VLM_MODEL ..."

mkdir -p "$VLM_MODEL_CACHE"

if [[ "$USE_CHINA_MIRROR" == "true" ]]; then
    log_info "尝试通过魔搭 ModelScope 下载模型（国内 fastest）..."

    # 先安装 modelscope
    pip install --target="$SITE_PACKAGES" $PIP_INDEX --upgrade \
        "modelscope>=1.15.0" --no-cache-dir

    python3 -c "
import sys
sys.path.insert(0, '$SITE_PACKAGES')

from modelscope import snapshot_download

# Qwen 模型在魔搭的命名空间：qwen/Qwen2.5-VL-3B-Instruct
model_id = '$VLM_MODEL'
# 如果是 HuggingFace 格式的 model_id，转换为魔搭格式
if model_id.startswith('Qwen/'):
    modelscope_id = model_id.replace('Qwen/', 'qwen/')
else:
    modelscope_id = model_id

print(f'使用魔搭下载: {modelscope_id}')
model_dir = snapshot_download(
    modelscope_id,
    cache_dir='$MODELSCOPE_CACHE',
)
print(f'模型已下载到: {model_dir}')

# 记录本地模型路径，供 run_all.sh / detector.py 直接使用
path_file = '/workspace/.vlm_model_path'
with open(path_file, 'w') as f:
    f.write(f'export VLM_MODEL={model_dir}\\n')
print(f'本地模型路径已记录到: {path_file}')

# 验证可以加载
from transformers import Qwen2_5_VLForConditionalGeneration, AutoProcessor
print('验证模型加载...')
model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
    model_dir,
    torch_dtype='auto',
    device_map='auto',
    trust_remote_code=True,
)
processor = AutoProcessor.from_pretrained(
    model_dir,
    trust_remote_code=True,
)
print('模型加载验证成功！')
"
else
    log_info "使用 HuggingFace 官方源下载模型..."
    export HF_HOME="$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
from huggingface_hub import snapshot_download

print(f'正在下载模型: $VLM_MODEL ...')
# 先下载到本地缓存
model_dir = snapshot_download(
    repo_id='$VLM_MODEL',
    cache_dir='$VLM_MODEL_CACHE/hub',
)
print(f'模型已下载到: {model_dir}')

# 记录本地模型路径
path_file = '/workspace/.vlm_model_path'
with open(path_file, 'w') as f:
    f.write(f'export VLM_MODEL={model_dir}\\n')
print(f'本地模型路径已记录到: {path_file}')

# 验证加载
model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
    model_dir,
    torch_dtype='auto',
    device_map='auto',
    trust_remote_code=True,
)
processor = AutoProcessor.from_pretrained(
    model_dir,
    trust_remote_code=True,
)
print('模型加载验证成功！')
"
fi

# -----------------------------------------------------------------------------
# 6. 完成
# -----------------------------------------------------------------------------
log_step "6. 安装完成！"
log_info ""
log_info "============================================"
log_info "  VLM 安装摘要"
log_info "============================================"
log_info "模型: $VLM_MODEL"
log_info "HF 缓存: $VLM_MODEL_CACHE"
log_info "魔搭缓存: $MODELSCOPE_CACHE"
log_info "Python 包路径: $SITE_PACKAGES"
log_info "国内镜像: $USE_CHINA_MIRROR"
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 "============================================"