#!/usr/bin/env bash
# =============================================================================
# fermihdi-discover — inventory a host for FermiHDI installation
# =============================================================================
# Emits a YAML facts document describing CPU topology, IOMMU groups, candidate
# storage devices and NICs. fermihdi-configure consumes it to build cluster.yaml.
#
# Reads /sys and /proc only. No lspci, no numactl, no Python. Runs on a bare
# host before anything is installed.
#
# Usage:
#   ./fermihdi-discover                 YAML facts to stdout
#   ./fermihdi-discover -o facts.yaml   write to a file
#   ./fermihdi-discover --format=tsv    tab-separated, for shell consumers
#   ./fermihdi-discover --check-deps    report tool availability and exit
#
# The tsv form is what fermihdi-configure reads, so it can stay pure shell
# rather than depending on a YAML parser. One record per line, first field is
# the record type.
#
# Read-only. Changes nothing on the host.
# =============================================================================
set -uo pipefail

OUT=""
CHECK_ONLY=0
FORMAT=yaml
while [ $# -gt 0 ]; do
    case "$1" in
        -o|--output) OUT="${2:-}"; shift 2 ;;
        --format) FORMAT="${2:-yaml}"; shift 2 ;;
        --format=*) FORMAT="${1#*=}"; shift ;;
        --check-deps) CHECK_ONLY=1; shift ;;
        -h|--help) sed -n '2,20p' "$0" | sed 's/^# \?//'; exit 0 ;;
        *) echo "unknown argument: $1" >&2; exit 2 ;;
    esac
done

# ── Dependencies ─────────────────────────────────────────────────────────────
# Required tools are all util-linux or coreutils and present on any Ubuntu
# install. They are still checked, because failing here with a package name is
# kinder than failing later with an empty inventory.
REQUIRED="lsblk findmnt lscpu awk sed grep sort tr cat"
OPTIONAL="ip"
declare -A PKG=( [lsblk]=util-linux [findmnt]=util-linux [lscpu]=util-linux [ip]=iproute2 )

missing=""
for t in $REQUIRED; do command -v "$t" >/dev/null 2>&1 || missing="$missing $t"; done

if [ -n "$missing" ]; then
    pkgs=""
    for t in $missing; do pkgs="$pkgs ${PKG[$t]:-$t}"; done
    pkgs=$(echo "$pkgs" | tr ' ' '\n' | sort -u | tr '\n' ' ')
    echo "Missing required tools:$missing" >&2
    echo "Install with:  sudo apt-get install -y$pkgs" >&2
    if [ -t 0 ] && [ -t 1 ]; then
        printf 'Install now? [y/N] ' >&2
        read -r ans </dev/tty
        case "$ans" in
            [yY]*) sudo apt-get update -qq && sudo apt-get install -y $pkgs || exit 1 ;;
            *) exit 1 ;;
        esac
    else
        exit 1
    fi
fi

for t in $OPTIONAL; do
    command -v "$t" >/dev/null 2>&1 || echo "note: $t not found — network facts will be reduced" >&2
done

if [ "$CHECK_ONLY" = 1 ]; then
    for t in $REQUIRED $OPTIONAL; do
        printf '  %-10s %s\n' "$t" "$(command -v "$t" >/dev/null 2>&1 && echo present || echo MISSING)"
    done
    exit 0
fi

[ -n "$OUT" ] && exec > "$OUT"

# ── Helpers ──────────────────────────────────────────────────────────────────
r() { [ -r "$1" ] && tr -d '\0\n' < "$1" 2>/dev/null || true; }   # read a sysfs value

# IOMMU group for a PCI address, empty when IOMMU is off.
iommu_group_of() {
    local g; g=$(readlink -f "/sys/bus/pci/devices/$1/iommu_group" 2>/dev/null) || return 0
    [ -n "$g" ] && basename "$g" || true
}

# NUMA node a PCI device is attached to. "-1" when the platform does not say.
# Per-node memory and hugepage state. Hugepages are reserved per NUMA node, so
# a host total is not enough to tell whether an instance pinned to node 1 can
# actually get its pages -- node 0 may hold all of them.
node_mem_mb() { # node_mem_mb <n>
    local kb; kb=$(awk '/MemTotal:/{print $4}' "/sys/devices/system/node/node$1/meminfo" 2>/dev/null)
    echo $(( ${kb:-0} / 1024 ))
}
node_hp() { # node_hp <n> <2048kB|1048576kB>
    local v; v=$(r "/sys/devices/system/node/node$1/hugepages/hugepages-$2/nr_hugepages")
    echo "${v:-0}"
}

numa_of() {
    local n; n=$(r "/sys/bus/pci/devices/$1/numa_node"); [ -n "$n" ] && echo "$n" || echo "-1"
}

# Current kernel driver bound to a PCI address.
driver_of() {
    local d; d=$(readlink -f "/sys/bus/pci/devices/$1/driver" 2>/dev/null) || return 0
    [ -n "$d" ] && basename "$d" || true
}

# The block device holding / — never offer this for VFIO.
root_src=$(findmnt -no SOURCE / 2>/dev/null || true)
root_disk=$(lsblk -no PKNAME "$root_src" 2>/dev/null | head -1 || true)
[ -z "$root_disk" ] && root_disk=$(basename "${root_src:-none}" | sed 's/[0-9]*$//')

# The interface carrying the default route — never offer this for VFIO either.
default_if=""
command -v ip >/dev/null 2>&1 && default_if=$(ip -o route show default 2>/dev/null | awk '{print $5; exit}')

logical=$(getconf _NPROCESSORS_ONLN 2>/dev/null || echo 0)
sockets=$(lscpu 2>/dev/null | awk -F: '/^Socket\(s\)/{gsub(/ /,"",$2);print $2}')
per_socket=$(lscpu 2>/dev/null | awk -F: '/^Core\(s\) per socket/{gsub(/ /,"",$2);print $2}')
tpc=$(lscpu 2>/dev/null | awk -F: '/^Thread\(s\) per core/{gsub(/ /,"",$2);print $2}')
numa=$(lscpu 2>/dev/null | awk -F: '/^NUMA node\(s\)/{gsub(/ /,"",$2);print $2}')

emit_tsv() {
    # One record per line, tab separated, first field is the record type.
    # Deliberately flat: a shell consumer can read it with `while IFS=$'\t' read`.
    printf 'host\t%s\t%s\t%s\n' "$(hostname)" "$(uname -r)" "$(grep -qiE hypervisor /proc/cpuinfo && echo virtual || echo bare)"
    printf 'cpu\t%s\t%s\t%s\t%s\t%s\n' "${logical:-0}" "${sockets:-1}" "${per_socket:-0}" "${tpc:-1}" "${numa:-1}"
    for c in /sys/devices/system/cpu/cpu[0-9]*; do
        n=${c##*/cpu}; sib=$(r "$c/topology/thread_siblings_list")
        [ -n "$sib" ] && printf 'sibling\t%s\t%s\n' "$n" "$sib"
    done
    for nd in /sys/devices/system/node/node[0-9]*; do
        [ -d "$nd" ] || continue
        _n=${nd##*/node}
        printf 'numa\t%s\t%s\t%s\t%s\t%s\n' "$_n" "$(r "$nd/cpulist")" \
            "$(node_mem_mb "$_n")" "$(node_hp "$_n" 2048kB)" "$(node_hp "$_n" 1048576kB)"
    done
    if [ -d /sys/kernel/iommu_groups ] && [ -n "$(ls -A /sys/kernel/iommu_groups 2>/dev/null)" ]; then
        printf 'iommu\ttrue\n'
    else
        printf 'iommu\tfalse\n'
    fi
    printf 'memory\t%s\t%s\n' "$(awk '/^MemTotal:/{print $2}' /proc/meminfo)" "$(awk '/^MemAvailable:/{print $2}' /proc/meminfo)"
    printf 'hugepages\t2m\t%s\t%s\n' "$(r /sys/kernel/mm/hugepages/hugepages-2048kB/nr_hugepages)" "$(r /sys/kernel/mm/hugepages/hugepages-2048kB/free_hugepages)"
    printf 'hugepages\t1g\t%s\t%s\n' "$(r /sys/kernel/mm/hugepages/hugepages-1048576kB/nr_hugepages)" "$(r /sys/kernel/mm/hugepages/hugepages-1048576kB/free_hugepages)"

    lsblk -dn -P -o NAME,SIZE,MODEL,TRAN,TYPE 2>/dev/null | while read -r line; do
        NAME=""; SIZE=""; MODEL=""; TRAN=""; TYPE=""; eval "$line"
        [ "$TYPE" = "disk" ] || continue
        holds_root=false; [ "$NAME" = "$root_disk" ] && holds_root=true
        mounted=$(lsblk -no MOUNTPOINTS "/dev/$NAME" 2>/dev/null | tr -d " \n" | head -c1)
        has_fs=$(lsblk -no FSTYPE "/dev/$NAME" 2>/dev/null | tr -d " \n" | head -c1)
        pci=""; grp=""; drv=""
        if [ -e "/sys/block/$NAME/device" ]; then
            real=$(readlink -f "/sys/block/$NAME/device" 2>/dev/null)
            while [ -n "$real" ] && [ "$real" != "/" ]; do
                base=${real##*/}
                case "$base" in
                    [0-9a-f][0-9a-f][0-9a-f][0-9a-f]:[0-9a-f][0-9a-f]:[0-9a-f][0-9a-f].[0-9a-f]) pci=$base; break ;;
                esac
                real=$(dirname "$real")
            done
            if [ -n "$pci" ]; then grp=$(iommu_group_of "$pci"); drv=$(driver_of "$pci"); fi
        fi
        case "$NAME" in
            nvme*) kind=nvme ;;
            *)     kind=$(echo "${TRAN:-unknown}" | tr "A-Z" "a-z") ;;
        esac
        elig=true; why=""
        if [ "$holds_root" = true ]; then elig=false; why="holds the root filesystem"
        elif [ -n "$mounted" ]; then elig=false; why="currently mounted"; fi
        nn="-1"; [ -n "$pci" ] && nn=$(numa_of "$pci")
        # name path size model kind pci group driver eligible reason has_fs root numa
        printf 'storage\t%s\t/dev/%s\t%s\t%s\t%s\t%s\t%s\t%s\t%s\t%s\t%s\t%s\t%s\n' \
            "$NAME" "$NAME" "$SIZE" "${MODEL:-unknown}" "$kind" "${pci:--}" "${grp:--}" "${drv:--}" \
            "$elig" "${why:--}" "$([ -n "$has_fs" ] && echo true || echo false)" "$holds_root" "$nn"
    done

    for n in /sys/class/net/*; do
        ifn=${n##*/}; [ "$ifn" = "lo" ] && continue; [ -e "$n/device" ] || continue
        pci=$(basename "$(readlink -f "$n/device")")
        case "$pci" in [0-9a-f][0-9a-f][0-9a-f][0-9a-f]:*) ;; *) continue ;; esac
        isdef=false; [ "$ifn" = "$default_if" ] && isdef=true
        addr=""; command -v ip >/dev/null 2>&1 && addr=$(ip -o -4 addr show "$ifn" 2>/dev/null | awk "{print \$4; exit}")
        grp=$(iommu_group_of "$pci")
        # iface pci driver state ipv4 group default_route eligible numa
        printf 'nic\t%s\t%s\t%s\t%s\t%s\t%s\t%s\t%s\t%s\n' \
            "$ifn" "$pci" "$(driver_of "$pci")" "$(r "$n/operstate")" "${addr:--}" "${grp:--}" \
            "$isdef" "$([ "$isdef" = true ] && echo false || echo true)" "$(numa_of "$pci")"
    done
}

if [ "$FORMAT" = "tsv" ]; then emit_tsv; exit 0; fi

echo "# fermihdi-discover facts — generated $(date -u +%Y-%m-%dT%H:%M:%SZ)"
echo "# Read-only inventory. Nothing on this host was modified."
echo "version: 1"

# ── Host ─────────────────────────────────────────────────────────────────────
echo "host:"
echo "  name: $(hostname)"
echo "  kernel: $(uname -r)"
echo "  distro: \"$( . /etc/os-release 2>/dev/null && echo "${PRETTY_NAME:-unknown}" )\""
echo "  virtualised: $(grep -qiE 'hypervisor' /proc/cpuinfo && echo true || echo false)"

# ── CPU ──────────────────────────────────────────────────────────────────────
echo "cpu:"
echo "  logical: ${logical:-0}"
echo "  sockets: ${sockets:-1}"
echo "  cores_per_socket: ${per_socket:-0}"
echo "  threads_per_core: ${tpc:-1}"
echo "  numa_nodes: ${numa:-1}"
echo "  numa:"
for nd in /sys/devices/system/node/node[0-9]*; do
    [ -d "$nd" ] || continue
    _n=${nd##*/node}
    echo "    \"$_n\":"
    echo "      cpulist: \"$(r "$nd/cpulist")\""
    echo "      memory_mb: $(node_mem_mb "$_n")"
    echo "      hugepages_2m: $(node_hp "$_n" 2048kB)"
    echo "      hugepages_1g: $(node_hp "$_n" 1048576kB)"
done
echo "  smt: $([ "${tpc:-1}" -gt 1 ] 2>/dev/null && echo true || echo false)"
echo "  isolated_now: \"$(r /sys/devices/system/cpu/isolated)\""
# SMT sibling pairs matter: a poll-mode thread must not share a physical core.
echo "  siblings:"
for c in /sys/devices/system/cpu/cpu[0-9]*; do
    n=${c##*/cpu}
    s=$(r "$c/topology/thread_siblings_list")
    [ -n "$s" ] && echo "    \"$n\": \"$s\""
done | sort -t'"' -k2 -n

# ── IOMMU ────────────────────────────────────────────────────────────────────
echo "iommu:"
if [ -d /sys/kernel/iommu_groups ] && [ -n "$(ls -A /sys/kernel/iommu_groups 2>/dev/null)" ]; then
    echo "  enabled: true"
    echo "  groups:"
    for g in /sys/kernel/iommu_groups/*; do
        gid=${g##*/}
        devs=""
        for d in "$g"/devices/*; do devs="$devs \"${d##*/}\","; done
        echo "    \"$gid\": [${devs%,} ]"
    done
else
    echo "  enabled: false"
    echo "  groups: {}"
    echo "  # No IOMMU groups. VFIO cannot be used until intel_iommu=on (or"
    echo "  # amd_iommu=on) plus iommu=pt are set in GRUB and the host rebooted."
fi

# ── Hugepages ────────────────────────────────────────────────────────────────
mem_total_kb=$(awk '/^MemTotal:/{print $2}' /proc/meminfo)
mem_avail_kb=$(awk '/^MemAvailable:/{print $2}' /proc/meminfo)
echo "memory:"
echo "  total_mb: $(( ${mem_total_kb:-0} / 1024 ))"
echo "  available_mb: $(( ${mem_avail_kb:-0} / 1024 ))"
echo "hugepages:"
for sz in 2048kB 1048576kB; do
    d=/sys/kernel/mm/hugepages/hugepages-$sz
    label=$([ "$sz" = 2048kB ] && echo "2m" || echo "1g")
    if [ -d "$d" ]; then
        echo "  $label: { supported: true, total: $(r $d/nr_hugepages), free: $(r $d/free_hugepages) }"
    else
        echo "  $label: { supported: false, total: 0, free: 0 }"
    fi
done
echo "  mounted: $(findmnt -no TARGET -t hugetlbfs 2>/dev/null | tr '\n' ' ' | sed 's/ $//' | sed 's/^$/none/')"

# ── Storage ──────────────────────────────────────────────────────────────────
# Every block device that is not the root disk is a candidate. NVMe is reported
# with its PCI address because an SN needs it bound to vfio-pci; everything else
# is reported by path for the kernel data path.
echo "storage:"
# -P gives quoted key="value" pairs. Positional parsing breaks on models like
# "Virtual Disk", which silently shifts every later field.
lsblk -dn -P -o NAME,SIZE,MODEL,TRAN,TYPE 2>/dev/null | while read -r line; do
    NAME=""; SIZE=""; MODEL=""; TRAN=""; TYPE=""
    eval "$line"
    [ "$TYPE" = "disk" ] || continue
    name=$NAME

    holds_root=false
    [ "$name" = "$root_disk" ] && holds_root=true
    mounted=$(lsblk -no MOUNTPOINTS "/dev/$name" 2>/dev/null | tr -d ' \n' | head -c1)
    has_fs=$(lsblk -no FSTYPE "/dev/$name" 2>/dev/null | tr -d ' \n' | head -c1)

    pci=""; grp=""; drv=""
    if [ -e "/sys/block/$name/device" ]; then
        real=$(readlink -f "/sys/block/$name/device" 2>/dev/null)
        while [ -n "$real" ] && [ "$real" != "/" ]; do
            base=${real##*/}
            case "$base" in
                [0-9a-f][0-9a-f][0-9a-f][0-9a-f]:[0-9a-f][0-9a-f]:[0-9a-f][0-9a-f].[0-9a-f])
                    pci=$base; break ;;
            esac
            real=$(dirname "$real")
        done
        if [ -n "$pci" ]; then grp=$(iommu_group_of "$pci"); drv=$(driver_of "$pci"); fi
    fi

    # An NVMe controller must be handed to SPDK through vfio-pci. Anything else
    # -- SATA, SAS, FC, iSCSI, NVMe-oF -- is used through the kernel.
    case "$name" in
        nvme*) kind=nvme; needs_vfio=true ;;
        *)     kind=$(echo "${TRAN:-unknown}" | tr 'A-Z' 'a-z'); needs_vfio=false ;;
    esac

    echo "  - name: $name"
    echo "    path: /dev/$name"
    echo "    size: \"$SIZE\""
    echo "    model: \"$(echo "${MODEL:-unknown}" | sed 's/"/\\"/g')\""
    echo "    kind: $kind"
    [ -n "$pci" ] && echo "    pci: \"$pci\""
    [ -n "$grp" ] && echo "    iommu_group: $grp"
    [ -n "$drv" ] && echo "    current_driver: $drv"
    [ -n "$pci" ] && echo "    numa_node: $(numa_of "$pci")"
    echo "    holds_root: $holds_root"
    echo "    mounted: $([ -n "$mounted" ] && echo true || echo false)"
    echo "    has_filesystem: $([ -n "$has_fs" ] && echo true || echo false)"

    if [ "$needs_vfio" = true ]; then
        echo "    requires_vfio: true"
        if [ -n "$drv" ] && [ "$drv" != "vfio-pci" ]; then
            echo "    driver_change_required: \"$drv -> vfio-pci\""
        fi
    else
        echo "    requires_vfio: false"
    fi

    if [ "$holds_root" = true ]; then
        echo "    eligible: false"
        echo "    reason: \"holds the root filesystem\""
    elif [ -n "$mounted" ]; then
        echo "    eligible: false"
        echo "    reason: \"currently mounted -- unmount before assigning it\""
    else
        echo "    eligible: true"
        # The Storage Node takes the whole device and writes its own layout over
        # it. Anything already there is gone, and there is no undo.
        if [ -n "$has_fs" ]; then
            echo "    destructive: true"
            echo "    warning: \"DATA LOSS: this device has a filesystem. A Storage Node consumes the whole device and overwrites it.\""
        else
            echo "    destructive: true"
            echo "    warning: \"A Storage Node consumes the whole device and overwrites it.\""
        fi
    fi
done

# ── Network ──────────────────────────────────────────────────────────────────
echo "network:"
for n in /sys/class/net/*; do
    ifn=${n##*/}
    [ "$ifn" = "lo" ] && continue
    [ -e "$n/device" ] || continue          # skip virtual: docker0, veth, tailscale
    pci=$(basename "$(readlink -f "$n/device")")
    case "$pci" in
        [0-9a-f][0-9a-f][0-9a-f][0-9a-f]:*) ;;
        *) continue ;;
    esac
    is_default=false
    [ "$ifn" = "$default_if" ] && is_default=true
    addr=""
    command -v ip >/dev/null 2>&1 && addr=$(ip -o -4 addr show "$ifn" 2>/dev/null | awk '{print $4; exit}')
    echo "  - iface: $ifn"
    echo "    pci: \"$pci\""
    echo "    driver: $(driver_of "$pci")"
    echo "    state: $(r "$n/operstate")"
    echo "    mac: $(r "$n/address")"
    [ -n "$addr" ] && echo "    ipv4: \"$addr\""
    g=$(iommu_group_of "$pci"); [ -n "$g" ] && echo "    iommu_group: $g"
    echo "    numa_node: $(numa_of "$pci")"
    echo "    default_route: $is_default"
    if [ "$is_default" = true ]; then
        echo "    eligible_for_dpdk: false"
        echo "    reason: \"carries the default route — binding it to vfio-pci would cut off this host\""
    else
        echo "    eligible_for_dpdk: true"
    fi
done
