#!/bin/bash
# Configure Ollama to use /mnt/data for model storage
# Run this ON THE GPU VM

echo "=== Fixing Ollama Storage Location ==="
echo ""

# Check current disk usage
echo "Current disk usage:"
df -h | grep -E "Filesystem|/dev/sda"
echo ""

# Create models directory on /mnt/data
echo "Creating Ollama models directory on /mnt/data..."
sudo mkdir -p /mnt/data/ollama/models
sudo chown -R ollama:ollama /mnt/data/ollama 2>/dev/null || sudo chown -R $(whoami):$(whoami) /mnt/data/ollama
echo "✓ Directory created: /mnt/data/ollama/models"
echo ""

# Check if there are existing models to move
if [ -d ~/.ollama/models ] && [ "$(ls -A ~/.ollama/models 2>/dev/null)" ]; then
    echo "Found existing models in ~/.ollama/models"
    echo "Moving to /mnt/data/ollama/models..."
    sudo mv ~/.ollama/models/* /mnt/data/ollama/models/ 2>/dev/null
    echo "✓ Models moved"
elif [ -d /usr/share/ollama/models ] && [ "$(ls -A /usr/share/ollama/models 2>/dev/null)" ]; then
    echo "Found existing models in /usr/share/ollama/models"
    echo "Moving to /mnt/data/ollama/models..."
    sudo mv /usr/share/ollama/models/* /mnt/data/ollama/models/ 2>/dev/null
    echo "✓ Models moved"
else
    echo "No existing models found to move"
fi

echo ""

# Update systemd service to use new location
echo "Updating systemd service configuration..."
sudo mkdir -p /etc/systemd/system/ollama.service.d

# Check if override.conf exists and update it, or create new
if [ -f /etc/systemd/system/ollama.service.d/override.conf ]; then
    echo "Updating existing override.conf..."
    # Add OLLAMA_MODELS if not already there
    if ! grep -q "OLLAMA_MODELS" /etc/systemd/system/ollama.service.d/override.conf; then
        sudo sed -i '/\[Service\]/a Environment="OLLAMA_MODELS=/mnt/data/ollama/models"' /etc/systemd/system/ollama.service.d/override.conf
    else
        sudo sed -i 's|OLLAMA_MODELS=.*|OLLAMA_MODELS=/mnt/data/ollama/models|' /etc/systemd/system/ollama.service.d/override.conf
    fi
else
    echo "Creating new override.conf..."
    sudo tee /etc/systemd/system/ollama.service.d/override.conf > /dev/null <<EOF
[Service]
Environment="OLLAMA_HOST=0.0.0.0:11434"
Environment="OLLAMA_MODELS=/mnt/data/ollama/models"
EOF
fi

echo "✓ Systemd configuration updated"
echo ""

# Reload and restart
echo "Reloading systemd and restarting Ollama..."
sudo systemctl daemon-reload
sudo systemctl restart ollama
sleep 3

echo "✓ Ollama restarted with new storage location"
echo ""

# Verify
echo "Verifying configuration:"
echo "  Storage location: /mnt/data/ollama/models"
echo "  Disk space available:"
df -h /mnt/data | tail -1
echo ""
echo "  Checking if Ollama is running:"
systemctl is-active ollama && echo "  ✓ Ollama is running" || echo "  ✗ Ollama is not running"

echo ""
echo "=== Done! ==="
echo "You can now pull models:"
echo "  ollama pull qwen2:latest"
echo "  ollama pull qwen2.5:14b"
echo "  ollama pull llama3.1:8b"
echo "  ollama pull qwen2.5:7b"