#!/usr/bin/env bash
# ============================================================
# Velocity-OS — Model Hydration Script
# Uses s5cmd for high-throughput parallel S3 → NVMe sync.
# Run once at initial install; safe to re-run for updates.
#
# Models synced (all stored on NVMe, never in Docker images):
#   - Wan 2.2 (ComfyUI video/image model)
#   - Qwen-Image 2512 (ComfyUI poster/image model)
#   - Qwen3.6 35B A3B (SGLang LLM — MIG slice 0)
#
# Requires: s5cmd, AWS credentials with S3 read access
# ============================================================
set -euo pipefail

# ── Configuration ────────────────────────────────────────────
S3_BUCKET="${S3_MODEL_BUCKET:-s3://desineuron-models}"
LOCAL_BASE="/opt/dlami/nvme/models"
S5CMD_CONCURRENCY="${S5CMD_CONCURRENCY:-32}"  # Tune to NVMe write IOPS

# ── Ensure directories exist ─────────────────────────────────
mkdir -p \
  "${LOCAL_BASE}/comfy/wan2.2" \
  "${LOCAL_BASE}/comfy/qwen-image-2512" \
  "${LOCAL_BASE}/llm/qwen3.6-35b-a3b"

# ── Check s5cmd installed ────────────────────────────────────
if ! command -v s5cmd &> /dev/null; then
  echo "ERROR: s5cmd not found. Install from https://github.com/peak/s5cmd"
  echo "  curl -L https://github.com/peak/s5cmd/releases/latest/download/s5cmd_Linux_x86_64.tar.gz | tar xz -C /usr/local/bin"
  exit 1
fi

echo "=== Velocity-OS Model Hydration ==="
echo "Source: ${S3_BUCKET}"
echo "Target: ${LOCAL_BASE}"
echo "s5cmd workers: ${S5CMD_CONCURRENCY}"
echo ""

# ── Wan 2.2 (ComfyUI — MIG slice 1) ──────────────────────────
echo "[1/3] Syncing Wan 2.2..."
s5cmd \
  --numworkers "${S5CMD_CONCURRENCY}" \
  --credentials-file /etc/velocity/aws-credentials \
  sync \
  "${S3_BUCKET}/wan2.2/*" \
  "${LOCAL_BASE}/comfy/wan2.2/"
echo "  ✓ Wan 2.2 synced."

# ── Qwen-Image 2512 (ComfyUI — MIG slice 1) ──────────────────
echo "[2/3] Syncing Qwen-Image 2512..."
s5cmd \
  --numworkers "${S5CMD_CONCURRENCY}" \
  --credentials-file /etc/velocity/aws-credentials \
  sync \
  "${S3_BUCKET}/qwen-image-2512/*" \
  "${LOCAL_BASE}/comfy/qwen-image-2512/"
echo "  ✓ Qwen-Image 2512 synced."

# ── Qwen3.6 35B A3B (SGLang — MIG slice 0) ───────────────────
echo "[3/3] Syncing Qwen3.6 35B A3B (LLM — ~70GB, be patient)..."
s5cmd \
  --numworkers "${S5CMD_CONCURRENCY}" \
  --credentials-file /etc/velocity/aws-credentials \
  sync \
  "${S3_BUCKET}/qwen3.6-35b-a3b/*" \
  "${LOCAL_BASE}/llm/qwen3.6-35b-a3b/"
echo "  ✓ Qwen3.6 35B synced."

# ── Verify checksums (optional — if .sha256 files exist in S3) ─
echo ""
echo "=== Verifying checksums ==="
for dir in "${LOCAL_BASE}/comfy/wan2.2" "${LOCAL_BASE}/comfy/qwen-image-2512" "${LOCAL_BASE}/llm/qwen3.6-35b-a3b"; do
  if ls "${dir}"/*.sha256 2>/dev/null | head -1 | grep -q sha256; then
    echo "  Checking ${dir}..."
    (cd "${dir}" && sha256sum -c ./*.sha256 --quiet) && echo "  ✓ ${dir} checksums OK"
  fi
done

echo ""
echo "=== Model hydration complete ==="
echo "NVMe usage:"
du -sh "${LOCAL_BASE}"/*/* 2>/dev/null || true