#!/usr/bin/env bash
#
# ICC AI/HPC Stack Installer v2.0
# Copyright (c) 2026 International Computer Concepts (ICC)
# https://icc-usa.com
#
# A curated, production-grade AI/ML software stack for Ubuntu 24.04 LTS
# with NVIDIA GPU acceleration.
#
# Usage:
#   wget -nv -O- https://your-server/install-icc-ai-stack.sh | bash
#   -- or --
#   chmod +x install-icc-ai-stack.sh && sudo ./install-icc-ai-stack.sh [OPTIONS]
#
# Core options:
#   --driver-only         Install NVIDIA driver + CUDA only (no frameworks)
#   --no-jupyter          Skip Jupyter installation
#   --no-frameworks       Skip PyTorch/TensorFlow/JAX
#   --cuda-version VER    Pin CUDA version (e.g., 12.6)
#   --driver-version V    Pin driver branch (e.g., 560)
#
# Optional modules (off by default):
#   --with-docker         Docker CE + NVIDIA Container Toolkit
#   --with-profiling      NVIDIA DCGM + Nsight Systems/Compute + compute-sanitizer
#   --with-inference      vLLM, Triton client libs, llama.cpp (CUDA build)
#   --with-cluster        Open MPI (CUDA-aware), NCCL tests, bandwidth test
#   --with-all            Enable all optional modules
#
# General:
#   --dry-run             Show what would be installed without doing it
#   --uninstall           Remove the ICC AI stack
#   --yes                 Skip confirmation prompts
#   --help                Show this help
#
set -euo pipefail
IFS=$'\n\t'

# ─────────────────────────────────────────────────────────────────────────────
# Configuration — edit these to pin versions or change defaults
# ─────────────────────────────────────────────────────────────────────────────
readonly STACK_VERSION="2.0.0"
readonly STACK_NAME="ICC AI/HPC Stack"
readonly SUPPORTED_OS="Ubuntu 24.04"
readonly LOG_FILE="/var/log/icc-ai-stack-install.log"
readonly STATE_DIR="/etc/icc-ai-stack"
readonly STATE_FILE="${STATE_DIR}/installed-components"

# Default versions (set to "latest" to use repo defaults)
DEFAULT_CUDA_VERSION="12.8"
DEFAULT_DRIVER_BRANCH="570"
DEFAULT_PYTHON_VERSION="3.12"

# Frameworks — set to "latest" or pin (e.g., "2.4.0")
DEFAULT_PYTORCH_VERSION="latest"
DEFAULT_TENSORFLOW_VERSION="latest"
DEFAULT_JAX_VERSION="latest"

# ─────────────────────────────────────────────────────────────────────────────
# Runtime flags
# ─────────────────────────────────────────────────────────────────────────────
OPT_DRIVER_ONLY=false
OPT_NO_JUPYTER=false
OPT_NO_FRAMEWORKS=false
OPT_DRY_RUN=false
OPT_UNINSTALL=false
OPT_YES=false
OPT_CUDA_VERSION="${DEFAULT_CUDA_VERSION}"
OPT_DRIVER_VERSION="${DEFAULT_DRIVER_BRANCH}"

# Optional modules (off by default)
OPT_DOCKER=false
OPT_PROFILING=false
OPT_INFERENCE=false
OPT_CLUSTER=false

# ─────────────────────────────────────────────────────────────────────────────
# Colors & output helpers
# ─────────────────────────────────────────────────────────────────────────────
if [[ -t 1 ]]; then
    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'
else
    RED='' GREEN='' YELLOW='' BLUE='' CYAN='' BOLD='' NC=''
fi

banner() {
    echo -e "${CYAN}"
    cat << 'EOF'
    ╔══════════════════════════════════════════════════════════════╗
    ║                                                              ║
    ║     ██╗ ██████╗ ██████╗     █████╗ ██╗    ███████╗████████╗  ║
    ║     ██║██╔════╝██╔════╝    ██╔══██╗██║    ██╔════╝╚══██╔══╝  ║
    ║     ██║██║     ██║         ███████║██║    ███████╗   ██║     ║
    ║     ██║██║     ██║         ██╔══██║██║    ╚════██║   ██║     ║
    ║     ██║╚██████╗╚██████╗    ██║  ██║██║    ███████║   ██║     ║
    ║     ╚═╝ ╚═════╝ ╚═════╝    ╚═╝  ╚═╝╚═╝    ╚══════╝   ╚═╝     ║
    ║                                                              ║
    ║           AI / HPC Stack for Ubuntu 24.04 LTS                ║
    ║           International Computer Concepts                    ║
    ╚══════════════════════════════════════════════════════════════╝
EOF
    echo -e "${NC}"
    echo -e "  ${BOLD}Version:${NC} ${STACK_VERSION}"
    echo ""
}

log()   { echo -e "${GREEN}[ICC]${NC} $*"; echo "[$(date '+%Y-%m-%d %H:%M:%S')] $*" >> "${LOG_FILE}" 2>/dev/null || true; }
warn()  { echo -e "${YELLOW}[WARN]${NC} $*"; echo "[$(date '+%Y-%m-%d %H:%M:%S')] WARN: $*" >> "${LOG_FILE}" 2>/dev/null || true; }
err()   { echo -e "${RED}[ERROR]${NC} $*" >&2; echo "[$(date '+%Y-%m-%d %H:%M:%S')] ERROR: $*" >> "${LOG_FILE}" 2>/dev/null || true; }
info()  { echo -e "${BLUE}[INFO]${NC} $*"; }
dry()   { echo -e "${YELLOW}[DRY-RUN]${NC} Would run: $*"; }

die() { err "$*"; exit 1; }

# Count total phases dynamically
count_phases() {
    local n=6  # base: prereqs, driver, cuda, frameworks, python-ecosystem, jupyter
    [[ "${OPT_DOCKER}" == true ]]    && ((n++))
    [[ "${OPT_PROFILING}" == true ]] && ((n++))
    [[ "${OPT_INFERENCE}" == true ]] && ((n++))
    [[ "${OPT_CLUSTER}" == true ]]   && ((n++))
    echo "${n}"
}

CURRENT_PHASE=0
TOTAL_PHASES=6

next_phase() {
    ((CURRENT_PHASE++))
    log "Phase ${CURRENT_PHASE}/${TOTAL_PHASES}: $1"
}

# ─────────────────────────────────────────────────────────────────────────────
# Argument parsing
# ─────────────────────────────────────────────────────────────────────────────
parse_args() {
    while [[ $# -gt 0 ]]; do
        case "$1" in
            --driver-only)     OPT_DRIVER_ONLY=true ;;
            --no-jupyter)      OPT_NO_JUPYTER=true ;;
            --no-frameworks)   OPT_NO_FRAMEWORKS=true ;;
            --dry-run)         OPT_DRY_RUN=true ;;
            --uninstall)       OPT_UNINSTALL=true ;;
            --yes|-y)          OPT_YES=true ;;
            --with-docker)     OPT_DOCKER=true ;;
            --with-profiling)  OPT_PROFILING=true ;;
            --with-inference)  OPT_INFERENCE=true ;;
            --with-cluster)    OPT_CLUSTER=true ;;
            --with-all)
                OPT_DOCKER=true
                OPT_PROFILING=true
                OPT_INFERENCE=true
                OPT_CLUSTER=true
                ;;
            --cuda-version)
                shift; OPT_CUDA_VERSION="${1:-}"
                [[ -z "${OPT_CUDA_VERSION}" ]] && die "--cuda-version requires a value"
                ;;
            --driver-version)
                shift; OPT_DRIVER_VERSION="${1:-}"
                [[ -z "${OPT_DRIVER_VERSION}" ]] && die "--driver-version requires a value"
                ;;
            --help|-h)
                banner
                sed -n '2,/^set /p' "$0" | grep '^#' | sed 's/^# \?//'
                exit 0
                ;;
            *) die "Unknown option: $1 (use --help)" ;;
        esac
        shift
    done
}

# ─────────────────────────────────────────────────────────────────────────────
# Preflight checks
# ─────────────────────────────────────────────────────────────────────────────
check_root() {
    if [[ $EUID -ne 0 ]]; then
        die "This script must be run as root. Use: sudo $0"
    fi
}

check_os() {
    if [[ ! -f /etc/os-release ]]; then
        die "Cannot detect OS. /etc/os-release not found."
    fi
    source /etc/os-release
    if [[ "${ID}" != "ubuntu" ]] || [[ ! "${VERSION_ID}" =~ ^24\.04 ]]; then
        die "Unsupported OS: ${PRETTY_NAME}. This script requires ${SUPPORTED_OS}."
    fi
    log "Detected: ${PRETTY_NAME}"
}

check_gpu() {
    if ! lspci 2>/dev/null | grep -qi 'nvidia'; then
        warn "No NVIDIA GPU detected via lspci."
        warn "Installation will proceed, but GPU features won't work without hardware."
        if [[ "${OPT_YES}" != true ]]; then
            read -rp "Continue anyway? [y/N] " ans
            [[ "${ans}" =~ ^[Yy] ]] || exit 0
        fi
    else
        local gpu_info
        gpu_info=$(lspci | grep -i nvidia | head -5)
        log "Detected NVIDIA GPU(s):"
        echo "${gpu_info}" | while read -r line; do info "  ${line}"; done
    fi
}

check_disk_space() {
    local required=15
    [[ "${OPT_DOCKER}" == true ]]    && ((required += 5))
    [[ "${OPT_INFERENCE}" == true ]] && ((required += 8))
    [[ "${OPT_CLUSTER}" == true ]]   && ((required += 3))
    [[ "${OPT_PROFILING}" == true ]] && ((required += 2))

    local available_gb
    available_gb=$(df -BG / | awk 'NR==2 {print $4}' | tr -d 'G')
    if [[ "${available_gb}" -lt "${required}" ]]; then
        die "Insufficient disk space: ${available_gb}GB available, ${required}GB+ recommended."
    fi
    log "Disk space: ${available_gb}GB available (${required}GB required)"
}

check_internet() {
    if ! curl -sf --connect-timeout 5 https://developer.download.nvidia.com > /dev/null 2>&1; then
        die "Cannot reach NVIDIA servers. Check your internet connection."
    fi
    log "Internet connectivity: OK"
}

check_existing_install() {
    local warnings=()
    if dpkg -l 2>/dev/null | grep -q 'nvidia-driver'; then
        warnings+=("Existing NVIDIA driver detected")
    fi
    if command -v conda &>/dev/null; then
        warnings+=("Conda installation detected — may conflict with system Python packages")
    fi
    if command -v docker &>/dev/null && [[ "${OPT_DOCKER}" == true ]]; then
        warnings+=("Docker already installed — will be upgraded/reconfigured")
    fi
    if [[ -f "${STATE_FILE}" ]]; then
        warnings+=("Previous ICC AI Stack installation detected")
    fi
    if [[ ${#warnings[@]} -gt 0 ]]; then
        warn "Pre-existing installations found:"
        for w in "${warnings[@]}"; do warn "  • ${w}"; done
        if [[ "${OPT_YES}" != true ]]; then
            read -rp "Continue? Existing components will be upgraded. [y/N] " ans
            [[ "${ans}" =~ ^[Yy] ]] || exit 0
        fi
    fi
}

# ─────────────────────────────────────────────────────────────────────────────
# Installation helpers
# ─────────────────────────────────────────────────────────────────────────────
run_or_dry() {
    if [[ "${OPT_DRY_RUN}" == true ]]; then
        dry "$*"
    else
        "$@"
    fi
}

apt_install() {
    run_or_dry apt-get install -y --no-install-recommends "$@"
}

record_component() {
    if [[ "${OPT_DRY_RUN}" != true ]]; then
        mkdir -p "${STATE_DIR}"
        echo "$1=$(date -Iseconds)" >> "${STATE_FILE}"
    fi
}

pip_install() {
    run_or_dry pip3 install --no-cache-dir "$@"
}

# ─────────────────────────────────────────────────────────────────────────────
# Phase: System prerequisites
# ─────────────────────────────────────────────────────────────────────────────
install_prerequisites() {
    next_phase "Installing system prerequisites..."

    run_or_dry apt-get update -qq

    apt_install \
        build-essential \
        gcc g++ make cmake \
        dkms \
        linux-headers-"$(uname -r)" \
        curl wget gnupg2 \
        ca-certificates \
        apt-transport-https \
        software-properties-common \
        pkg-config \
        libssl-dev libffi-dev \
        zlib1g-dev libbz2-dev libreadline-dev libsqlite3-dev \
        liblzma-dev libncurses5-dev libncursesw5-dev \
        libxml2-dev libxslt1-dev \
        git unzip \
        htop tmux \
        numactl \
        lsof \
        pciutils \
        "python${DEFAULT_PYTHON_VERSION}" \
        "python${DEFAULT_PYTHON_VERSION}-dev" \
        "python${DEFAULT_PYTHON_VERSION}-venv" \
        python3-pip \
        python3-setuptools \
        python3-wheel

    # Ensure python3 points to the right version
    run_or_dry update-alternatives --install /usr/bin/python3 python3 \
        "/usr/bin/python${DEFAULT_PYTHON_VERSION}" 1 2>/dev/null || true
    run_or_dry update-alternatives --install /usr/bin/python python \
        "/usr/bin/python${DEFAULT_PYTHON_VERSION}" 1 2>/dev/null || true

    # Install uv (fast pip/venv replacement)
    if [[ "${OPT_DRY_RUN}" != true ]]; then
        curl -LsSf https://astral.sh/uv/install.sh | sh 2>/dev/null || warn "uv install failed"
    else
        dry "curl ... | sh  (install uv)"
    fi

    record_component "prerequisites"
    log "Phase ${CURRENT_PHASE} complete."
}

# ─────────────────────────────────────────────────────────────────────────────
# Phase: NVIDIA driver
# ─────────────────────────────────────────────────────────────────────────────
install_nvidia_driver() {
    next_phase "Installing NVIDIA driver (branch ${OPT_DRIVER_VERSION})..."

    local keyring_pkg="cuda-keyring_1.1-1_all.deb"
    local keyring_url="https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2404/x86_64/${keyring_pkg}"

    if [[ "${OPT_DRY_RUN}" != true ]]; then
        wget -q "${keyring_url}" -O "/tmp/${keyring_pkg}"
        dpkg -i "/tmp/${keyring_pkg}"
        rm -f "/tmp/${keyring_pkg}"
        apt-get update -qq
    else
        dry "wget + dpkg -i ${keyring_url}"
        dry "apt-get update"
    fi

    local driver_pkg="nvidia-driver-${OPT_DRIVER_VERSION}"
    if ! apt-cache show "${driver_pkg}" &>/dev/null && [[ "${OPT_DRY_RUN}" != true ]]; then
        warn "Driver package ${driver_pkg} not found. Falling back to nvidia-driver-570."
        driver_pkg="nvidia-driver-570"
    fi

    apt_install "${driver_pkg}"

    # Fabric Manager for multi-GPU NVLink/NVSwitch systems
    if lspci | grep -qi 'NVSwitch\|nvlink'; then
        log "NVSwitch/NVLink detected — installing Fabric Manager..."
        apt_install nvidia-fabricmanager-"${OPT_DRIVER_VERSION}" || true
    fi

    # Persistence daemon
    apt_install nvidia-persistenced || true

    record_component "nvidia-driver:${OPT_DRIVER_VERSION}"
    log "Phase ${CURRENT_PHASE} complete."
}

# ─────────────────────────────────────────────────────────────────────────────
# Phase: CUDA Toolkit + libraries
# ─────────────────────────────────────────────────────────────────────────────
install_cuda() {
    next_phase "Installing CUDA ${OPT_CUDA_VERSION} toolkit and libraries..."

    local cuda_pkg_ver="${OPT_CUDA_VERSION//./-}"
    local cuda_major="${OPT_CUDA_VERSION%%.*}"

    apt_install \
        "cuda-toolkit-${cuda_pkg_ver}" \
        "cuda-tools-${cuda_pkg_ver}" \
        "cuda-compiler-${cuda_pkg_ver}"

    # cuDNN
    log "Installing cuDNN..."
    apt_install \
        "libcudnn9-cuda-${cuda_major}" \
        "libcudnn9-dev-cuda-${cuda_major}" || \
    apt_install libcudnn9 libcudnn9-dev || \
        warn "cuDNN installation failed — may need manual install"

    # NCCL
    log "Installing NCCL..."
    apt_install "libnccl2" "libnccl-dev" || warn "NCCL installation had issues"

    # TensorRT
    log "Installing TensorRT..."
    apt_install tensorrt libnvinfer-dev libnvonnxparsers-dev || \
        warn "TensorRT installation had issues — may need manual install"

    # cuSPARSELt
    apt_install libcusparselt0 libcusparselt-dev 2>/dev/null || true

    # CUDA samples
    apt_install "cuda-samples-${cuda_pkg_ver}" 2>/dev/null || true

    # Environment setup
    if [[ "${OPT_DRY_RUN}" != true ]]; then
        cat > /etc/profile.d/icc-cuda.sh << 'CUDA_ENV'
# ICC AI Stack — CUDA environment
export CUDA_HOME=/usr/local/cuda
export PATH="${CUDA_HOME}/bin${PATH:+:${PATH}}"
export LD_LIBRARY_PATH="${CUDA_HOME}/lib64${LD_LIBRARY_PATH:+:${LD_LIBRARY_PATH}}"
CUDA_ENV
        chmod 644 /etc/profile.d/icc-cuda.sh
        echo "/usr/local/cuda/lib64" > /etc/ld.so.conf.d/icc-cuda.conf
        ldconfig
    fi

    record_component "cuda:${OPT_CUDA_VERSION}"
    record_component "cudnn"
    record_component "nccl"
    record_component "tensorrt"
    log "Phase ${CURRENT_PHASE} complete."
}

# ─────────────────────────────────────────────────────────────────────────────
# Phase: Deep learning frameworks
# ─────────────────────────────────────────────────────────────────────────────
install_frameworks() {
    if [[ "${OPT_DRIVER_ONLY}" == true ]] || [[ "${OPT_NO_FRAMEWORKS}" == true ]]; then
        next_phase "Skipping frameworks (--driver-only or --no-frameworks)."
        return
    fi

    next_phase "Installing deep learning frameworks..."

    pip_install --upgrade pip setuptools wheel

    # PyTorch
    log "Installing PyTorch..."
    local cu_tag="cu${OPT_CUDA_VERSION//./}"
    if [[ "${DEFAULT_PYTORCH_VERSION}" == "latest" ]]; then
        pip_install torch torchvision torchaudio \
            --index-url "https://download.pytorch.org/whl/${cu_tag}"
    else
        pip_install "torch==${DEFAULT_PYTORCH_VERSION}" torchvision torchaudio \
            --index-url "https://download.pytorch.org/whl/${cu_tag}"
    fi
    record_component "pytorch"

    # TensorFlow
    log "Installing TensorFlow..."
    if [[ "${DEFAULT_TENSORFLOW_VERSION}" == "latest" ]]; then
        pip_install 'tensorflow[and-cuda]'
    else
        pip_install "tensorflow[and-cuda]==${DEFAULT_TENSORFLOW_VERSION}"
    fi
    record_component "tensorflow"

    # JAX
    log "Installing JAX with CUDA support..."
    if [[ "${DEFAULT_JAX_VERSION}" == "latest" ]]; then
        pip_install "jax[cuda12]"
    else
        pip_install "jax[cuda12]==${DEFAULT_JAX_VERSION}"
    fi
    record_component "jax"

    log "Phase ${CURRENT_PHASE} complete."
}

# ─────────────────────────────────────────────────────────────────────────────
# Phase: Scientific Python & ML ecosystem
# ─────────────────────────────────────────────────────────────────────────────
install_python_ecosystem() {
    if [[ "${OPT_DRIVER_ONLY}" == true ]]; then
        next_phase "Skipping Python ecosystem (--driver-only)."
        return
    fi

    next_phase "Installing Python ML/scientific ecosystem..."

    # Core scientific stack
    pip_install \
        numpy scipy pandas scikit-learn \
        matplotlib seaborn Pillow sympy h5py

    # ML / LLM utilities
    pip_install \
        transformers datasets accelerate \
        safetensors tokenizers sentencepiece protobuf \
        peft bitsandbytes \
        "huggingface-hub[cli]"

    # Computer vision
    pip_install opencv-python-headless

    # Distributed / kernel compilation
    pip_install deepspeed || warn "DeepSpeed install failed"
    pip_install flash-attn --no-build-isolation 2>/dev/null || warn "flash-attn build failed (may need manual install)"
    pip_install triton || warn "OpenAI Triton install failed"

    # Monitoring
    pip_install tensorboard wandb nvitop gpustat

    # ONNX
    pip_install onnx onnxruntime-gpu

    record_component "python-ecosystem"
    log "Phase ${CURRENT_PHASE} complete."
}

# ─────────────────────────────────────────────────────────────────────────────
# Phase: Jupyter
# ─────────────────────────────────────────────────────────────────────────────
install_jupyter() {
    if [[ "${OPT_DRIVER_ONLY}" == true ]] || [[ "${OPT_NO_JUPYTER}" == true ]]; then
        next_phase "Skipping Jupyter."
        return
    fi

    next_phase "Installing JupyterLab..."

    pip_install jupyterlab notebook ipywidgets jupyterlab-git

    if [[ "${OPT_DRY_RUN}" != true ]]; then
        cat > /etc/systemd/system/jupyterlab.service << 'JUPYTER_SVC'
[Unit]
Description=JupyterLab Server (ICC AI Stack)
After=network.target

[Service]
Type=simple
User=root
ExecStart=/usr/bin/jupyter lab \
    --ip=0.0.0.0 \
    --port=8888 \
    --no-browser \
    --allow-root \
    --NotebookApp.token='' \
    --NotebookApp.password=''
Restart=on-failure
RestartSec=10

[Install]
WantedBy=multi-user.target
JUPYTER_SVC

        warn "JupyterLab service created WITHOUT authentication."
        warn "Run: jupyter lab password   to set a password before enabling."
        info "Enable with: systemctl enable --now jupyterlab"
    fi

    record_component "jupyter"
    log "Phase ${CURRENT_PHASE} complete."
}

# ═════════════════════════════════════════════════════════════════════════════
# OPTIONAL MODULES
# ═════════════════════════════════════════════════════════════════════════════

# ─────────────────────────────────────────────────────────────────────────────
# --with-docker : Docker CE + NVIDIA Container Toolkit
# ─────────────────────────────────────────────────────────────────────────────
install_docker() {
    [[ "${OPT_DOCKER}" != true ]] && return

    next_phase "Installing Docker CE + NVIDIA Container Toolkit..."

    # ── Docker CE ────────────────────────────────────────────────────────
    if ! command -v docker &>/dev/null; then
        log "Adding Docker repository..."
        if [[ "${OPT_DRY_RUN}" != true ]]; then
            install -m 0755 -d /etc/apt/keyrings
            curl -fsSL https://download.docker.com/linux/ubuntu/gpg \
                -o /etc/apt/keyrings/docker.asc
            chmod a+r /etc/apt/keyrings/docker.asc

            echo \
                "deb [arch=$(dpkg --print-architecture) signed-by=/etc/apt/keyrings/docker.asc] \
                https://download.docker.com/linux/ubuntu \
                $(. /etc/os-release && echo "${VERSION_CODENAME}") stable" \
                > /etc/apt/sources.list.d/docker.list

            apt-get update -qq
        else
            dry "Add Docker apt repository"
        fi

        apt_install \
            docker-ce \
            docker-ce-cli \
            containerd.io \
            docker-buildx-plugin \
            docker-compose-plugin
    else
        log "Docker already installed — skipping Docker CE install."
    fi

    # ── NVIDIA Container Toolkit ─────────────────────────────────────────
    log "Adding NVIDIA Container Toolkit repository..."
    if [[ "${OPT_DRY_RUN}" != true ]]; then
        curl -fsSL https://nvidia.github.io/libnvidia-container/gpgkey \
            | gpg --dearmor -o /usr/share/keyrings/nvidia-container-toolkit-keyring.gpg 2>/dev/null || true

        curl -s -L https://nvidia.github.io/libnvidia-container/stable/deb/nvidia-container-toolkit.list \
            | sed 's#deb https://#deb [signed-by=/usr/share/keyrings/nvidia-container-toolkit-keyring.gpg] https://#g' \
            > /etc/apt/sources.list.d/nvidia-container-toolkit.list

        apt-get update -qq
    else
        dry "Add NVIDIA Container Toolkit apt repository"
    fi

    apt_install nvidia-container-toolkit

    # Configure Docker runtime
    if [[ "${OPT_DRY_RUN}" != true ]]; then
        nvidia-ctk runtime configure --runtime=docker 2>/dev/null || true
        systemctl restart docker 2>/dev/null || true

        # Set nvidia as default runtime
        if [[ -f /etc/docker/daemon.json ]]; then
            python3 -c "
import json
with open('/etc/docker/daemon.json') as f:
    cfg = json.load(f)
cfg['default-runtime'] = 'nvidia'
with open('/etc/docker/daemon.json', 'w') as f:
    json.dump(cfg, f, indent=2)
" 2>/dev/null || warn "Could not set nvidia as default Docker runtime"
            systemctl restart docker 2>/dev/null || true
        fi
    fi

    record_component "docker"
    record_component "nvidia-container-toolkit"
    log "Phase ${CURRENT_PHASE} complete."
}

# ─────────────────────────────────────────────────────────────────────────────
# --with-profiling : NVIDIA DCGM + Nsight Systems/Compute + debugging tools
# ─────────────────────────────────────────────────────────────────────────────
install_profiling() {
    [[ "${OPT_PROFILING}" != true ]] && return

    next_phase "Installing DCGM + Nsight profiling suite..."

    local cuda_pkg_ver="${OPT_CUDA_VERSION//./-}"

    # ── DCGM (Data Center GPU Manager) ───────────────────────────────────
    log "Installing NVIDIA DCGM..."
    apt_install datacenter-gpu-manager 2>/dev/null || \
    apt_install nvidia-dcgm 2>/dev/null || \
        warn "DCGM package not found — install manually from NVIDIA"

    # Enable DCGM service
    if [[ "${OPT_DRY_RUN}" != true ]]; then
        systemctl enable nvidia-dcgm 2>/dev/null || true
        systemctl start nvidia-dcgm 2>/dev/null || true
    fi

    # DCGM Python bindings (pairs with Prometheus/sysmon)
    pip_install nvidia-dcgm 2>/dev/null || true

    # ── Nsight Systems ───────────────────────────────────────────────────
    log "Installing Nsight Systems..."
    apt_install nsight-systems 2>/dev/null || \
    apt_install nvidia-nsight-systems 2>/dev/null || \
        warn "Nsight Systems not found in repos — download from developer.nvidia.com"

    # ── Nsight Compute ───────────────────────────────────────────────────
    log "Installing Nsight Compute..."
    apt_install nsight-compute 2>/dev/null || \
    apt_install nvidia-nsight-compute 2>/dev/null || \
        warn "Nsight Compute not found in repos — download from developer.nvidia.com"

    # ── compute-sanitizer & cuda-gdb ─────────────────────────────────────
    log "Installing CUDA debugging tools..."
    apt_install \
        "cuda-sanitizer-${cuda_pkg_ver}" \
        "cuda-gdb-${cuda_pkg_ver}" 2>/dev/null || \
    apt_install compute-sanitizer cuda-gdb 2>/dev/null || true

    # ── NVIDIA perf tools ────────────────────────────────────────────────
    apt_install nvidia-utils-"${OPT_DRIVER_VERSION}" 2>/dev/null || true

    record_component "dcgm"
    record_component "nsight-systems"
    record_component "nsight-compute"
    record_component "compute-sanitizer"
    log "Phase ${CURRENT_PHASE} complete."
}

# ─────────────────────────────────────────────────────────────────────────────
# --with-inference : vLLM, Triton client, llama.cpp (CUDA build)
# ─────────────────────────────────────────────────────────────────────────────
install_inference() {
    [[ "${OPT_INFERENCE}" != true ]] && return

    next_phase "Installing inference stack..."

    # ── vLLM ─────────────────────────────────────────────────────────────
    log "Installing vLLM..."
    pip_install vllm || warn "vLLM installation failed"

    # ── Triton Inference Server client libraries ─────────────────────────
    log "Installing Triton client libraries..."
    pip_install "tritonclient[all]" || \
    pip_install "tritonclient[http,grpc]" || \
        warn "Triton client install failed"

    # ── llama.cpp (CUDA build) ───────────────────────────────────────────
    log "Building llama.cpp with CUDA support..."
    local LLAMACPP_DIR="/opt/llama.cpp"
    if [[ "${OPT_DRY_RUN}" != true ]]; then
        if [[ -d "${LLAMACPP_DIR}" ]]; then
            log "llama.cpp directory exists — pulling latest..."
            cd "${LLAMACPP_DIR}" && git pull --ff-only 2>/dev/null || true
        else
            git clone --depth 1 https://github.com/ggerganov/llama.cpp.git "${LLAMACPP_DIR}"
        fi
        cd "${LLAMACPP_DIR}"

        # Source CUDA env
        export CUDA_HOME=/usr/local/cuda
        export PATH="${CUDA_HOME}/bin:${PATH}"

        # CMake build with CUDA — broad compute capability coverage
        cmake -B build \
            -DGGML_CUDA=ON \
            -DCMAKE_CUDA_ARCHITECTURES="70;75;80;86;89;90" \
            -DCMAKE_BUILD_TYPE=Release 2>&1 | tail -5
        cmake --build build --config Release -j"$(nproc)" 2>&1 | tail -5

        # Symlink binaries to PATH
        for bin in llama-server llama-cli llama-quantize llama-bench; do
            if [[ -x "${LLAMACPP_DIR}/build/bin/${bin}" ]]; then
                ln -sf "${LLAMACPP_DIR}/build/bin/${bin}" "/usr/local/bin/${bin}"
            fi
        done

        cd /tmp
    else
        dry "git clone llama.cpp && cmake -DGGML_CUDA=ON && make -j$(nproc)"
    fi

    # ── TGI client ───────────────────────────────────────────────────────
    pip_install text-generation 2>/dev/null || true

    # ── Create llama-server systemd template ─────────────────────────────
    if [[ "${OPT_DRY_RUN}" != true ]]; then
        cat > /etc/systemd/system/llama-server@.service << 'LLAMA_SVC'
[Unit]
Description=llama.cpp server - %i (ICC AI Stack)
After=network.target

[Service]
Type=simple
ExecStart=/usr/local/bin/llama-server \
    -m /opt/models/%i \
    --port 8080 \
    --n-gpu-layers 99 \
    --ctx-size 4096
Restart=on-failure
RestartSec=10
Environment="CUDA_VISIBLE_DEVICES=0"

[Install]
WantedBy=multi-user.target
LLAMA_SVC

        mkdir -p /opt/models
        info "llama-server systemd template created."
        info "Usage: systemctl start llama-server@model-name.gguf"
        info "Place models in /opt/models/"
    fi

    record_component "vllm"
    record_component "triton-client"
    record_component "llama-cpp"
    log "Phase ${CURRENT_PHASE} complete."
}

# ─────────────────────────────────────────────────────────────────────────────
# --with-cluster : Open MPI + NCCL tests + bandwidth test + cluster utils
# ─────────────────────────────────────────────────────────────────────────────
install_cluster() {
    [[ "${OPT_CLUSTER}" != true ]] && return

    next_phase "Installing cluster & multi-node tools..."

    # ── Open MPI ─────────────────────────────────────────────────────────
    log "Installing Open MPI..."
    apt_install \
        openmpi-bin \
        openmpi-common \
        libopenmpi-dev 2>/dev/null || \
    apt_install openmpi-bin libopenmpi-dev || warn "OpenMPI install failed"

    # Check CUDA awareness
    if [[ "${OPT_DRY_RUN}" != true ]]; then
        if ompi_info --parsable --all 2>/dev/null | grep -qi 'mpi_built_with_cuda_support:value:true'; then
            log "OpenMPI CUDA support: ENABLED"
        else
            warn "System OpenMPI may not have CUDA support compiled in."
            warn "For production multi-node, consider building OpenMPI from source with --with-cuda."
        fi
    fi

    # ── NCCL Tests ───────────────────────────────────────────────────────
    log "Building NCCL tests..."
    local NCCL_TESTS_DIR="/opt/nccl-tests"
    if [[ "${OPT_DRY_RUN}" != true ]]; then
        export CUDA_HOME=/usr/local/cuda
        export PATH="${CUDA_HOME}/bin:${PATH}"
        export LD_LIBRARY_PATH="${CUDA_HOME}/lib64:${LD_LIBRARY_PATH:-}"

        if [[ -d "${NCCL_TESTS_DIR}" ]]; then
            cd "${NCCL_TESTS_DIR}" && git pull --ff-only 2>/dev/null || true
        else
            git clone --depth 1 https://github.com/NVIDIA/nccl-tests.git "${NCCL_TESTS_DIR}"
        fi
        cd "${NCCL_TESTS_DIR}"

        make -j"$(nproc)" \
            CUDA_HOME=/usr/local/cuda \
            MPI=1 \
            MPI_HOME=/usr/lib/x86_64-linux-gnu/openmpi 2>&1 | tail -5 || \
        make -j"$(nproc)" \
            CUDA_HOME=/usr/local/cuda 2>&1 | tail -5 || \
            warn "NCCL tests build failed"

        # Symlink test binaries
        if [[ -d "${NCCL_TESTS_DIR}/build" ]]; then
            for bin in "${NCCL_TESTS_DIR}/build/"*_perf; do
                [[ -x "${bin}" ]] && ln -sf "${bin}" /usr/local/bin/ 2>/dev/null || true
            done
        fi
        cd /tmp
    else
        dry "git clone nccl-tests && make CUDA_HOME=/usr/local/cuda MPI=1"
    fi

    # ── CUDA Bandwidth Test ──────────────────────────────────────────────
    log "Installing bandwidth test..."
    if [[ "${OPT_DRY_RUN}" != true ]]; then
        local samples_dir="/usr/local/cuda/extras/demo_suite"
        if [[ -x "${samples_dir}/bandwidthTest" ]]; then
            ln -sf "${samples_dir}/bandwidthTest" /usr/local/bin/cuda-bandwidth-test
            log "bandwidthTest linked from CUDA samples"
        else
            warn "bandwidthTest not found in CUDA extras — may need manual build"
        fi

        # Also link deviceQuery if available
        if [[ -x "${samples_dir}/deviceQuery" ]]; then
            ln -sf "${samples_dir}/deviceQuery" /usr/local/bin/cuda-device-query
        fi
    fi

    # ── Parallel SSH / cluster management ────────────────────────────────
    apt_install pssh pdsh 2>/dev/null || true
    pip_install fabric paramiko 2>/dev/null || true

    record_component "openmpi"
    record_component "nccl-tests"
    record_component "cluster-tools"
    log "Phase ${CURRENT_PHASE} complete."
}

# ─────────────────────────────────────────────────────────────────────────────
# Post-install validation
# ─────────────────────────────────────────────────────────────────────────────
validate_install() {
    log "Running post-installation validation..."
    echo ""

    local pass=0 fail=0

    check_cmd() {
        local label="$1" cmd="$2"
        if eval "${cmd}" &>/dev/null; then
            echo -e "  ${GREEN}✓${NC} ${label}"
            ((pass++))
        else
            echo -e "  ${RED}✗${NC} ${label}"
            ((fail++))
        fi
    }

    echo -e "${BOLD}Core Components:${NC}"
    check_cmd "NVIDIA Driver"        "nvidia-smi"
    check_cmd "CUDA Compiler (nvcc)" "nvcc --version"
    check_cmd "Python 3"             "python3 --version"
    check_cmd "uv"                   "uv --version 2>/dev/null || test -x ~/.cargo/bin/uv"

    if [[ "${OPT_DRIVER_ONLY}" != true ]] && [[ "${OPT_NO_FRAMEWORKS}" != true ]]; then
        echo ""
        echo -e "${BOLD}Frameworks:${NC}"
        check_cmd "PyTorch"          "python3 -c 'import torch; print(torch.__version__)'"
        check_cmd "PyTorch CUDA"     "python3 -c 'import torch; assert torch.cuda.is_available()'"
        check_cmd "TensorFlow"       "python3 -c 'import tensorflow as tf; print(tf.__version__)'"
        check_cmd "TF GPU"           "python3 -c 'import tensorflow as tf; assert len(tf.config.list_physical_devices(\"GPU\")) > 0'"
        check_cmd "JAX"              "python3 -c 'import jax; print(jax.__version__)'"
    fi

    if [[ "${OPT_DRIVER_ONLY}" != true ]]; then
        echo ""
        echo -e "${BOLD}Python Ecosystem:${NC}"
        check_cmd "NumPy"            "python3 -c 'import numpy'"
        check_cmd "Transformers"     "python3 -c 'import transformers'"
        check_cmd "PEFT"             "python3 -c 'import peft'"
        check_cmd "bitsandbytes"     "python3 -c 'import bitsandbytes'"
        check_cmd "HF CLI"           "huggingface-cli --help"
        check_cmd "DeepSpeed"        "python3 -c 'import deepspeed'"
        check_cmd "OpenAI Triton"    "python3 -c 'import triton'"
    fi

    if [[ "${OPT_NO_JUPYTER}" != true ]] && [[ "${OPT_DRIVER_ONLY}" != true ]]; then
        check_cmd "JupyterLab"       "jupyter lab --version"
    fi

    # ── Optional module checks ───────────────────────────────────────────
    if [[ "${OPT_DOCKER}" == true ]]; then
        echo ""
        echo -e "${BOLD}Docker:${NC}"
        check_cmd "Docker"                    "docker --version"
        check_cmd "Docker Compose"            "docker compose version"
        check_cmd "NVIDIA Container Toolkit"  "nvidia-ctk --version"
    fi

    if [[ "${OPT_PROFILING}" == true ]]; then
        echo ""
        echo -e "${BOLD}Profiling & Debugging:${NC}"
        check_cmd "DCGM"                 "nv-hostengine --version 2>/dev/null || dcgmi discovery -l 2>/dev/null"
        check_cmd "Nsight Systems (nsys)" "nsys --version"
        check_cmd "Nsight Compute (ncu)"  "ncu --version"
        check_cmd "compute-sanitizer"     "compute-sanitizer --version"
        check_cmd "cuda-gdb"              "cuda-gdb --version"
    fi

    if [[ "${OPT_INFERENCE}" == true ]]; then
        echo ""
        echo -e "${BOLD}Inference:${NC}"
        check_cmd "vLLM"              "python3 -c 'import vllm; print(vllm.__version__)'"
        check_cmd "Triton client"     "python3 -c 'import tritonclient'"
        check_cmd "llama-server"      "test -x /usr/local/bin/llama-server"
        check_cmd "llama-cli"         "test -x /usr/local/bin/llama-cli"
        check_cmd "llama-quantize"    "test -x /usr/local/bin/llama-quantize"
        check_cmd "llama-bench"       "test -x /usr/local/bin/llama-bench"
    fi

    if [[ "${OPT_CLUSTER}" == true ]]; then
        echo ""
        echo -e "${BOLD}Cluster & Multi-Node:${NC}"
        check_cmd "Open MPI (mpirun)"    "mpirun --version"
        check_cmd "all_reduce_perf"      "test -x /usr/local/bin/all_reduce_perf"
        check_cmd "all_gather_perf"      "test -x /usr/local/bin/all_gather_perf"
        check_cmd "reduce_scatter_perf"  "test -x /usr/local/bin/reduce_scatter_perf"
    fi

    echo ""
    echo -e "${BOLD}Results:${NC} ${GREEN}${pass} passed${NC}, ${RED}${fail} failed${NC}"

    if [[ ${fail} -gt 0 ]]; then
        warn "Some components failed validation. Check ${LOG_FILE} for details."
        warn "A reboot may be required for the NVIDIA driver to load."
    fi

    # GPU info
    if nvidia-smi &>/dev/null; then
        echo ""
        echo -e "${BOLD}GPU Information:${NC}"
        nvidia-smi --query-gpu=index,name,driver_version,memory.total,compute_cap \
            --format=csv,noheader 2>/dev/null | while IFS=',' read -r idx name drv mem cap; do
            info "  GPU ${idx}:${name} | Driver:${drv} | VRAM:${mem} | Compute:${cap}"
        done
    fi
}

# ─────────────────────────────────────────────────────────────────────────────
# Environment summary
# ─────────────────────────────────────────────────────────────────────────────
print_summary() {
    echo ""
    echo -e "${CYAN}═══════════════════════════════════════════════════════════════${NC}"
    echo -e "${BOLD} ICC AI/HPC Stack v${STACK_VERSION} — Installation Complete${NC}"
    echo -e "${CYAN}═══════════════════════════════════════════════════════════════${NC}"
    echo ""
    echo -e " ${BOLD}CUDA:${NC}        ${OPT_CUDA_VERSION}"
    echo -e " ${BOLD}Driver:${NC}      ${OPT_DRIVER_VERSION}"

    if [[ "${OPT_DRIVER_ONLY}" != true ]] && [[ "${OPT_NO_FRAMEWORKS}" != true ]]; then
        local pt_ver tf_ver jax_ver
        pt_ver=$(python3 -c 'import torch; print(torch.__version__)' 2>/dev/null || echo "N/A")
        tf_ver=$(python3 -c 'import tensorflow as tf; print(tf.__version__)' 2>/dev/null || echo "N/A")
        jax_ver=$(python3 -c 'import jax; print(jax.__version__)' 2>/dev/null || echo "N/A")
        echo -e " ${BOLD}PyTorch:${NC}     ${pt_ver}"
        echo -e " ${BOLD}TensorFlow:${NC}  ${tf_ver}"
        echo -e " ${BOLD}JAX:${NC}         ${jax_ver}"
    fi

    echo ""
    echo -e " ${BOLD}Optional modules:${NC}"
    echo -e "   Docker + NVIDIA CTK:     $( [[ "${OPT_DOCKER}" == true ]]    && echo "${GREEN}✓ Installed${NC}" || echo "${BLUE}— skipped${NC}" )"
    echo -e "   DCGM + Nsight:           $( [[ "${OPT_PROFILING}" == true ]] && echo "${GREEN}✓ Installed${NC}" || echo "${BLUE}— skipped${NC}" )"
    echo -e "   Inference (vLLM/llama):  $( [[ "${OPT_INFERENCE}" == true ]] && echo "${GREEN}✓ Installed${NC}" || echo "${BLUE}— skipped${NC}" )"
    echo -e "   Cluster (MPI/NCCL):     $( [[ "${OPT_CLUSTER}" == true ]]   && echo "${GREEN}✓ Installed${NC}" || echo "${BLUE}— skipped${NC}" )"

    echo ""
    echo -e " ${BOLD}Files:${NC}"
    echo -e "   Log:        ${LOG_FILE}"
    echo -e "   State:      ${STATE_FILE}"
    echo -e "   CUDA env:   /etc/profile.d/icc-cuda.sh"

    if [[ "${OPT_DOCKER}" == true ]]; then
        echo ""
        echo -e " ${BOLD}Docker GPU test:${NC}"
        echo "   docker run --rm --gpus all nvidia/cuda:${OPT_CUDA_VERSION}.0-base-ubuntu24.04 nvidia-smi"
    fi

    if [[ "${OPT_INFERENCE}" == true ]]; then
        echo ""
        echo -e " ${BOLD}Inference quick start:${NC}"
        echo "   # llama.cpp"
        echo "   llama-server -m /opt/models/model.gguf --port 8080 --n-gpu-layers 99"
        echo "   # or as a service:"
        echo "   systemctl start llama-server@model-name.gguf"
        echo ""
        echo "   # vLLM (HuggingFace model)"
        echo "   vllm serve meta-llama/Llama-3-8B-Instruct --port 8000"
    fi

    if [[ "${OPT_CLUSTER}" == true ]]; then
        echo ""
        echo -e " ${BOLD}Cluster validation:${NC}"
        echo "   # Single-node multi-GPU NCCL test"
        echo "   all_reduce_perf -b 8 -e 256M -f 2 -g <num_gpus>"
        echo ""
        echo "   # Multi-node (2 hosts, 1 GPU each)"
        echo "   mpirun -np 2 -H host1:1,host2:1 --allow-run-as-root \\"
        echo "     all_reduce_perf -b 8 -e 256M -f 2 -g 1"
    fi

    if [[ "${OPT_PROFILING}" == true ]]; then
        echo ""
        echo -e " ${BOLD}Profiling quick start:${NC}"
        echo "   # GPU health check"
        echo "   dcgmi diag -r 1"
        echo ""
        echo "   # Profile a training run"
        echo "   nsys profile -o report python3 train.py"
        echo "   ncu --set full -o kernel_report python3 train.py"
    fi

    echo ""
    echo -e " ${YELLOW}⚠  A reboot is recommended to ensure the NVIDIA driver loads.${NC}"
    echo -e " ${YELLOW}⚠  Run 'source /etc/profile.d/icc-cuda.sh' or re-login for PATH.${NC}"
    echo ""
    echo -e " ${BOLD}Quick test:${NC}"
    echo "   nvidia-smi"
    echo "   python3 -c \"import torch; print(torch.cuda.get_device_name(0))\""
    echo ""
}

# ─────────────────────────────────────────────────────────────────────────────
# Uninstall
# ─────────────────────────────────────────────────────────────────────────────
do_uninstall() {
    log "Uninstalling ICC AI/HPC Stack..."

    if [[ "${OPT_YES}" != true ]]; then
        warn "This will remove NVIDIA drivers, CUDA, frameworks, and all optional modules."
        read -rp "Are you sure? [y/N] " ans
        [[ "${ans}" =~ ^[Yy] ]] || exit 0
    fi

    # Python packages
    log "Removing Python packages..."
    pip3 uninstall -y \
        torch torchvision torchaudio \
        tensorflow \
        jax jaxlib \
        jupyterlab notebook \
        transformers datasets accelerate peft bitsandbytes \
        deepspeed flash-attn triton \
        vllm tritonclient text-generation \
        nvitop gpustat tensorboard wandb \
        onnx onnxruntime-gpu \
        nvidia-dcgm \
        fabric paramiko \
        2>/dev/null || true

    # NVIDIA/CUDA/system packages
    log "Removing NVIDIA/CUDA packages..."
    apt-get remove --purge -y \
        'cuda*' 'nvidia*' 'libcudnn*' 'libnccl*' \
        'tensorrt*' 'libnvinfer*' \
        'nsight*' 'datacenter-gpu-manager*' \
        nvidia-container-toolkit \
        2>/dev/null || true
    apt-get autoremove -y 2>/dev/null || true

    # Built-from-source artifacts
    log "Removing built artifacts..."
    rm -rf /opt/llama.cpp /opt/nccl-tests /opt/models
    rm -f /usr/local/bin/llama-server /usr/local/bin/llama-cli
    rm -f /usr/local/bin/llama-quantize /usr/local/bin/llama-bench
    rm -f /usr/local/bin/all_reduce_perf /usr/local/bin/all_gather_perf
    rm -f /usr/local/bin/reduce_scatter_perf /usr/local/bin/alltoall_perf
    rm -f /usr/local/bin/broadcast_perf /usr/local/bin/sendrecv_perf
    rm -f /usr/local/bin/cuda-bandwidth-test /usr/local/bin/cuda-device-query

    # Environment and service files
    rm -f /etc/profile.d/icc-cuda.sh
    rm -f /etc/ld.so.conf.d/icc-cuda.conf
    rm -f /etc/systemd/system/jupyterlab.service
    rm -f /etc/systemd/system/llama-server@.service
    rm -rf "${STATE_DIR}"
    ldconfig 2>/dev/null || true
    systemctl daemon-reload 2>/dev/null || true

    log "Uninstall complete. Reboot recommended."
}

# ─────────────────────────────────────────────────────────────────────────────
# Main
# ─────────────────────────────────────────────────────────────────────────────
main() {
    parse_args "$@"
    banner

    if [[ "${OPT_UNINSTALL}" == true ]]; then
        check_root
        do_uninstall
        exit 0
    fi

    # Compute total phases
    TOTAL_PHASES=$(count_phases)

    # Build module list string
    local modules_str=""
    [[ "${OPT_DOCKER}" == true ]]    && modules_str+="Docker, "
    [[ "${OPT_PROFILING}" == true ]] && modules_str+="Profiling, "
    [[ "${OPT_INFERENCE}" == true ]] && modules_str+="Inference, "
    [[ "${OPT_CLUSTER}" == true ]]   && modules_str+="Cluster, "
    modules_str="${modules_str%, }"
    [[ -z "${modules_str}" ]] && modules_str="None"

    # Confirmation
    echo -e "${BOLD}Installation plan:${NC}"
    echo "  Target:      ${SUPPORTED_OS} ($(uname -r))"
    echo "  CUDA:        ${OPT_CUDA_VERSION}"
    echo "  Driver:      ${OPT_DRIVER_VERSION}"
    echo "  Frameworks:  $( [[ "${OPT_DRIVER_ONLY}" == true || "${OPT_NO_FRAMEWORKS}" == true ]] && echo 'SKIP' || echo 'PyTorch + TensorFlow + JAX' )"
    echo "  Jupyter:     $( [[ "${OPT_DRIVER_ONLY}" == true || "${OPT_NO_JUPYTER}" == true ]] && echo 'SKIP' || echo 'Yes' )"
    echo "  Modules:     ${modules_str}"
    echo "  Phases:      ${TOTAL_PHASES}"
    echo "  Dry run:     ${OPT_DRY_RUN}"
    echo ""

    if [[ "${OPT_YES}" != true ]] && [[ "${OPT_DRY_RUN}" != true ]]; then
        read -rp "Proceed with installation? [Y/n] " ans
        [[ "${ans}" =~ ^[Nn] ]] && exit 0
    fi

    check_root
    check_os
    check_gpu
    check_disk_space
    check_internet
    check_existing_install

    local start_time
    start_time=$(date +%s)

    # Core phases
    install_prerequisites
    install_nvidia_driver
    install_cuda
    install_frameworks
    install_python_ecosystem
    install_jupyter

    # Optional modules (each is a no-op if flag not set)
    install_docker
    install_profiling
    install_inference
    install_cluster

    local elapsed=$(( $(date +%s) - start_time ))
    log "Total installation time: $((elapsed / 60))m $((elapsed % 60))s"

    if [[ "${OPT_DRY_RUN}" != true ]]; then
        validate_install
        print_summary
    else
        echo ""
        log "Dry run complete. No changes were made."
    fi
}

main "$@"