#!/bin/bash
# Startup script for AudioCraft Studio
# Used in Docker container and RunPod

set -e

echo "=========================================="
echo "  AudioCraft Studio"
echo "=========================================="

# Create directories if they don't exist
mkdir -p "${AUDIOCRAFT_OUTPUT_DIR:-/workspace/outputs}"
mkdir -p "${AUDIOCRAFT_DATA_DIR:-/workspace/data}"
mkdir -p "${AUDIOCRAFT_MODEL_CACHE:-/workspace/models}"

# Check GPU availability
echo "Checking GPU..."
if command -v nvidia-smi &> /dev/null; then
    nvidia-smi --query-gpu=name,memory.total,memory.free --format=csv
else
    echo "Warning: nvidia-smi not found"
fi

# Check Python and dependencies
echo "Python version:"
python --version

echo "PyTorch version:"
python -c "import torch; print(f'PyTorch: {torch.__version__}, CUDA: {torch.cuda.is_available()}')"

# Check AudioCraft installation
echo "AudioCraft version:"
python -c "import audiocraft; print(audiocraft.__version__)" 2>/dev/null || echo "AudioCraft installed from source"

# Generate API key if not exists
if [ ! -f "${AUDIOCRAFT_DATA_DIR:-/workspace/data}/.api_key" ]; then
    echo "Generating API key..."
    python -c "
from src.api.auth import get_key_manager
km = get_key_manager()
if not km.has_key():
    key = km.generate_new_key()
    print(f'Generated API key: {key}')
    print('Store this key securely - it will not be shown again!')
"
fi

# Start the application
echo "Starting AudioCraft Studio..."
echo "Gradio UI: http://0.0.0.0:${AUDIOCRAFT_GRADIO_PORT:-7860}"
echo "REST API:  http://0.0.0.0:${AUDIOCRAFT_API_PORT:-8000}"
echo "API Docs:  http://0.0.0.0:${AUDIOCRAFT_API_PORT:-8000}/api/docs"
echo "=========================================="

exec python main.py "$@"