#!/usr/bin/env bash
# =============================================================================
# DeepSeek-Math-V2 → GGUF Q8_0 conversion pipeline
#
# Uses the HF Storage Bucket mounted at /data as persistent disk.
# All large files (FP8 shards, BF16 shards, GGUF) live on /data — the
# ephemeral overlay disk is never used for model data.
#
# Steps:
#   1. Stream: download FP8 shard → cast BF16 → save to /data/bf16 → delete FP8
#   2. Convert: /data/bf16 → Q8_0 GGUF → /data/output/
#   3. (Optional) upload GGUF to a separate HF repo
#
# Space Variables to set:
#   HF_TOKEN          – HuggingFace token (Secret)
#   HF_DATASET_REPO   – source/dest dataset repo e.g. memmywinks/DeepSeek_Math_V2
#
# Skip flags (set to 1 to skip a step):
#   SKIP_CAST=1       BF16 shards already in /data/bf16
#   SKIP_CONVERT=1    GGUF already exists
#   SKIP_UPLOAD=1     skip final upload (default: 1, upload manually if needed)
# =============================================================================

set -euo pipefail

RED='\033[0;31m'; GREEN='\033[0;32m'; YELLOW='\033[1;33m'; BLUE='\033[0;34m'; NC='\033[0m'
log()  { echo -e "${BLUE}[$(date '+%H:%M:%S')] $*${NC}"; }
ok()   { echo -e "${GREEN}[$(date '+%H:%M:%S')] ✔ $*${NC}"; }
warn() { echo -e "${YELLOW}[$(date '+%H:%M:%S')] ⚠ $*${NC}"; }
err()  { echo -e "${RED}[$(date '+%H:%M:%S')] ✘ $*${NC}" >&2; exit 1; }

# ── Paths — all under /data (the bucket mount) ────────────────────────────────
SRC_REPO="${SRC_REPO:-deepseek-ai/DeepSeek-Math-V2}"
HF_DATASET_REPO="${HF_DATASET_REPO:-}"
LLAMA_CPP_DIR="${LLAMA_CPP_DIR:-/opt/llama.cpp}"

DATA_DIR="${DATA_DIR:-/data}"
FP8_DIR="${FP8_DIR:-$DATA_DIR/fp8}"          # temp FP8 shard (deleted after cast)
BF16_DIR="${BF16_DIR:-$DATA_DIR/bf16}"        # BF16 shards (persisted on bucket)
GGUF_Q8="${GGUF_Q8:-$DATA_DIR/output/deepseek-math-v2-q8_0.gguf}"
SCRATCH_DIR="${SCRATCH_DIR:-$DATA_DIR/scratch}"

SKIP_CAST="${SKIP_CAST:-0}"
SKIP_CONVERT="${SKIP_CONVERT:-0}"
SKIP_UPLOAD="${SKIP_UPLOAD:-1}"

[[ -z "$HF_DATASET_REPO" ]] && err "HF_DATASET_REPO is not set."
[[ -z "${HF_TOKEN:-}" ]]    && warn "HF_TOKEN not set"

mkdir -p "$FP8_DIR" "$BF16_DIR" "$SCRATCH_DIR" "$(dirname "$GGUF_Q8")"

log "========================================================"
log " DeepSeek-Math-V2 → GGUF Q8_0"
log "========================================================"
log "Source   : $SRC_REPO"
log "BF16 dir : $BF16_DIR   (on bucket)"
log "Output   : $GGUF_Q8    (on bucket)"
log "Disk (/data): $(df -h "$DATA_DIR" | tail -1)"
log ""

# ── Step 1: stream FP8 → BF16, save to bucket ────────────────────────────────
if [[ "$SKIP_CAST" == "1" ]]; then
    warn "Step 1 skipped (SKIP_CAST=1)"
else
    log "Step 1/3 – Streaming cast → saving BF16 shards to $BF16_DIR"
    log "Each shard: download (~4 GB FP8) → cast → save (~9 GB BF16) → delete FP8"

    python3 /workspace/scripts/streaming_cast.py \
        --src-repo   "$SRC_REPO" \
        --dst-repo   "$HF_DATASET_REPO" \
        --bf16-dir   "$BF16_DIR" \
        --work-dir   "$SCRATCH_DIR" \
        --resume
    ok "All BF16 shards saved to $BF16_DIR"
fi

# ── Step 2: BF16 dir (on bucket) → Q8_0 GGUF ────────────────────────────────
if [[ "$SKIP_CONVERT" == "1" ]]; then
    warn "Step 2 skipped (SKIP_CONVERT=1)"
else
    log "Step 2/3 – Converting $BF16_DIR → Q8_0 GGUF"
    [[ -d "$BF16_DIR" ]] || err "BF16 dir not found at $BF16_DIR"

    CONVERT_SCRIPT="$LLAMA_CPP_DIR/convert_hf_to_gguf.py"
    [[ -f "$CONVERT_SCRIPT" ]] || err "convert_hf_to_gguf.py not found at $CONVERT_SCRIPT"

    python3 "$CONVERT_SCRIPT" \
        "$BF16_DIR" \
        --outtype q8_0 \
        --outfile "$GGUF_Q8"

    ok "GGUF written → $GGUF_Q8  ($(du -sh "$GGUF_Q8" | cut -f1))"
fi

# ── Step 3: upload GGUF to HF (optional — it's already on the bucket) ─────────
if [[ "$SKIP_UPLOAD" == "1" ]]; then
    warn "Step 3 skipped (SKIP_UPLOAD=1)"
    log "GGUF is at $GGUF_Q8 on your bucket — accessible via HF dataset repo"
else
    log "Step 3/3 – Uploading GGUF to $HF_DATASET_REPO..."
    huggingface-cli upload-large-folder \
        "$HF_DATASET_REPO" \
        "$(dirname "$GGUF_Q8")" \
        --repo-type dataset \
        --include "*.gguf"
    ok "Upload complete → https://huggingface.co/datasets/$HF_DATASET_REPO"
fi

log ""
ok " Pipeline complete!"
echo "GGUF location: $GGUF_Q8"