#!/usr/bin/env bash
# Download the source Qwen3.5-9B checkpoint — the ONLY place this model
# family's MTP head (15 mtp.* tensors) exists. frank-dragoon runs it as the
# vLLM cross-model MTP draft (MiMo-V2.6-Distill-Qwen-9B itself ships zero
# mtp tensors; that is why SGLang's NEXTN on frank-zealot draws 0.00
# acceptance). Public repo — no auth header on the command line.
#
# NOTE the odd upstream shard names (model.safetensors-0000N-of-00004...).
#
# Usage:
#   scripts/download-frank-dragoon-qwen35.sh [DEST_DIR] [--verify-only]
set -euo pipefail

REPO="Qwen/Qwen3.5-9B"
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}"

# name bytes sha256(- = non-LFS, size-only check)
FILES=(
  "model.safetensors-00001-of-00004.safetensors 5276436216 db6f444b43d318c92f360a13a25561a6a65b10c0631b8ed305a426dbaa6c380e"
  "model.safetensors-00002-of-00004.safetensors 5335161512 31c7d7e2dd5d207840b31cc59083c8f4c4718959149e0358c0364052bb9a0330"
  "model.safetensors-00003-of-00004.safetensors 5368717440 7ec36ba3a4176a44c3c0876ad80c56a2f70c84bf008d82e9501df642f17dadec"
  "model.safetensors-00004-of-00004.safetensors 3325995712 b62b0c4cd7e44edee103ee8f4fe225f246d5e768e07bfd5f25b63a8aa1fdd0c6"
  "model.safetensors.index.json 79657 -"
  "config.json 3126 -"
  "chat_template.jinja 7756 -"
  "tokenizer.json 12807982 5f9e4d4901a92b997e463c1f46055088b6cca5ca61a6522d1b9f64c4bb81cb42"
  "tokenizer_config.json 16710 -"
  "vocab.json 6722759 -"
  "merges.txt 3353259 -"
  "preprocessor_config.json 390 -"
  "video_preprocessor_config.json 385 -"
)

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"