#!/bin/bash
# -*- coding: utf-8 -*-
"""
XiaLi 云 GPU 一键搭建脚本
=========================
适用于 AutoDL / 矩池云 / 恒源云 等国内云GPU平台

用法:
  bash setup_cloud.sh                    # 完整搭建
  bash setup_cloud.sh --train-only       # 只拉代码，不下载数据
"""

set -e

echo "========================================"
echo "  XiaLi 云 GPU 训练环境搭建"
echo "========================================"

REPO_URL="https://gitee.com/shuiyu1123/xiaoli-llm-chat.git"
WORK_DIR="/root/xiaoli-llm-chat"

# 1. 克隆代码
echo ""
echo "[1/4] 克隆代码..."
if [ -d "$WORK_DIR" ]; then
    echo "  目录已存在，拉取更新..."
    cd "$WORK_DIR" && git pull
else
    git clone "$REPO_URL" "$WORK_DIR"
    cd "$WORK_DIR"
fi

# 2. 安装依赖
echo ""
echo "[2/4] 安装依赖..."
pip install torch numpy -q
# 检查 torch 是否有 CUDA
python3 -c "import torch; print(f'CUDA可用: {torch.cuda.is_available()}, GPU数: {torch.cuda.device_count()}')" 2>/dev/null || {
    echo "  ⚠️ 检测到无 CUDA，尝试安装 GPU 版 PyTorch..."
    pip install torch --index-url https://download.pytorch.org/whl/cu118 -q
}
pip install modelscope -q

# 确认词表存在
echo ""
echo "[3/4] 检查词表..."
ls -la vocab.json 2>/dev/null || {
    echo "  ⚠️ vocab.json 不存在！请确认仓库完整"
    exit 1
}
python3 -c "
import json
with open('vocab.json') as f:
    v = json.load(f)
print(f'  词表: {len(v)} tokens, vocab_size={max(v.values())+1}')
"

# 4. 确认训练脚本
echo ""
echo "[4/4] 确认训练脚本..."
ls -la train_cloud.py && echo "  ✅ train_cloud.py 就绪"

# 显存信息
echo ""
echo "========================================"
echo "  环境就绪！"
echo "========================================"
python3 -c "
import torch
if torch.cuda.is_available():
    total = torch.cuda.get_device_properties(0).total_mem / 1024**3
    name = torch.cuda.get_device_name(0)
    print(f'  GPU: {name} ({total:.1f} GiB)')
else:
    print('  ⚠️ 未检测到 GPU！')
"
echo ""
echo "启动训练命令:"
echo "  cd $WORK_DIR"
echo ""
echo "  # 单卡训练（完整流程）"
echo "  python3 train_cloud.py --epochs 5 --align-epochs 20"
echo ""
echo "  # 从检查点续训"
echo "  python3 train_cloud.py --resume --epochs 10"
echo ""
echo "  # 24G显存用更小batch"
echo "  python3 train_cloud.py --batch-size 32"
echo ""
echo "========================================"