#!/bin/bash
# setup.sh - Setup script for DeepCoder deployment

set -e

echo "🚀 DeepCoder Model Setup"
echo "========================"

# Create necessary directories
echo "📁 Creating directories..."
mkdir -p models cache logs ssl

# Set permissions
chmod 755 models cache logs
chmod 700 ssl

# Check for GPU support
echo "🔍 Checking GPU support..."
if command -v nvidia-smi &> /dev/null; then
    echo "✅ NVIDIA GPU detected:"
    nvidia-smi --query-gpu=gpu_name,memory.total --format=csv,noheader
    
    # Check for Docker GPU support
    if docker run --rm --gpus all nvidia/cuda:11.8-base nvidia-smi &> /dev/null; then
        echo "✅ Docker GPU support verified"
        export GPU_SUPPORT=true
    else
        echo "⚠️  Docker GPU support not available"
        export GPU_SUPPORT=false
    fi
else
    echo "⚠️  No GPU detected. Running on CPU."
    export GPU_SUPPORT=false
fi

# Build and start containers
echo "🏗️  Building Docker containers..."
docker-compose build

echo "🚀 Starting services..."
if [ "$GPU_SUPPORT" = true ]; then
    docker-compose up -d
else
    # Remove GPU requirements for CPU-only deployment
    sed 's/devices:/# devices:/g' docker-compose.yml | \
    sed 's/- driver: nvidia/# - driver: nvidia/g' | \
    sed 's/count: 1/# count: 1/g' | \
    sed 's/capabilities: \[gpu\]/# capabilities: [gpu]/g' > docker-compose-cpu.yml
    docker-compose -f docker-compose-cpu.yml up -d
fi

# Wait for services to be ready
echo "⏳ Waiting for services to start..."
sleep 30

# Health check
echo "🏥 Performing health check..."
for i in {1..10}; do
    if curl -f http://localhost:8000/health > /dev/null 2>&1; then
        echo "✅ DeepCoder API is healthy!"
        break
    else
        echo "⏳ Waiting for API to be ready... (attempt $i/10)"
        sleep 10
    fi
done

# Show status
echo "📊 Service Status:"
docker-compose ps

echo ""
echo "🎉 DeepCoder setup complete!"
echo "API endpoint: http://localhost:8000"
echo "Health check: http://localhost:8000/health"
echo "Model info: http://localhost:8000/model/info"
echo ""
echo "To test the API:"
echo "curl -X POST http://localhost:8000/generate \\"
echo "  -H 'Content-Type: application/json' \\"
echo "  -d '{\"prompt\": \"def fibonacci(n):\", \"max_tokens\": 200}'"