#!/bin/bash
set -e

# Check GPU status and compute capability
echo "Checking GPU status..."
nvidia-smi || echo "Warning: nvidia-smi failed. GPU might not be available"
echo "GPU Compute Capability:"
nvidia-smi --query-gpu=compute_cap --format=csv,noheader || echo "Warning: Could not get compute capability"

# Create temporary directory for llamafiler
mkdir -p /tmp/llamafiler

# Start nginx
echo "Starting nginx..."
/usr/sbin/nginx -c /etc/nginx/nginx.conf
if [ $? -ne 0 ]; then
    echo "Failed to start nginx"
    exit 1
fi

# Verify nginx is running
if ! ps aux | grep nginx | grep -v grep > /dev/null; then
    echo "Nginx failed to start"
    exit 1
fi
echo "Nginx started successfully"

# Start the models
echo "Starting models..."
TMPDIR=/tmp/llamacpp ./llama-server -m $HOME/models/gemma-2b.gguf -ngl 999 --host 0.0.0.0 --port 8082 &
GEMMA_PID=$!

TMPDIR=/tmp/llamacpp ./llama-server --embedding -m $HOME/models/embeddings.gguf -ngl 999 --host 0.0.0.0 --port 8081 &
EMBEDDINGS_PID=$!

# Wait for models to be ready
echo "Waiting for models to be ready..."
START_TIME=$SECONDS
TIMEOUT=600  # 10 minutes

wait_for_models() {
    CHAT_HEALTH=$(curl -s http://127.0.0.1:8082/health)
    EMBED_HEALTH=$(curl -s http://127.0.0.1:8081/health)
    
    [[ "$CHAT_HEALTH" == *"\"status\":\"ok\""* ]] && [[ "$EMBED_HEALTH" == *"\"status\":\"ok\""* ]]
}

until wait_for_models; do
    ELAPSED=$((SECONDS - START_TIME))
    if [ $ELAPSED -gt $TIMEOUT ]; then
        echo "Timeout after ${TIMEOUT} seconds"
        exit 1
    fi
    if ! kill -0 $GEMMA_PID 2>/dev/null || ! kill -0 $EMBEDDINGS_PID 2>/dev/null; then
        echo "Model process died"
        exit 1
    fi
    echo "Waiting for models... (${ELAPSED}s elapsed)"
    sleep 2
done

# Start AI Assistant
echo "Models ready after ${ELAPSED}s. Starting AI Assistant..."
cd $HOME/app/aia
PORT=4000 node app/main.bundle.js &
AIA_PID=$!

# Keep container running
wait $GEMMA_PID