#!/usr/bin/env bash
# pull-models.sh — Pull Ollama models into the HomeAI stack
#
# Usage:
#   ./scripts/pull-models.sh                    # pulls llama3.2:3b (default)
#   ./scripts/pull-models.sh mistral:7b         # pulls a specific model
#   ./scripts/pull-models.sh llama3.2:3b mistral:7b  # pulls multiple models

set -euo pipefail

CONTAINER="homeai-ollama"
DEFAULT_MODEL="llama3.2:3b"

# Use provided models or fall back to default
MODELS=("${@:-$DEFAULT_MODEL}")

# Check that the Ollama container is running
if ! docker inspect "$CONTAINER" --format '{{.State.Running}}' 2>/dev/null | grep -q true; then
    echo "Error: Container '$CONTAINER' is not running."
    echo "Start the stack first: docker compose up -d"
    exit 1
fi

# Wait for Ollama to be healthy
echo "Waiting for Ollama to be ready..."
for i in $(seq 1 30); do
    if docker exec "$CONTAINER" bash -c 'echo > /dev/tcp/localhost/11434' 2>/dev/null; then
        break
    fi
    if [ "$i" -eq 30 ]; then
        echo "Error: Ollama did not become healthy after 30 seconds."
        exit 1
    fi
    sleep 1
done
echo "Ollama is ready."

# Pull each model
for MODEL in "${MODELS[@]}"; do
    echo ""
    echo "Pulling $MODEL ..."
    echo "(This may take 10-40 minutes depending on model size and your connection)"
    echo ""
    docker exec "$CONTAINER" ollama pull "$MODEL"
    echo ""
    echo "$MODEL pulled successfully."
done

# List installed models
echo ""
echo "=== Installed Models ==="
docker exec "$CONTAINER" ollama list
echo ""
echo "Done. Open WebUI at http://localhost:${WEBUI_PORT:-3000} should now show your models."