bootstrap.sh

#!/bin/bash
#
# Easy Bootstrap Script for Single-Node Setup
#
# This script simplifies the process of bootstrapping the system on a single
# server for development or as the initial control node for a new cluster.
# It runs the main Ansible playbook using a local inventory file.

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

# --- Help Menu ---
show_help() {
    echo "Usage: $0 [options]"
    echo ""
    echo "This script runs the main Ansible playbook to bootstrap the system."
    echo ""
    echo "Options:"
    echo "  --role <role>                Specify the role for this node (all, controller, worker). Default: all."
    echo "  --controller-ip <ip>         Required if --role is 'worker'. IP address of the controller node."
    echo "  --tags <tags>                Comma-separated list of Ansible tags to run."
    echo "  --user <user>                Specify the target user for Ansible. Default: pipecatapp."
    echo "  --purge-jobs                 Stop and purge all running Nomad jobs."
    echo "  --clean-git                  Clean the repository of all untracked files (interactive prompt)."
    echo "  --system-cleanup             Perform a full system cleanup (Purge Jobs, Clean System, Clean Git), with interactive prompts."
    echo "  --verbose [level]            Set verbosity level (0-4). Default 0, or 3 if flag is used without value."
    echo "  --debug                      Alias for --verbose 4."
    echo "  --leave-services-running     Do not clean up Nomad and Consul data on startup."
    echo "  --external-model-server      Skip large model downloads and builds, assuming an external server."
    echo "  --continue                   Resume from the last successfully completed playbook."
    echo "  --benchmark                  Run benchmark tests."
    echo "  --deploy-docker              Deploy the pipecat application using Docker (Default)."
    echo "  --run-local                  Deploy the pipecat application using local raw_exec (for debugging)."
    echo "  --home-assistant-debug       Enable debug mode for Home Assistant."
    echo "  --container                  Run the entire infrastructure inside a single large container."
    echo "  --watch <target>             Pause for inspection after the specified target (task/role) completes."
    echo "  -h, --help                   Display this help message and exit."
}

# --- Initialize flags ---
USE_CONTAINER=false
DO_CLEAN_GIT=false
DO_SYSTEM_CLEANUP=false
DO_PURGE_JOBS=false
VERBOSE_LEVEL=0
ROLE=""
CONTROLLER_IP=""

# --- Profile System Resources ---
profile_system() {
    echo -e "\n${BOLD}=== Profiling System Resources ===${NC}"

    local CPU_CORES=$(nproc 2>/dev/null || echo 1)
    local RAM_KB=$(awk '/MemTotal/ {print $2}' /proc/meminfo 2>/dev/null || echo 0)
    local RAM_GB=$(( RAM_KB / 1024 / 1024 ))

    echo -e "Detected CPU Cores: ${CYAN}${CPU_CORES}${NC}"
    echo -e "Detected Total RAM: ${CYAN}${RAM_GB} GB${NC}"

    # Auto-detect role if not explicitly set
    if [ -z "$ROLE" ]; then
        if [ "$RAM_GB" -le 4 ]; then
            echo -e "${YELLOW}⚠️  Low resource machine detected ($RAM_GB GB RAM). Defaulting role to 'worker' and enabling external models.${NC}"
            ROLE="worker"
            PROCESSED_ARGS+=("--role" "worker" "--external-model-server")
        elif [ "$RAM_GB" -ge 8 ] && [ "$CPU_CORES" -ge 4 ]; then
            echo -e "${GREEN}✅ Powerful machine detected. Defaulting role to 'all'.${NC}"
            ROLE="all"
            PROCESSED_ARGS+=("--role" "all")
        else
            echo -e "${CYAN}ℹ️  Standard machine detected. Defaulting role to 'worker'.${NC}"
            ROLE="worker"
            PROCESSED_ARGS+=("--role" "worker")
        fi
    else
        echo -e "Role explicitly set to: ${CYAN}${ROLE}${NC}"
    fi
}

# --- Parse command-line arguments for wrapper logic ---
# We use a while loop to handle optional values for flags like --verbose
ARGS=("$@")
PROCESSED_ARGS=()
SKIP_NEXT=false

for ((i=0; i<${#ARGS[@]}; i++)); do
    arg="${ARGS[$i]}"

    if [ "$SKIP_NEXT" = true ]; then
        SKIP_NEXT=false
        continue
    fi

    case $arg in
        --system-cleanup)
            DO_SYSTEM_CLEANUP=true
            DO_PURGE_JOBS=true
            DO_CLEAN_GIT=true
            ;;
        --clean-git|--clean) # Support legacy --clean just in case, but map to clean-git
            DO_CLEAN_GIT=true
            # Don't pass to provisioning
            ;;
        --purge-jobs)
            DO_PURGE_JOBS=true
            # Don't pass to provisioning as a direct arg, we handle logic
            ;;
        --container)
            USE_CONTAINER=true
            PROCESSED_ARGS+=("$arg")
            ;;
        --debug)
            VERBOSE_LEVEL=4
            PROCESSED_ARGS+=("$arg")
            ;;
        --verbose)
            # Check next arg
            NEXT_ARG="${ARGS[$((i+1))]}"
            if [[ -n "$NEXT_ARG" && ! "$NEXT_ARG" =~ ^- ]]; then
                VERBOSE_LEVEL="$NEXT_ARG"
                PROCESSED_ARGS+=("--verbose" "$VERBOSE_LEVEL")
                SKIP_NEXT=true
            else
                VERBOSE_LEVEL=3
                PROCESSED_ARGS+=("--verbose" "3")
            fi
            ;;
        --role)
            NEXT_ARG="${ARGS[$((i+1))]}"
            if [[ -n "$NEXT_ARG" && ! "$NEXT_ARG" =~ ^- ]]; then
                ROLE="$NEXT_ARG"
                PROCESSED_ARGS+=("--role" "$ROLE")
                SKIP_NEXT=true
            fi
            ;;
        -h|--help)
            show_help
            exit 0
            ;;
        *)
            PROCESSED_ARGS+=("$arg")
            ;;
    esac
done

# Run system profiling before we proceed
profile_system

# --- Move to the script's directory ---
SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd )
cd "$SCRIPT_DIR"

LOG_FILE="bootstrap_debug.log"
> "$LOG_FILE"

# --- Helper: Run Step ---
run_step() {
    local desc="$1"
    local cmd="$2"

    # Levels 3 and 4 show output
    if [ "$VERBOSE_LEVEL" -ge 3 ]; then
        echo -e "\n${BOLD}${CYAN}--- Running ${desc} ---${NC}"
        # Execute and tee to log
        eval "$cmd" 2>&1 | tee -a "$LOG_FILE"
        local status=${PIPESTATUS[0]} # Capture exit code of the evaluated command

        if [ $status -eq 0 ]; then
            echo -e "${GREEN}✅ ${desc} complete.${NC}"
        else
            echo -e "${RED}❌ ${desc} failed.${NC}"
            return $status
        fi
    else
        echo -n -e "⏳ ${desc}..."

        local tmp_log=$(mktemp)
        eval "$cmd" > "$tmp_log" 2>&1
        local status=$?

        cat "$tmp_log" >> "$LOG_FILE"

        if [ $status -eq 0 ]; then
            echo -e "\r\033[K${GREEN}✅ ${desc} Complete${NC}"
        else
            echo -e "\r\033[K${RED}❌ ${desc} Failed${NC}"
            echo -e "${YELLOW}--- Error Log ---${NC}"
            cat "$tmp_log"
            echo -e "${YELLOW}-----------------${NC}"
            rm "$tmp_log"
            return $status
        fi
        rm "$tmp_log"
    fi
}

ask_confirm() {
    local prompt="$1"
    read -p "$prompt [y/N] " -n 1 -r
    echo
    if [[ $REPLY =~ ^[Yy]$ ]]; then
        return 0
    else
        return 1
    fi
}

# --- Environment Setup (Reusable) ---
VENV_DIR="$SCRIPT_DIR/.venv"

ensure_python_environment() {
    echo -e "\n${BOLD}=== Environment Setup ===${NC}"

    setup_venv() {
        if [ ! -d "$VENV_DIR" ]; then
            python3 -m venv "$VENV_DIR"
        fi
    }
    run_step "Creating Python virtual environment" "setup_venv"

    # Activate venv for this script execution
    source "$VENV_DIR/bin/activate"

    run_step "Upgrading pip" "pip install --upgrade pip"

    if [ -f "requirements-dev.txt" ]; then
        run_step "Installing Python dependencies" "pip install --no-cache-dir -r requirements-dev.txt"
    else
        echo "⚠️  Warning: requirements-dev.txt not found. Skipping dependency installation."
    fi

    run_step "Installing Ansible Core" "pip install ansible-core pyyaml"

    # --- Find Ansible Playbook executable ---
    # Since we are in a venv, these are guaranteed to be in the path
    ANSIBLE_GALAXY_EXEC="$(which ansible-galaxy)"

    # Install Ansible collections
    if [ -x "$ANSIBLE_GALAXY_EXEC" ]; then
        run_step "Installing Ansible collections" "$ANSIBLE_GALAXY_EXEC collection install community.general ansible.posix community.docker"
    else
        echo "Error: ansible-galaxy not found in venv." >&2
        exit 1
    fi
}

# --- Cleanup Actions ---

perform_purge_jobs() {
    echo -e "\n${BOLD}${YELLOW}⚠️  Purge Jobs initiated.${NC}"
    if ask_confirm "Are you sure you want to stop and purge all Nomad jobs?"; then
        # Ensure we have the python environment to run the script
        ensure_python_environment

        echo "Running provisioning script to purge jobs..."
        # Pass --purge-jobs and --only-purge
        python3 scripts/provisioning.py --purge-jobs --only-purge
        local status=$?
        if [ $status -eq 0 ]; then
            echo -e "${GREEN}✅ Jobs purged.${NC}"
        else
             echo -e "${RED}❌ Job purge failed.${NC}"
        fi
    else
        echo "Job purge cancelled."
    fi
}

perform_system_cleanup() {
    echo -e "\n${BOLD}${YELLOW}⚠️  System Cleanup initiated (Docker, Apt, Logs).${NC}"
    if ask_confirm "Are you sure you want to aggressively clean system resources?"; then
        if [ -x "scripts/cleanup.sh" ]; then
            # We assume the user has sudo if they are running this
            run_step "Running system cleanup script" "sudo ./scripts/cleanup.sh"
        else
            echo -e "${RED}❌ scripts/cleanup.sh not found or not executable.${NC}"
        fi
    else
        echo "System cleanup cancelled."
    fi
}

perform_git_clean() {
    echo -e "\n${BOLD}${YELLOW}⚠️  Git Clean initiated.${NC}"
    echo "This will permanently delete all untracked files."
    echo "--------------------------------------------------"
    git clean -ndx
    echo "--------------------------------------------------"

    if ask_confirm "Are you sure you want to permanently delete these files?"; then
        if ! run_step "Cleaning repository" "git clean -fdx"; then
            echo -e "\n${YELLOW}⚠️  Standard cleanup failed. This is often due to files created with sudo.${NC}"
            if ask_confirm "Do you want to try cleaning with sudo?"; then
                # Ensure sudo credentials
                if ! sudo -n true 2>/dev/null; then
                    sudo -v
                fi
                run_step "Cleaning repository (with sudo)" "sudo git clean -fdx"
            else
                echo "Cleanup skipped."
            fi
        fi
    else
        echo "Git clean cancelled."
    fi
}

# --- Execute Cleanup Actions ---
# We execute these BEFORE everything else to ensure a clean slate if requested.

if [ "$DO_PURGE_JOBS" = true ]; then
    perform_purge_jobs
fi

if [ "$DO_SYSTEM_CLEANUP" = true ]; then
    perform_system_cleanup
fi

if [ "$DO_CLEAN_GIT" = true ]; then
    perform_git_clean
fi

# If the user only requested cleanup, we might want to stop here?
# But typically bootstrap means "setup". If I wanted to JUST clean, I might not expect it to start building again.
# However, for now we follow the pattern: cleanup then proceed.


# --- Container Mode ---
if [ "$USE_CONTAINER" = true ]; then
    echo -e "${BOLD}--- Running in Container Mode ---${NC}"

    # --- Host-side Cluster Detection ---
    check_for_existing_cluster() {
        echo "--- Checking for existing cluster on host ---"
        CLUSTER_EXISTS=false
        if nc -z localhost 4646 2>/dev/null || nc -z localhost 8500 2>/dev/null; then
            echo "✅ Detected existing Nomad/Consul service on the host."
            CLUSTER_EXISTS=true
        else
            echo "ℹ️  No existing cluster detected on the host. Will start a new one."
        fi
    }

    # Check if we are already inside the container
    if [ -f "/.dockerenv" ] && [ "$(hostname)" = "pipecat-dev-runner" ]; then
        echo "✅ Already inside the container. Proceeding with bootstrap..."
    else
        IMAGE_NAME="pipecat-dev-container"
        CONTAINER_NAME="pipecat-dev-runner"

        check_for_existing_cluster

        echo "Building container image: $IMAGE_NAME..."
        if ! docker build -t "$IMAGE_NAME" docker/dev_container/; then
            echo "❌ Failed to build container image."
            exit 1
        fi

        echo "Checking for existing container..."
        if docker ps -a --format '{{.Names}}' | grep -q "^${CONTAINER_NAME}$"; then
            echo "Removing existing container..."
            docker rm -f "$CONTAINER_NAME"
        fi

        # --- Dynamically configure and start container ---
        HOST_IP=$(hostname -I | awk '{print $1}')
        if [ -z "$HOST_IP" ]; then
            echo "❌ Error: Could not determine host IP address."
            exit 1
        fi

        DOCKER_RUN_CMD=(docker run -d --privileged --name "$CONTAINER_NAME" \
            --hostname "$CONTAINER_NAME" \
            -v /sys/fs/cgroup:/sys/fs/cgroup:rw --cgroupns=host \
            -v "$SCRIPT_DIR":/opt/cluster-infra -e "HOST_IP=$HOST_IP")

        if [ "$CLUSTER_EXISTS" = true ]; then
            echo "Configuring container as a WORKER to join the existing cluster."
        else
            echo "Configuring container as a new CONTROLLER."
            DOCKER_RUN_CMD+=(-p 4646:4646 -p 8500:8500 -p 8081:8081 -p 8000:8000)
        fi

        DOCKER_RUN_CMD+=("$IMAGE_NAME")

        echo "Starting container..."
        if ! "${DOCKER_RUN_CMD[@]}"; then
            echo "❌ Failed to start container."
            exit 1
        fi

        echo "Waiting for container to initialize..."
        sleep 5

        echo "Executing bootstrap inside the container..."
        # Pass processed args
        docker exec -it "$CONTAINER_NAME" /bin/bash -c "cd /opt/cluster-infra && ./bootstrap.sh ${PROCESSED_ARGS[*]}"
        EXIT_CODE=$?

        echo "Container bootstrap finished with exit code: $EXIT_CODE"
        echo "You can access the container using: docker exec -it $CONTAINER_NAME /bin/bash"
        echo "To stop and remove the container: docker rm -f $CONTAINER_NAME"

        exit $EXIT_CODE
    fi
fi

# --- Run Initial Machine Setup ---
echo -e "${BOLD}=== System Bootstrap ===${NC}"
if [ -f "initial-setup/setup.sh" ]; then
    if [ -f "/.dockerenv" ] && [ "$(hostname)" = "pipecat-dev-runner" ]; then
         echo "🐳 Container environment detected. Skipping initial machine setup (setup.sh)."
    else
        # We need to ensure sudo doesn't hang on prompt hidden by redirection
        if sudo -n true 2>/dev/null; then
            # Sudo is already cached
            run_step "Initial machine setup" "sudo bash initial-setup/setup.sh"
        else
            echo "You may be prompted for your sudo password to run the initial setup script."
            sudo -v
            run_step "Initial machine setup" "sudo bash initial-setup/setup.sh"
        fi
    fi
else
    echo "⚠️  Warning: initial-setup/setup.sh not found. Skipping pre-configuration."
fi

# --- Install Python dependencies (Virtual Environment) ---
# Ensure environment is ready (it might have been set up by purge_jobs, or deleted by git clean)
ensure_python_environment

# --- Run Provisioning Script ---
# echo "🚀 Handing over to Python provisioning script..."
# Pass processed arguments to the python script.
python3 scripts/provisioning.py "${PROCESSED_ARGS[@]}"
EXIT_CODE=$?

if [ $EXIT_CODE -eq 0 ]; then
    echo -e "\n${GREEN}✨ Bootstrap complete.${NC}"
else
    echo -e "\n${RED}❌ Bootstrap failed with exit code $EXIT_CODE.${NC}"
fi

exit $EXIT_CODE