setup-openaudio.sh

#!/bin/bash

# OpenAudio Setup Script
#
# Works in two modes:
#   1. Standalone (curl | bash): Downloads compose files from GitHub, pulls
#      pre-built images from DockerHub, downloads models via containers.
#      No git clone, no local Python, no building.
#
#   2. Cloned repo: Detects local source, offers build-from-source or
#      pull-from-DockerHub. Uses local files directly.
#
# Usage:
#   One-liner:  curl -fsSL https://raw.githubusercontent.com/davidamacey/OpenAudio/main/setup-openaudio.sh | bash
#   From repo:  ./setup-openaudio.sh [--yes]

set -e

#######################
# COLOR DEFINITIONS
#######################

RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
CYAN='\033[0;36m'
BOLD='\033[1m'
NC='\033[0m'

log_info() { echo -e "${BLUE}[INFO]${NC} $1"; }
log_success() { echo -e "${GREEN}[OK]${NC} $1"; }
log_warn() { echo -e "${YELLOW}[WARN]${NC} $1"; }
log_error() { echo -e "${RED}[ERROR]${NC} $1"; }
log_step() { echo -e "${CYAN}[STEP]${NC} $1"; }
log_header() { echo -e "\n${BOLD}=== $1 ===${NC}\n"; }

#######################
# CONSTANTS
#######################

GITHUB_RAW_BASE="https://raw.githubusercontent.com/davidamacey/OpenAudio/main"
INSTALL_DIR_DEFAULT="$HOME/openaudio"

# Files to download in standalone mode
STANDALONE_FILES=(
    "docker-compose.yml"
    "docker-compose.prod.yml"
    "docker-compose.gpu.yml"
    ".env.example"
    "openaudio.sh"
    "scripts/common.sh"
    "scripts/download-models.sh"
)

# Infrastructure images to pull
INFRA_IMAGES=(
    "pgvector/pgvector:pg17"
    "redis:8.2.2-alpine3.22"
    "minio/minio:RELEASE.2025-09-07T16-13-09Z"
)

# Application images from DockerHub
APP_IMAGES=(
    "davidamacey/openaudio-backend"
    "davidamacey/openaudio-torch"
    "davidamacey/openaudio-sam"
    "davidamacey/openaudio-enhance"
    "davidamacey/openaudio-frontend"
)

#######################
# PARSE ARGUMENTS
#######################

AUTO_CONFIRM=false
for arg in "$@"; do
    case "$arg" in
        --yes|-y) AUTO_CONFIRM=true ;;
    esac
done

#######################
# INTERACTIVE INPUT
#######################

# When running via `curl | bash`, stdin is the pipe, not the terminal.
# We must read user input from /dev/tty explicitly. If /dev/tty is not
# available (e.g., headless CI), fall back to non-interactive defaults.
INTERACTIVE=false
if [ -t 0 ]; then
    # Normal execution: stdin is a terminal
    INTERACTIVE=true
elif [ -e /dev/tty ]; then
    # Piped execution (curl | bash) but terminal is available
    INTERACTIVE=true
fi

# Read a line of user input, handling both normal and piped execution.
# Falls back to empty string if non-interactive.
prompt_user() {
    local PROMPT="$1"
    local RESULT=""
    echo -n "$PROMPT"
    if [ -t 0 ]; then
        read -r RESULT
    elif [ -e /dev/tty ]; then
        read -r RESULT < /dev/tty
    fi
    echo "$RESULT"
}

#######################
# CONTEXT DETECTION
#######################

detect_context() {
    # Determine if we're running from a cloned repo or standalone (curl | bash)
    if [ -f "docker-compose.yml" ] && [ -f "backend/Dockerfile" ]; then
        IS_REPO=true
        PROJECT_DIR="$(pwd)"
        log_info "Detected cloned repository at $PROJECT_DIR"
    else
        IS_REPO=false
        log_info "Standalone mode (no local repo detected)"
    fi
}

#######################
# PREREQUISITE CHECKS
#######################

check_prerequisites() {
    log_header "Checking Prerequisites"

    local MISSING=()

    # Check Docker
    log_step "Checking Docker..."
    if command -v docker &> /dev/null; then
        DOCKER_VERSION=$(docker --version | cut -d' ' -f3 | tr -d ',')
        log_success "Docker installed: $DOCKER_VERSION"

        if ! docker info > /dev/null 2>&1; then
            log_warn "Docker daemon is not running"
            MISSING+=("Docker daemon not running")
        fi
    else
        log_error "Docker not found"
        MISSING+=("Docker")
    fi

    # Check docker compose
    log_step "Checking Docker Compose..."
    if docker compose version &> /dev/null; then
        COMPOSE_VERSION=$(docker compose version | cut -d' ' -f4)
        log_success "Docker Compose (plugin) installed: $COMPOSE_VERSION"
    elif command -v docker-compose &> /dev/null; then
        COMPOSE_VERSION=$(docker-compose --version | cut -d' ' -f4 | tr -d ',')
        log_success "Docker Compose installed: $COMPOSE_VERSION"
    else
        log_error "Docker Compose not found"
        MISSING+=("Docker Compose")
    fi

    # Check curl (needed for standalone downloads)
    log_step "Checking curl..."
    if command -v curl &> /dev/null; then
        log_success "curl available"
    else
        log_error "curl not found"
        MISSING+=("curl")
    fi

    if [ ${#MISSING[@]} -gt 0 ]; then
        log_error "Missing prerequisites: ${MISSING[*]}"
        echo ""
        echo "Installation guides:"
        echo "  Docker: https://docs.docker.com/get-docker/"
        echo "  curl:   apt install curl / yum install curl"
        exit 1
    fi

    log_success "All prerequisites satisfied"
}

#######################
# GPU DETECTION
#######################

detect_gpu() {
    log_header "Detecting GPU Configuration"

    GPU_AVAILABLE="false"

    if command -v nvidia-smi &> /dev/null && nvidia-smi &> /dev/null; then
        log_success "NVIDIA GPU detected"

        GPU_NAME=$(nvidia-smi --query-gpu=name --format=csv,noheader | head -1)
        GPU_MEMORY=$(nvidia-smi --query-gpu=memory.total --format=csv,noheader | head -1)
        DRIVER_VERSION=$(nvidia-smi --query-gpu=driver_version --format=csv,noheader | head -1)

        log_info "GPU: $GPU_NAME"
        log_info "Memory: $GPU_MEMORY"
        log_info "Driver: $DRIVER_VERSION"

        if command -v nvcc &> /dev/null; then
            CUDA_VERSION=$(nvcc --version | grep release | sed 's/.*release //' | cut -d',' -f1)
            log_info "CUDA: $CUDA_VERSION"
        else
            log_info "CUDA toolkit not installed (not required for Docker)"
        fi

        if docker info 2>/dev/null | grep -q nvidia; then
            log_success "NVIDIA Container Toolkit available"
            GPU_AVAILABLE="true"
        else
            log_warn "NVIDIA Container Toolkit not installed"
            log_info "Install: https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/install-guide.html"
        fi
    else
        log_info "No NVIDIA GPU detected - will use CPU mode"
    fi
}

#######################
# STANDALONE: DOWNLOAD FILES FROM GITHUB
#######################

download_github_file() {
    local FILE_PATH="$1"
    local DEST="$2"
    local URL="${GITHUB_RAW_BASE}/${FILE_PATH}"
    local MAX_RETRIES=3
    local RETRY=0

    while [ $RETRY -lt $MAX_RETRIES ]; do
        if curl -fsSL --connect-timeout 10 --max-time 30 "$URL" -o "$DEST" 2>/dev/null; then
            if [ -s "$DEST" ]; then
                return 0
            fi
            rm -f "$DEST"
        fi

        RETRY=$((RETRY + 1))
        if [ $RETRY -lt $MAX_RETRIES ]; then
            log_warn "Download failed for $FILE_PATH, retrying ($RETRY/$MAX_RETRIES)..."
            sleep 2
        fi
    done

    log_error "Failed to download $FILE_PATH after $MAX_RETRIES attempts"
    return 1
}

download_standalone_files() {
    log_header "Downloading Configuration Files"
    log_info "Fetching files from GitHub..."

    for FILE_PATH in "${STANDALONE_FILES[@]}"; do
        # Create subdirectory if needed
        local DIR
        DIR=$(dirname "$FILE_PATH")
        if [ "$DIR" != "." ]; then
            mkdir -p "$DIR"
        fi

        log_step "Downloading $FILE_PATH..."
        if download_github_file "$FILE_PATH" "$FILE_PATH"; then
            log_success "Downloaded $FILE_PATH"
        else
            log_error "Failed to download $FILE_PATH"
            echo "Please check your internet connection and try again."
            echo "URL: ${GITHUB_RAW_BASE}/${FILE_PATH}"
            exit 1
        fi
    done

    # Make scripts executable
    chmod +x openaudio.sh 2>/dev/null || true
    chmod +x scripts/download-models.sh 2>/dev/null || true
    chmod +x scripts/common.sh 2>/dev/null || true

    # Validate critical files
    log_step "Validating downloaded files..."

    if ! grep -q "services:" docker-compose.yml; then
        log_error "docker-compose.yml appears invalid (no 'services:' section)"
        exit 1
    fi

    if ! grep -q "services:" docker-compose.prod.yml; then
        log_error "docker-compose.prod.yml appears invalid (no 'services:' section)"
        exit 1
    fi

    if ! grep -q "services:\|deploy:" docker-compose.gpu.yml 2>/dev/null; then
        log_warn "docker-compose.gpu.yml may not have downloaded correctly - GPU mode may not work"
    fi

    log_success "All configuration files downloaded and validated"
}

#######################
# INSTALL DIRECTORY (STANDALONE)
#######################

setup_install_dir() {
    log_header "Setting Up Install Directory"

    if [ "$AUTO_CONFIRM" = true ] || [ "$INTERACTIVE" = false ]; then
        PROJECT_DIR="$INSTALL_DIR_DEFAULT"
    else
        USER_DIR=$(prompt_user "Install directory [$INSTALL_DIR_DEFAULT]: ")
        PROJECT_DIR="${USER_DIR:-$INSTALL_DIR_DEFAULT}"
    fi

    mkdir -p "$PROJECT_DIR"
    cd "$PROJECT_DIR"
    log_success "Install directory: $PROJECT_DIR"
}

#######################
# ENVIRONMENT SETUP
#######################

# Generate a random hex password (safe for .env files and sed - no special chars).
generate_password() {
    local BYTES="$1"
    if command -v openssl &> /dev/null; then
        openssl rand -hex "$BYTES"
    elif command -v python3 &> /dev/null; then
        python3 -c "import secrets; print(secrets.token_hex($BYTES))"
    else
        < /dev/urandom tr -dc 'a-f0-9' | head -c "$(( BYTES * 2 ))"
    fi
}

# Portable sed-in-place: writes to a temp file then moves it back.
# Avoids macOS vs GNU sed -i incompatibility and special character issues.
sed_inplace() {
    local PATTERN="$1"
    local FILE="$2"
    local TMPFILE
    TMPFILE=$(mktemp)
    sed "$PATTERN" "$FILE" > "$TMPFILE" && mv "$TMPFILE" "$FILE"
}

# Set a KEY=VALUE pair in the .env file. Creates the key if it doesn't exist.
# Uses exact string replacement (not regex) for the value to avoid special char issues.
set_env_value() {
    local KEY="$1"
    local VALUE="$2"
    local FILE="${3:-.env}"

    if grep -q "^${KEY}=" "$FILE" 2>/dev/null; then
        # Key exists: replace the line using awk (safe for any value content)
        local TMPFILE
        TMPFILE=$(mktemp)
        awk -v key="$KEY" -v val="$VALUE" 'BEGIN{FS=OFS="="} $1==key{$0=key"="val} 1' "$FILE" > "$TMPFILE" && mv "$TMPFILE" "$FILE"
    else
        # Key doesn't exist: append
        echo "${KEY}=${VALUE}" >> "$FILE"
    fi
}

setup_environment() {
    log_header "Setting Up Environment"

    # Check for existing .env
    if [ -f ".env" ]; then
        log_warn ".env file already exists"
        if [ "$AUTO_CONFIRM" = true ] || [ "$INTERACTIVE" = false ]; then
            log_info "Keeping existing .env file (non-interactive mode)"
            return
        else
            CONFIRM=$(prompt_user "Overwrite? (y/N): ")
            if [[ ! "$CONFIRM" =~ ^[Yy]$ ]]; then
                log_info "Keeping existing .env file"
                return
            fi
        fi
    fi

    if [ ! -f ".env.example" ]; then
        log_error ".env.example not found"
        exit 1
    fi

    log_step "Creating .env from .env.example..."
    cp .env.example .env

    # Generate random passwords (hex-only, safe for env files)
    log_step "Generating secure passwords..."

    set_env_value "POSTGRES_PASSWORD" "$(generate_password 16)"
    set_env_value "MINIO_ROOT_PASSWORD" "$(generate_password 16)"
    set_env_value "REDIS_PASSWORD" "$(generate_password 16)"
    set_env_value "JWT_SECRET_KEY" "$(generate_password 32)"
    set_env_value "ENCRYPTION_KEY" "$(generate_password 32)"

    # Set default admin credentials (user can change after first login)
    set_env_value "DEFAULT_ADMIN_EMAIL" "admin@openaudio.local"
    set_env_value "DEFAULT_ADMIN_PASSWORD" "changeme"

    # Set GPU configuration
    if [ "$GPU_AVAILABLE" = "true" ]; then
        set_env_value "USE_GPU" "true"
        log_info "GPU support enabled in .env"
    fi

    # Prompt for HF_TOKEN
    log_step "HuggingFace token configuration..."
    echo ""
    echo "Some models (PyAnnote, SAM Audio) require a HuggingFace token."
    echo "Get one at: https://huggingface.co/settings/tokens"
    echo "Accept model licenses at:"
    echo "  - https://huggingface.co/pyannote/speaker-diarization-community-1"
    echo "  - https://huggingface.co/facebook/sam-audio-base"
    echo ""

    if [ "$AUTO_CONFIRM" = true ] || [ "$INTERACTIVE" = false ]; then
        log_info "Skipping HF_TOKEN prompt (non-interactive mode)"
    else
        USER_HF_TOKEN=$(prompt_user "Enter HuggingFace token (or press Enter to skip): ")
        if [ -n "$USER_HF_TOKEN" ]; then
            set_env_value "HF_TOKEN" "$USER_HF_TOKEN"
            export HF_TOKEN="$USER_HF_TOKEN"
            log_success "HF_TOKEN saved to .env"
        else
            log_info "Skipping HF_TOKEN - gated models will be skipped during download"
        fi
    fi

    log_success "Environment file created"
}

#######################
# DIRECTORY SETUP
#######################

setup_directories() {
    log_header "Creating Directory Structure"

    DIRS=(
        "models/huggingface"
        "models/torch"
        "models/beats"
        "models/modelscope"
        "data/uploads"
        "data/outputs"
        "data/exports"
        "data/temp"
        "logs"
        "backups"
    )

    for DIR in "${DIRS[@]}"; do
        if [ ! -d "$DIR" ]; then
            mkdir -p "$DIR"
            log_info "Created: $DIR"
        else
            log_info "Exists: $DIR"
        fi
    done

    # Set permissions for model cache (for non-root container user UID 1000)
    log_step "Setting directory permissions..."
    chmod -R 775 models data logs backups 2>/dev/null || true
    chown -R 1000:1000 models/ 2>/dev/null || log_warn "Could not set models/ ownership to UID 1000 - you may need: sudo chown -R 1000:1000 models/"

    log_success "Directory structure ready"
}

#######################
# DOCKER IMAGES
#######################

pull_docker_images() {
    log_header "Pulling Docker Images"

    # Pull infrastructure images
    log_step "Pulling infrastructure images..."
    for IMAGE in "${INFRA_IMAGES[@]}"; do
        log_info "Pulling $IMAGE..."
        docker pull "$IMAGE" || log_warn "Failed to pull $IMAGE"
    done

    # Pull application images from DockerHub
    log_step "Pulling application images from DockerHub..."

    # Read IMAGE_TAG from .env if set, default to latest
    local TAG="latest"
    if [ -f ".env" ]; then
        local ENV_TAG
        ENV_TAG=$(grep "^IMAGE_TAG=" .env 2>/dev/null | cut -d'=' -f2)
        if [ -n "$ENV_TAG" ] && [ "$ENV_TAG" != "latest" ]; then
            TAG="$ENV_TAG"
        fi
    fi

    for IMAGE in "${APP_IMAGES[@]}"; do
        log_info "Pulling $IMAGE:$TAG..."
        docker pull "$IMAGE:$TAG" || log_warn "Failed to pull $IMAGE:$TAG"
    done

    log_success "Docker images pulled"
}

build_from_source() {
    log_header "Building Docker Images from Source"

    log_step "Building application Docker images..."
    docker compose build || {
        log_warn "Some images failed to build"
    }

    log_success "Docker images built"
}

#######################
# DATABASE INIT
#######################

init_database() {
    log_header "Initializing Database"

    local COMPOSE_FILES
    if [ "$IS_REPO" = true ] && [ "$BUILD_MODE" = "source" ]; then
        COMPOSE_FILES=""
    else
        COMPOSE_FILES="-f docker-compose.yml -f docker-compose.prod.yml"
    fi

    log_step "Starting database services..."
    # shellcheck disable=SC2086
    docker compose $COMPOSE_FILES up -d postgres redis minio

    log_step "Waiting for PostgreSQL to be ready..."
    sleep 10

    RETRIES=30
    # shellcheck disable=SC2086
    until docker compose $COMPOSE_FILES exec -T postgres pg_isready -U postgres > /dev/null 2>&1 || [ $RETRIES -eq 0 ]; do
        log_info "Waiting for PostgreSQL... ($RETRIES attempts left)"
        sleep 2
        RETRIES=$((RETRIES - 1))
    done

    if [ $RETRIES -eq 0 ]; then
        log_error "PostgreSQL failed to start"
        exit 1
    fi

    log_success "PostgreSQL is ready"
    log_info "MinIO buckets will be created when backend starts"
    log_success "Database services initialized"
}

#######################
# MODEL DOWNLOAD
#######################

download_models() {
    log_header "Downloading ML Models"

    if [ ! -f "./scripts/download-models.sh" ]; then
        log_warn "Model download script not found"
        return 0
    fi

    if [ "$AUTO_CONFIRM" = true ]; then
        log_info "Auto-confirming model download (--yes)"
        ./scripts/download-models.sh
    elif [ "$INTERACTIVE" = false ]; then
        log_info "Skipping model download (non-interactive mode)"
        log_info "Run './scripts/download-models.sh' later to download models"
    else
        CONFIRM=$(prompt_user "Download ML models now? This may take a while. (y/N): ")
        if [[ "$CONFIRM" =~ ^[Yy]$ ]]; then
            ./scripts/download-models.sh
        else
            log_info "Skipping model download"
            log_info "Run './scripts/download-models.sh' later to download models"
        fi
    fi
}

#######################
# FINAL SUMMARY
#######################

print_summary() {
    log_header "Setup Complete!"

    echo -e "${GREEN}OpenAudio has been set up successfully!${NC}"
    echo ""
    echo "Next steps:"
    echo ""

    if [ "$IS_REPO" = true ]; then
        echo "  1. Start the application:"
        echo -e "     ${CYAN}./openaudio.sh start${NC}          (dev mode, builds from source)"
        echo -e "     ${CYAN}./openaudio.sh start prod${NC}     (prod mode, uses DockerHub images)"
    else
        echo "  1. Start the application:"
        echo -e "     ${CYAN}cd $PROJECT_DIR${NC}"
        echo -e "     ${CYAN}./openaudio.sh start prod${NC}"
    fi
    echo ""
    echo "  2. Access the services:"
    echo -e "     Frontend:          ${CYAN}http://localhost:5473${NC}"
    echo -e "     Backend API:       ${CYAN}http://localhost:5474${NC}"
    echo -e "     API Docs:          ${CYAN}http://localhost:5474/docs${NC}"
    echo -e "     Flower Monitor:    ${CYAN}http://localhost:5475${NC}"
    echo -e "     MinIO Console:     ${CYAN}http://localhost:5479${NC}"
    echo ""
    echo "  3. Default login:"
    echo -e "     Email:    ${CYAN}admin@openaudio.local${NC}"
    echo -e "     Password: ${CYAN}changeme${NC}"
    echo "     (change your password after first login)"
    echo ""
    echo "  4. View logs:"
    echo -e "     ${CYAN}./openaudio.sh logs${NC}"
    echo ""
    echo "  5. Download ML models (if skipped):"
    echo -e "     ${CYAN}./scripts/download-models.sh${NC}"
    echo ""

    if [ "$GPU_AVAILABLE" = "true" ]; then
        echo -e "${GREEN}GPU Support:${NC} Enabled"
        echo "  Start with GPU: ./openaudio.sh start gpu"
    else
        echo -e "${YELLOW}GPU Support:${NC} Not available (CPU mode)"
    fi

    echo ""
    echo "For help, run: ./openaudio.sh help"
    echo ""
}

#######################
# MAIN
#######################

main() {
    echo ""
    echo -e "${BOLD}========================================${NC}"
    echo -e "${BOLD}       OpenAudio Setup                 ${NC}"
    echo -e "${BOLD}========================================${NC}"
    echo ""

    # Step 1: Detect context
    detect_context

    # Step 2: Check prerequisites
    check_prerequisites

    # Step 3: Detect GPU
    detect_gpu

    # Step 4: Setup install directory (standalone) or use current dir (repo)
    if [ "$IS_REPO" = false ]; then
        setup_install_dir
        download_standalone_files
    fi

    # Step 5: Setup environment (.env)
    setup_environment

    # Step 6: Create directory structure
    setup_directories

    # Step 7: Get Docker images
    if [ "$IS_REPO" = true ]; then
        # Repo mode: ask build-from-source or pull
        BUILD_MODE="source"
        if [ "$AUTO_CONFIRM" = false ] && [ "$INTERACTIVE" = true ]; then
            echo ""
            echo "Docker image source:"
            echo "  1) Build from source (development, live reload)"
            echo "  2) Pull from DockerHub (faster, pre-built)"
            IMAGE_CHOICE=$(prompt_user "Choose [1]: ")
            if [ "$IMAGE_CHOICE" = "2" ]; then
                BUILD_MODE="pull"
            fi
        fi

        if [ "$BUILD_MODE" = "pull" ]; then
            pull_docker_images
        else
            build_from_source
        fi
    else
        # Standalone: always pull from DockerHub
        BUILD_MODE="pull"
        pull_docker_images
    fi

    # Step 8: Initialize database
    init_database

    # Step 9: Download ML models
    download_models

    # Step 10: Stop database services started during init
    log_step "Stopping setup services..."
    if [ "$IS_REPO" = true ] && [ "$BUILD_MODE" = "source" ]; then
        docker compose down 2>/dev/null || true
    else
        docker compose -f docker-compose.yml -f docker-compose.prod.yml down 2>/dev/null || true
    fi

    # Step 11: Print summary
    print_summary
}

# Run main
main "$@"