#!/bin/sh
# Preseed Ollama with the models the stack should have available at startup.
# Runs once via the model-init service (see docker-compose.yml). Safe to
# re-run — already-present models are skipped.
#
# Two pull paths:
#   - s3_pull   — fetches a tarball from $S3_OLLAMA_BASE (your own mirror,
#                 created by mirror-ollama-model.sh) and extracts into
#                 Ollama's data dir. Faster + immune to upstream changes.
#                 Falls back to ollama pull if S3_OLLAMA_BASE is unset.
#   - pull      — standard `ollama pull` against registry.ollama.ai.

set -e

# Make sure curl is available — ollama/ollama:latest doesn't always include
# it, and s3_pull needs it. tar is in the base image.
if ! command -v curl >/dev/null 2>&1; then
    apt-get update -qq && apt-get install -y -qq curl ca-certificates >/dev/null
fi

S3_OLLAMA_BASE="${S3_OLLAMA_BASE:-}"
OLLAMA_DATA="/root/.ollama"

s3_pull() {
    name="$1"; archive="$2"
    if ollama list 2>/dev/null | awk 'NR>1 {print $1}' | grep -qx "$name"; then
        echo "✓ $name already present"
        return
    fi
    if [ -z "$S3_OLLAMA_BASE" ]; then
        echo "→ $name: S3_OLLAMA_BASE unset, falling back to ollama pull"
        ollama pull "$name"
        return
    fi
    url="${S3_OLLAMA_BASE%/}/$archive"
    echo "→ Downloading $name from $url…"
    curl -fL -C - --retry 3 -o "/tmp/$archive" "$url"
    tar -xzf "/tmp/$archive" -C "$OLLAMA_DATA/models/"
    rm -f "/tmp/$archive"
    echo "✓ $name installed (mirror)"
}

pull() {
    name="$1"
    if ollama list 2>/dev/null | awk 'NR>1 {print $1}' | grep -qx "$name"; then
        echo "✓ $name already present"
    else
        echo "→ Pulling $name from registry.ollama.ai…"
        ollama pull "$name"
    fi
}

# ─── S3-mirrored models ─────────────────────────────────────────────────────
# These live in your own bucket. Create the tarballs once with
# mirror-ollama-model.sh, upload to S3, then list them here.
s3_pull "huihui_ai/qwen3.5-abliterated:9b" "qwen3.5-abliterated-9b.tgz"

# huihui_ai/qwen3-vl-abliterated — Qwen 3 VL base abliteration (different
# fine-tune lineage than Qwen 3.5, so its tool-call template stays intact).
# Used as the Image Studio dispatcher: vision-capable, calls tools
# reliably, and doesn't refuse to dispatch on NSFW edit prompts. Pulled
# from registry; no S3 mirror entry yet.
pull "huihui_ai/qwen3-vl-abliterated:8b"

# ─── Direct registry pulls ──────────────────────────────────────────────────
for model in dolphin3:8b llama3.1:8b ministral-3:8b mistral-nemo:12b qwen3.6:latest; do
    pull "$model"
done

echo "Done."