#!/bin/bash

# Function to handle shutdown gracefully
cleanup() {
    echo "Shutting down Ollama server..."
    pkill -f "ollama serve"
    exit 0
}


# Set up signal handlers for graceful shutdown
trap cleanup SIGTERM SIGINT

# Start Ollama server in the background
echo "Starting Ollama server..."
ollama serve &
OLLAMA_PID=$!

# Wait for the server to be ready with timeout
echo "Waiting for Ollama server to start..."
TIMEOUT=60
COUNTER=0

while ! nc -z localhost 7860; do
    if [ $COUNTER -ge $TIMEOUT ]; then
        echo "ERROR: Ollama server failed to start within $TIMEOUT seconds"
        exit 1
    fi
    echo "Waiting for Ollama server to start... ($COUNTER/$TIMEOUT)"
    sleep 1
    ((COUNTER++))
done

echo "Ollama server is ready!"

# Pull the model with error handling
echo "Pulling the model..."
if ! ollama pull gemma3:270m; then
    echo "ERROR: Failed to pull model smollm2:135m"
    exit 1
fi

if ! ollama pull nomic-embed-text; then
    echo "ERROR: Failed to pull model gemma3:1b"
    exit 1
fi

echo "Model pulled successfully!"

# Keep the container running and wait for the ollama process
echo "Container is ready. Ollama server is running on port 7860."
wait $OLLAMA_PID