#!/bin/bash
#
# Purpose:
#   Quickly determine which Ansible inventory groups a host belongs to.
#   Accepts a partial hostname match (e.g., "webse" instead of "webserver.college.edu") and resolves it to
#   the first full hostname found in the inventory.
#   Parses the output of `ansible-inventory --graph` to efficiently return
#   group membership, including nested groups, with sub-second runtime.
#
# Usage:
#   ansgroups <hostname_substring>
#
# Example:
#   ansgroups wpdev
#
# Notes:
# - Uses a static inventory file at /home/sysadm/ansible-configuration/hosts
# - Requires ansible-inventory in PATH
# - Returns first match only
#
# 2025-07-10 - Kodiak Firesmith <firesmith@protonmail.com>

inventory_file="~/ansible/hosts"
input="$1"

if [[ -z "$input" || "$input" == "-h" || "$input" == "--help" ]]; then
  echo "Usage: $(basename "$0") <hostname_substring>"
  echo
  echo "Description:"
  echo "  Return the list of Ansible groups for the first inventory host"
  echo "  matching the given substring (e.g., 'wpdev' matches wpdev.whoi.edu)."
  exit 1
fi

# Find the first matching host
host=$(ansible-inventory -i "$inventory_file" --graph | awk -v pat="$input" '
  /[[:space:]]+\|--/ {
    line = gensub(/^.*\|--/, "", "g", $0)
    if (line !~ /^@/ && line ~ pat) {
      print line
      exit
    }
  }')

if [[ -z "$host" ]]; then
  echo "Error: No matching host found in inventory for pattern '$input'" >&2
  exit 2
fi

# Extract groups
groups=$(ansible-inventory -i "$inventory_file" --graph | awk -v host="$host" '
function print_parents(indent) {
  for (i = 0; i <= indent; i++) {
    if (level[i] && !seen[level[i]]++) print level[i]
  }
}
{
  indent = length($0) - length(gensub(/^[[:space:]]*/, "", "g", $0))
  line = gensub(/^.*\|--/, "", "g", $0)

  if (match(line, /^@(.+):$/, m)) {
    level[indent] = m[1]
    for (j in level) if (j > indent) delete level[j]
  } else if (line == host) {
    print_parents(indent)
  }
}' | sort -u)

# Separate 'all' from others
all_group=""
other_groups=()

while IFS= read -r group; do
  if [[ "$group" == "all" ]]; then
    all_group="[all]"
  else
    other_groups+=("$group")
  fi
done <<< "$groups"

IFS=$'\n' sorted_groups=($(sort <<<"${other_groups[*]}"))
unset IFS

# Output
echo "$host"
if [[ -n "$all_group" ]]; then
  echo "  - $all_group"
  for group in "${sorted_groups[@]}"; do
    echo "  - $group"
  done
else
  for group in "${sorted_groups[@]}"; do
    echo "  - $group"
  done
fi