#!/usr/bin/env bash
#
# Generate and apply a netplan config for RDMA devices.
#
# Usage:
#   gen-netplan.sh
#
# Script overrides:
#   METADATA_URL   GPU fabric NICs endpoint
#                  (default: http://169.254.169.254/v1/gpu-fabric-nics)
#   CONFIG_FILE    skip metadata; read NIC JSON from this file instead
#   RDMA_DEV_FILE  file containing `rdma dev` output   (default: run `rdma dev`)
#   NETDEV_MAP     file with "<ibdev> <netdev>" lines  (default: sysfs lookup)
#   RENDERER       netplan renderer                    (default: networkd)
#   NETPLAN_OUT    output path                         (default: /etc/netplan/90-rdma.yaml)
#   SKIP_APPLY     if set, skip `netplan apply`
set -euo pipefail
METADATA_URL=${METADATA_URL:-http://169.254.169.254/v1/gpu-fabric-nics}
RENDERER=${RENDERER:-networkd}
MTU=${MTU:-4200}
TABLE_BASE=${TABLE_BASE:-100}       # first VRF table id; +1 per interface
VRF_PREFIX=${VRF_PREFIX:-vrf-}      # VRF name = <prefix><netdev> (keep <=15 chars)
SUMMARY=${SUMMARY:-fd02::/24}       # on-link fabric summary route
NETPLAN_OUT=${NETPLAN_OUT:-/etc/netplan/90-rdma.yaml}
METADATA_RETRIES=${METADATA_RETRIES:-30}
METADATA_RETRY_SECS=${METADATA_RETRY_SECS:-2}
command -v jq >/dev/null 2>&1 || { echo "error: jq is required" >&2; exit 1; }
command -v curl >/dev/null 2>&1 || { echo "error: curl is required" >&2; exit 1; }
ensure_vrf_module() {
  if lsmod | awk '$1=="vrf" {found=1} END{exit !found}'; then
    echo "vrf module already loaded"
    return 0
  fi
}
# Fetch NIC JSON from metadata (or CONFIG_FILE). Exits 1 on a permanent "not found".
# Prints the config path on stdout when successful.
fetch_config() {
  local config http_code body attempt
  config=$(mktemp)
  if [[ -n ${CONFIG_FILE:-} ]]; then
    [[ -f $CONFIG_FILE ]] || { echo "error: config '$CONFIG_FILE' not found" >&2; exit 1; }
    cp "$CONFIG_FILE" "$config"
    printf '%s' "$config"
    return 0
  fi
  for attempt in $(seq 1 "$METADATA_RETRIES"); do
    body=$(mktemp)
    http_code=$(curl -sS -o "$body" -w '%{http_code}' --connect-timeout 2 --max-time 10 \
      "$METADATA_URL" || true)
    case $http_code in
      200)
        if ! jq -e . >/dev/null 2>&1 <"$body"; then
          echo "error: metadata returned invalid JSON" >&2
          rm -f "$body" "$config"
          exit 1
        fi
        mv "$body" "$config"
        printf '%s' "$config"
        return 0
        ;;
      404)
        # Transient while gpu-fabric.sh is still writing, or permanently absent.
        if grep -qi 'not found' "$body" 2>/dev/null || [[ ! -s $body ]]; then
          if [[ $attempt -lt $METADATA_RETRIES ]]; then
            echo "metadata: not found (attempt $attempt/$METADATA_RETRIES); retrying in ${METADATA_RETRY_SECS}s" >&2
            rm -f "$body"
            sleep "$METADATA_RETRY_SECS"
            continue
          fi
          echo "error: GPU fabric NICs not found at $METADATA_URL" >&2
          rm -f "$body" "$config"
          exit 1
        fi
        echo "error: unexpected 404 body from $METADATA_URL: $(cat "$body")" >&2
        rm -f "$body" "$config"
        exit 1
        ;;
      400)
        echo "error: hypervisor is not GPU fabric enabled (HTTP 400)" >&2
        rm -f "$body" "$config"
        exit 1
        ;;
      *)
        if [[ $attempt -lt $METADATA_RETRIES ]]; then
          echo "metadata: HTTP ${http_code:-000} (attempt $attempt/$METADATA_RETRIES); retrying in ${METADATA_RETRY_SECS}s" >&2
          rm -f "$body"
          sleep "$METADATA_RETRY_SECS"
          continue
        fi
        echo "error: failed to fetch $METADATA_URL (HTTP ${http_code:-000})" >&2
        rm -f "$body" "$config"
        exit 1
        ;;
    esac
  done
}
get_rdma_dev() {
  if [[ -n ${RDMA_DEV_FILE:-} ]]; then cat "$RDMA_DEV_FILE"; else rdma dev; fi
}
# last 6 hex digits, lowercase, no separators (portable: no bash4-isms)
last6() {
  local h=${1//:/}
  h=$(printf '%s' "$h" | tr 'A-F' 'a-f')
  printf '%s' "${h:${#h}-6}"
}
# netdev via sysfs (pipe-free; always returns 0, empty if not found)
netdev_via_sysfs() {
  local dev=$1 p d
  p="/sys/class/infiniband/$dev/device/net"
  [[ -d $p ]] || return 0
  for d in "$p"/*; do
    [[ -e $d ]] || break
    basename "$d"
    return 0
  done
}
# netdev via `rdma link` (fallback; parses "... mlx5_X/port ... netdev <name>")
netdev_via_rdmalink() {
  local dev=$1
  rdma link 2>/dev/null | awk -v d="$dev" '
    { split($2, a, "/"); if (a[1]==d) for (i=1;i<=NF;i++) if ($i=="netdev") { print $(i+1); exit } }'
}
# netdev name for an ib device (never aborts under set -e)
netdev_for_dev() {
  local dev=$1 nd=""
  if [[ -n ${NETDEV_MAP:-} ]]; then
    nd=$(awk -v d="$dev" '$1==d {print $2; exit}' "$NETDEV_MAP") || nd=""
  else
    nd=$(netdev_via_sysfs "$dev") || nd=""
    [[ -n $nd ]] || { nd=$(netdev_via_rdmalink "$dev") || nd=""; }
  fi
  printf '%s' "$nd"
}
emit_netplan() {
  local config=$1 out=$2
  local MAP ROWS matched idx
  MAP=$(mktemp)
  ROWS=$(mktemp)
  matched=0
  while read -r _idx dev rest; do
    [[ -n $dev ]] || continue
    dev=${dev%:}
    guid=$(printf '%s' "$rest" | sed -n 's/.*node_guid \([0-9a-fA-F:]*\).*/\1/p') || guid=""
    [[ -n $guid ]] || continue
    nd=$(netdev_for_dev "$dev") || nd=""
    if [[ -z $nd ]]; then
      echo "warning: could not resolve a netdev for $dev; skipping" >&2
      continue
    fi
    printf '%s %s\n' "$(last6 "$guid")" "$nd" >>"$MAP"
    matched=$((matched + 1))
  done < <(get_rdma_dev)
  if [[ $matched -eq 0 ]]; then
    echo "warning: resolved 0 RDMA netdevs (check 'rdma dev' / sysfs / 'rdma link')" >&2
  fi
  idx=0
  while IFS=$'\t' read -r guid address prefix; do
    table=$((TABLE_BASE + idx))
    idx=$((idx + 1))
    key=$(last6 "$guid")
    nd=$(awk -v k="$key" '$1==k {print $2; exit}' "$MAP")
    if [[ -z $nd ]]; then
      echo "warning: no RDMA device matches guid $guid (last6 $key); skipping" >&2
      continue
    fi
    # next hop = leaf side of the /64 (::2)
    via="${address%::*}::2"
    printf '%s\t%s\t%s\t%s\t%s\n' "$nd" "$address" "$prefix" "$table" "$via" >>"$ROWS"
  done < <(jq -r '.nics[] | [.guid, .address, .prefix] | @tsv' "$config")
  {
    printf 'network:\n  version: 2\n  renderer: %s\n' "$RENDERER"
    printf '  vrfs:\n'
    while IFS=$'\t' read -r nd address prefix table via; do
      printf '    %s%s:\n' "$VRF_PREFIX" "$nd"
      printf '      table: %s\n' "$table"
      printf '      interfaces:\n        - %s\n' "$nd"
    done <"$ROWS"
    printf '  ethernets:\n'
    while IFS=$'\t' read -r nd address prefix table via; do
      printf '    %s:\n' "$nd"
      printf '      ignore-carrier: true\n'
      printf '      addresses:\n        - %s%s\n' "$address" "$prefix"
      printf '      routes:\n        - to: %s\n          via: %s\n          on-link: true\n' "$SUMMARY" "$via"
      printf '      mtu: %s\n' "$MTU"
    done <"$ROWS"
  } >"$out"
  rm -f "$MAP" "$ROWS"
  echo "wrote $out"
}
ensure_vrf_module
CONFIG=$(fetch_config)
if [[ -z $CONFIG ]]; then
  echo "error: no GPU fabric NIC configuration to apply" >&2
  exit 1
fi
trap 'rm -f "$CONFIG"' EXIT
emit_netplan "$CONFIG" "$NETPLAN_OUT"
if [[ -z ${SKIP_APPLY:-} ]]; then
  echo "applying netplan"
  netplan apply
fi