#!/usr/bin/env bash
# ============================================================================
# restore_build.sh — 恢复已编译的 llama.cpp(CUDA/T4) 产物，跳过 40-60 分钟编译
#
# 下载优先级:
#   1) HuggingFace 公开仓（无需任何 token，推荐）
#   2) GitHub Release（私有仓，需要 GITHUB_TOKEN 环境变量）
#
# 用法（Colab 新会话，挂载 Drive 后）:
#   bash restore_build.sh
#   # 若 HF 仓库名不同: HF_REPO=<user>/Spark-X2.5-4B-build-cpp bash restore_build.sh
#   # 若走 GitHub:      GITHUB_TOKEN=ghp_xxx bash restore_build.sh
#
# 成功标准: $BASE_DIR/llama.cpp-spark/build/bin/llama-server 存在且可执行
# 之后直接跑 oneclick_start.sh 即可起服务（500k 上下文 + 公网隧道）
# ============================================================================
set -euo pipefail

BASE_DIR="${BASE_DIR:-/content/drive/MyDrive/spark-t4}"
FILE="llama.cpp-spark-build-t4-cuda.tar.gz"
HF_REPO="${HF_REPO:-gasschina/Spark-X2.5-4B-build-cpp}"
GH_REPO="${GH_REPO:-samaidev/spark4b}"
GH_TAG="${GH_TAG:-t4-cuda-v1}"
TMP="/content/${FILE}"

GREEN='\033[0;32m'; RED='\033[0;31m'; NC='\033[0m'
ok()   { echo -e "${GREEN}[restore]${NC} $*"; }
fail() { echo -e "${RED}[restore 失败]${NC} $*"; exit 1; }

[[ -d /content/drive/MyDrive ]] || fail "Drive 未挂载，请先在 Colab 挂载 Google Drive"

if [[ -x "$BASE_DIR/llama.cpp-spark/build/bin/llama-server" ]]; then
    ok "检测到编译产物已存在（$BASE_DIR/llama.cpp-spark/build/bin/llama-server），无需恢复"
    exit 0
fi

# ---------- 1) HF 公开仓 ----------
if [[ ! -s "$TMP" ]]; then
    rm -f "$TMP"
    if command -v hf >/dev/null 2>&1; then
        ok "尝试从 HuggingFace 下载: $HF_REPO :: $FILE"
        hf download "$HF_REPO" "$FILE" --local-dir /content/ >/dev/null 2>&1 || \
            echo -e "${RED}[restore]${NC} HF 下载未成功（仓库名不对或网络问题），继续尝试 GitHub"
    else
        echo -e "${RED}[restore]${NC} 未安装 hf CLI，跳过 HF 途径（pip install -U huggingface_hub 后可用）"
    fi
fi

# ---------- 2) GitHub Release（私有，需 token） ----------
if [[ ! -s "$TMP" && -n "${GITHUB_TOKEN:-}" ]]; then
    ok "尝试从 GitHub Release 下载: $GH_REPO@$GH_TAG :: $FILE"
    curl -fL --retry 3 --max-time 900 \
        -H "Authorization: token ${GITHUB_TOKEN}" \
        -o "$TMP" "https://github.com/${GH_REPO}/releases/download/${GH_TAG}/${FILE}" \
        || echo -e "${RED}[restore]${NC} GitHub 下载未成功"
fi

[[ -s "$TMP" ]] || fail "未能获取 $FILE。检查 HF_REPO 名称，或提供 GITHUB_TOKEN。"

# ---------- 3) 校验与解压 ----------
ok "校验并解压到 $BASE_DIR ..."
MD5_EXPECTED="${MD5_EXPECTED:-}"
if [[ -n "$MD5_EXPECTED" ]]; then
    MD5_ACTUAL="$(md5sum "$TMP" | awk '{print $1}')"
    [[ "$MD5_ACTUAL" == "$MD5_EXPECTED" ]] || fail "md5 不匹配: 期望 $MD5_EXPECTED 实际 $MD5_ACTUAL"
    ok "md5 校验通过: $MD5_ACTUAL"
fi
mkdir -p "$BASE_DIR"
tar xzf "$TMP" -C "$BASE_DIR"

[[ -x "$BASE_DIR/llama.cpp-spark/build/bin/llama-server" ]] || fail "解压后未找到 llama-server，压缩包内容异常"
ok "llama-server 就绪: $BASE_DIR/llama.cpp-spark/build/bin/llama-server"
ok "下一步: bash $BASE_DIR/oneclick_start.sh"