#!/usr/bin/env bash
# Download the z-lab Qwen3.5-9B-DFlash drafter for frank-dragoon's DFlash
# speculative-decode arm (vLLM --speculative-config method=dflash). Public
# repo — no auth header on the command line.
#
# The drafter is a 6-layer qwen3 draft whose dflash_config pins
# target_layer_ids [1,5,9,13,17,21,25,29] and block_size 16; its hidden_size
# 4096 / vocab_size 248320 match the MiMo-V2.6-Distill-Qwen-9B target exactly
# (both preserve the Qwen3.5-9B geometry), which is what lets vLLM load it as
# the draft for that target.
#
# Usage:
#   scripts/download-frank-dragoon-dflash.sh [DEST_DIR] [--verify-only]
set -euo pipefail

REPO="z-lab/Qwen3.5-9B-DFlash"
BASE="https://huggingface.co/${REPO}/resolve/main"

DEST=""
VERIFY_ONLY=0
for arg in "$@"; do
    case "$arg" in
        --verify-only) VERIFY_ONLY=1 ;;
        *) DEST="$arg" ;;
    esac
done
DEST="${DEST:-/home/dank/models/Qwen3.5-9B-DFlash}"

# name bytes sha256(- = non-LFS, size-only check)
FILES=(
  "model.safetensors 2583816465 0a42274b32554f48de1faa0d42824e9c2ceda649c30ae0a731cddf410dd698c7"
  "config.json 1222 -"
)

mkdir -p "$DEST"
status=0
for entry in "${FILES[@]}"; do
  read -r name size sha <<<"$entry"
  f="$DEST/$name"
  have=0
  if [[ -f "$f" ]] && [[ "$(stat -c%s "$f")" == "$size" ]]; then
    if [[ "$sha" == "-" ]] || sha256sum -c --quiet <(printf '%s  %s\n' "$sha" "$f"); then
      have=1
    fi
  fi
  if (( have )); then
    echo "OK        $name"
    continue
  fi
  if (( VERIFY_ONLY )); then
    echo "MISSING   $name"
    status=1
    continue
  fi
  echo "FETCHING  $name"
  tmp="$f.part"
  curl -fsSL --fail-with-body -o "$tmp" "$BASE/$name"
  mv "$tmp" "$f"
  if [[ "$(stat -c%s "$f")" != "$size" ]]; then
    echo "SIZE FAIL $name ($(stat -c%s "$f") != $size)"
    status=1
    continue
  fi
  if [[ "$sha" != "-" ]] && ! sha256sum -c --quiet <(printf '%s  %s\n' "$sha" "$f"); then
    echo "SHA FAIL  $name"
    status=1
    continue
  fi
  echo "VERIFIED  $name"
done
if (( status == 0 )); then
  echo "DOWNLOAD_DONE $DEST"
else
  echo "DOWNLOAD_FAILED $DEST" >&2
fi
exit "$status"