#!/bin/bash
set -e

# Support comma-separated list of models via OLLAMA_MODELS
# Falls back to single model via OLLAMA_MODEL for backwards compatibility
MODELS="${OLLAMA_MODELS:-${OLLAMA_MODEL:-llama-guard3:1b}}"

# Start Ollama server in background
ollama serve &
SERVER_PID=$!

# Wait for server to be ready
echo "Waiting for Ollama server to start..."
for i in {1..60}; do
    if curl -sf http://localhost:11434/ > /dev/null 2>&1; then
        echo "Ollama server is ready"
        break
    fi
    sleep 1
done

# Pull each model in the comma-separated list
IFS=',' read -ra MODEL_ARRAY <<< "$MODELS"
for MODEL in "${MODEL_ARRAY[@]}"; do
    # Trim whitespace
    MODEL=$(echo "$MODEL" | xargs)

    if [ -z "$MODEL" ]; then
        continue
    fi

    # Check if model exists, pull if not
    if ! ollama list | grep -q "^${MODEL}"; then
        echo "Pulling model: ${MODEL}"
        ollama pull "${MODEL}"
        echo "Model ${MODEL} pulled successfully"
    else
        echo "Model ${MODEL} already available"
    fi
done

echo "All models ready"

# Keep server running in foreground
wait $SERVER_PID