#!/bin/bash

set -e

# Function to display usage information
usage() {
  echo "Usage: $0 <model_name>"
  exit 1
}

# Check if a model name is provided as an argument
if [ "$#" -ne 1 ]; then
  usage
fi

MODEL_NAME=$1

# Validate model name to prevent command injection
if ! echo "$MODEL_NAME" | grep -qE '^[a-zA-Z0-9._:/-]+$'; then
    echo "ERROR: Invalid model name: $MODEL_NAME"
    exit 1
fi

# Start the main Ollama application
ollama serve &

# Wait for the Ollama application to be ready (optional, if necessary)
while ! ollama ls; do
  echo "Waiting for Ollama service to be ready..."
  sleep 10
done
echo "Ollama service is ready."

# Pull the model using ollama pull
echo "Pulling the $MODEL_NAME with ollama pull..."
# Check if the model was pulled successfully
if ollama pull "$MODEL_NAME"; then
  echo "Model pulled successfully."
else
  echo "Failed to pull model."
  exit 1
fi

# Run ollama forever.
sleep infinity