#!/bin/sh

# -----------------------
# Config
# -----------------------
NODES="watgpu-100.cs.uwaterloo.ca watgpu-200.cs.uwaterloo.ca watgpu-300.cs.uwaterloo.ca watgpu-400.cs.private.uwaterloo.ca watgpu-500.cs.private.uwaterloo.ca watgpu-600.cs.private.uwaterloo.ca"   # list of cluster nodes
SSH_USER="root"
MD_OUTPUT="./gpu_layout.md"

# -----------------------
# Start Markdown file
# -----------------------
{
    echo "# Cluster GPU Layout"
    echo
} > "$MD_OUTPUT"

# -----------------------
# Loop over nodes
# -----------------------
for host in $NODES; do
    echo "[INFO] Fetching GPU info from $host..." >&2

    GPUs=$(ssh -o BatchMode=yes -o ConnectTimeout=5 "${SSH_USER}@${host}" \
        "nvidia-smi --query-gpu=index,name,memory.total --format=csv,noheader" 2>/dev/null)

    {
        echo "## ${host%%.*}"
        echo
        if [ -z "$GPUs" ]; then
            echo "_No GPUs found_"
            echo
            continue
        fi
        echo "| Index | Model | Memory |"
        echo "|-------|-------|--------|"
        echo "$GPUs" | while IFS=, read -r idx name mem; do
            mem_gb=$(echo "$mem" | awk '{printf "%.1f", $1/1024}')
            echo "| $(echo "$idx" | xargs) | $(echo "$name" | xargs) | ${mem_gb} GB |"
        done
        echo
    } >> "$MD_OUTPUT"
done

echo "[INFO] GPU layout written to $MD_OUTPUT" >&2

echo "[INFO] Copying file into LXC container $LXC_NAME:$LXC_PATH" >&2
lxc file push "$MD_OUTPUT" "$watgpu/var/www/src/gpu_layout.md"

rm -f "$MD_OUTPUT"
echo "[INFO] Cleanup done. GPU layout updated inside container." >&2