#!/usr/bin/env bash
# manage-node.sh
# Usage: manage-node.sh <add|remove> --node NODE [--inventory PATH] [--k3s-server SERVER]
# Options:
#   --no-ansible     : skip running ./ansible.sh
#   --no-k8s         : skip k3s cluster operations (drain/delete / wait-for-ready)
#   --ssh-user USER  : SSH user for remote commands (default: ansible)
#   --ssh-key PATH   : SSH key to use (default: ~/.ssh/id_ed25519_ansible)
#   --ssh-host HOST  : SSH host to use for node-side uninstall (default: same as --node)
#   --vault PASSFILE : ansible vault password file path (default: .vault_pass)
#   --dry-run        : print actions but do not execute destructive steps
#   --force          : allow destructive actions (delete node, rm -rf)
#   --timeout N      : wait timeout in seconds for node readiness (default: 300)
set -euo pipefail

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
REPO_ROOT="$(cd "$SCRIPT_DIR/.." && pwd)"

# Defaults (override with env or flags)
SSH_USER="${SSH_USER:-ansible}"
SSH_KEY="${SSH_KEY:-$HOME/.ssh/id_ed25519_ansible}"
SSH_HOST="${SSH_HOST:-}"
K3S_SERVER="${K3S_SERVER:-myrddin.prole.org}"
ANSIBLE_SCRIPT="${ANSIBLE_SCRIPT:-$REPO_ROOT/ansible.sh}"
VAULT_PASS="${VAULT_PASS:-$REPO_ROOT/.vault_pass}"
DRY_RUN=false
FORCE=false
NO_ANSIBLE=false
NO_K8S=false
TIMEOUT=300

# candidate inventory locations (script will pick first that contains k3s_hosts)
INVENTORIES=(
  "$REPO_ROOT/inventory"
  "$REPO_ROOT/inventory/hosts.ini"
  "$REPO_ROOT/infrastructure/hosts"
  "$REPO_ROOT/infrastructure/ansible/hosts"
  "$REPO_ROOT/infrastructure/ansible/inventory"
  "$REPO_ROOT/infrastructure/ansible/inventory.ini"
  "$REPO_ROOT/hosts.ini"
  "$REPO_ROOT/ansible/inventory"
)

usage() {
  cat <<EOF
Usage: $0 <add|remove> --node NODE [options]

Examples:
  # dry-run removal
  $0 remove --node pi.prole.org --dry-run

  # remove node (do k8s drain + ansible update)
  $0 remove --node pi.prole.org --force

  # add node
  $0 add --node gandalf.prole.org

Options:
  --no-ansible     : skip running ansible
  --no-k8s         : skip k8s drain/delete / wait-for-ready
  --ssh-user USER
  --ssh-key PATH
  --ssh-host HOST  : SSH host for node uninstall (default: --node value)
  --k3s-server HOST : default myrddin.prole.org
  --vault PASSFILE
  --inventory PATH
  --dry-run
  --force
  --timeout N
EOF
  exit 1
}

# parse args
if [ $# -lt 1 ]; then usage; fi
ACTION="$1"; shift
NODE=""
INVENTORY_OVERRIDE=""
while [ $# -gt 0 ]; do
  case "$1" in
    --node) NODE="$2"; shift 2;;
    --ssh-user) SSH_USER="$2"; shift 2;;
    --ssh-key) SSH_KEY="$2"; shift 2;;
    --ssh-host) SSH_HOST="$2"; shift 2;;
    --k3s-server) K3S_SERVER="$2"; shift 2;;
    --inventory) INVENTORY_OVERRIDE="$2"; shift 2;;
    --vault) VAULT_PASS="$2"; shift 2;;
    --dry-run) DRY_RUN=true; shift;;
    --force) FORCE=true; shift;;
    --no-ansible) NO_ANSIBLE=true; shift;;
    --no-k8s) NO_K8S=true; shift;;
    --timeout) TIMEOUT="$2"; shift 2;;
    -h|--help) usage;;
    *) echo "Unknown arg: $1"; usage;;
  esac
done

[ -n "$NODE" ] || { echo "ERROR: --node required"; usage; }
if [ -z "$SSH_HOST" ]; then
  SSH_HOST="$NODE"
fi

echoinfo() { echo "==> $*"; }
echowarn() { echo "WARN: $*" >&2; }
echofatal() { echo "FATAL: $*" >&2; exit 2; }

run_or_echo() {
  if [ "$DRY_RUN" = true ]; then
    echo "[DRY-RUN] $*"
  else
    echo "+ $*"
    eval "$*"
  fi
}

# find inventory file that contains group 'k3s_hosts'
find_inventory() {
  if [ -n "$INVENTORY_OVERRIDE" ]; then
    echo "$INVENTORY_OVERRIDE"
    return 0
  fi
  for f in "${INVENTORIES[@]}"; do
    [ -e "$f" ] || continue
    # check whether it looks like an ini and has [k3s_hosts] or a YAML with k3s_hosts
    if grep -q -E '^\s*\[k3s_hosts\]' "$f" 2>/dev/null || grep -q 'k3s_hosts' "$f" 2>/dev/null; then
      echo "$f"
      return 0
    fi
  done
  # fallback: repo root inventory file
  if [ -e "$REPO_ROOT/inventory" ]; then
    echo "$REPO_ROOT/inventory"
    return 0
  fi
  echowarn "No inventory file found in standard locations. Use --inventory to specify."
  echo ""
  return 1
}

backup_file() {
  local file="$1"
  if [ -e "$file" ]; then
    local ts; ts="$(date -u +"%Y%m%dT%H%M%SZ")"
    local bak="${file}.manage-node.bak.${ts}"
    run_or_echo "cp -a \"$file\" \"$bak\""
    echo "$bak"
  fi
}

# Remove node from INI-style inventory group [k3s_hosts]
remove_node_from_ini_group() {
  local inventory="$1"
  local node="$2"
  local group="k3s_hosts"
  backup_file "$inventory" >/dev/null || true

  # Use awk to remove lines matching the node only when within the k3s_hosts group
  local tmp="${inventory}.tmp.$$"
  awk -v node="$node" -v grp="$group" '
    BEGIN{in_grp=0}
    /^[[:space:]]*\[[^]]+\][[:space:]]*$/ {
      header=$0
      sub(/^[[:space:]]*\[/, "", header)
      sub(/\][[:space:]]*$/, "", header)
      in_grp = (header == grp) ? 1 : 0
      print; next
    }
    {
      if(in_grp){
        line=$0
        sub(/^[[:space:]]+/, "", line)
        sub(/[[:space:]]+$/, "", line)
        if(line == node){ next } # skip exact node line
      }
      print
    }
  ' "$inventory" > "$tmp"
  run_or_echo "mv \"$tmp\" \"$inventory\""
  echoinfo "Removed $node from $inventory ($group)"
}

# Add node to INI-style inventory group [k3s_hosts]
add_node_to_ini_group() {
  local inventory="$1"
  local node="$2"
  local group="k3s_hosts"
  backup_file "$inventory" >/dev/null || true

  # if group exists, append node if not present
  if grep -q -E "^[[:space:]]*\[$group\]" "$inventory"; then
    if grep -q -E "^[[:space:]]*$node[[:space:]]*$" "$inventory"; then
      echoinfo "$node already present in $inventory"
      return 0
    fi
    # insert node after group header or at end of group
    awk -v node="$node" -v grp="$group" '
      BEGIN{in_grp=0; printed=0}
      /^[[:space:]]*\[[^]]+\][[:space:]]*$/ {
        if(in_grp && !printed){ print node; printed=1 } # add before next group
        header=$0
        sub(/^[[:space:]]*\[/, "", header)
        sub(/\][[:space:]]*$/, "", header)
        in_grp = (header == grp) ? 1 : 0
        print; next
      }
      { print }
      END {
        if(!printed && in_grp){ print node; printed=1 } # if file ended in group
        if(!printed && !in_grp){ print ""; print "[" grp "]"; print node } # group not found anywhere
      }
    ' "$inventory" > "${inventory}.tmp.$$"
    run_or_echo "mv \"${inventory}.tmp.$$\" \"$inventory\""
  else
    # group doesn't exist, append group + node
    cat >> "$inventory" <<EOF

[$group]
$node
EOF
    echoinfo "Appended group [$group] and added $node to $inventory"
  fi
}

# run ansible playbook for the provided limit or full k3s_hosts
run_ansible() {
  local limit="$1"
  if [ "$NO_ANSIBLE" = true ]; then
    echoinfo "Skipping ansible as requested (--no-ansible)"
    return 0
  fi
  if [ ! -x "$ANSIBLE_SCRIPT" ]; then
    echowarn "Ansible script $ANSIBLE_SCRIPT not found or not executable; skipping ansible step"
    return 0
  fi
  local cmd
  # Use vault pass file if exists
  if [ -f "$VAULT_PASS" ]; then
    cmd="$ANSIBLE_SCRIPT -v $VAULT_PASS --limit $limit"
  else
    cmd="$ANSIBLE_SCRIPT -v --limit $limit"
  fi
  run_or_echo "$cmd"
}

# k8s: cordon, drain, delete node
drain_and_delete_node() {
  local server="$1"
  local node="$2"
  echoinfo "Will cordon & drain $node via k3s server $server"
  if [ "$NO_K8S" = true ]; then
    echoinfo "Skipping k8s operations (--no-k8s)"
    return 0
  fi

  # Step 1: show pods on node
  echoinfo "Listing pods scheduled on $node"
  run_or_echo "ssh -o ControlMaster=no -i \"$SSH_KEY\" $SSH_USER@$server \"sudo k3s kubectl get pods -A --field-selector spec.nodeName=$node -o wide || true\""

  echoinfo "Cordoning $node"
  run_or_echo "ssh -o ControlMaster=no -i \"$SSH_KEY\" $SSH_USER@$server \"sudo k3s kubectl cordon $node\""

  echoinfo "Draining $node (this may take a while)"
  run_or_echo "ssh -o ControlMaster=no -i \"$SSH_KEY\" $SSH_USER@$server \"if sudo k3s kubectl drain --help 2>/dev/null | grep -q -- '--delete-emptydir-data'; then sudo k3s kubectl drain $node --ignore-daemonsets --delete-emptydir-data --force --grace-period=30 --timeout=10m || true; else sudo k3s kubectl drain $node --ignore-daemonsets --delete-local-data --force --grace-period=30 --timeout=10m || true; fi\""

  echoinfo "Deleting node object from k8s"
  run_or_echo "ssh -o ControlMaster=no -i \"$SSH_KEY\" $SSH_USER@$server \"sudo k3s kubectl delete node $node || true\""
}

# wait for node ready
wait_for_node_ready() {
  local server="$1"
  local node="$2"
  local timeout_secs="$3"
  local start_ts=$(date +%s)
  echoinfo "Waiting up to ${timeout_secs}s for $node to become Ready"
  while true; do
    if ssh -o ControlMaster=no -i "$SSH_KEY" "$SSH_USER@$server" "sudo k3s kubectl get node $node -o jsonpath='{.status.conditions[?(@.type==\"Ready\")].status}' 2>/dev/null" | grep -q "True"; then
      echoinfo "$node is Ready"
      return 0
    fi
    now=$(date +%s)
    if [ $((now - start_ts)) -gt "$timeout_secs" ]; then
      echofatal "Timeout waiting for $node to become Ready"
    fi
    sleep 5
  done
}

# uninstall k3s agent on node
uninstall_k3s_agent_on_node() {
  local ssh_host="$1"
  echoinfo "Attempting to run k3s-agent uninstall on SSH host $ssh_host"
  if [ "$DRY_RUN" = true ]; then
    echo "[DRY-RUN] ssh -i $SSH_KEY $SSH_USER@$ssh_host sudo /usr/local/bin/k3s-agent-uninstall.sh || sudo /usr/local/bin/k3s-uninstall.sh"
    return 0
  fi

  ssh -o ControlMaster=no -i "$SSH_KEY" "$SSH_USER@$ssh_host" "sudo /usr/local/bin/k3s-agent-uninstall.sh || sudo /usr/local/bin/k3s-uninstall.sh || true"
  # cleanup known drop-ins introduced by iscsi role
  ssh -o ControlMaster=no -i "$SSH_KEY" "$SSH_USER@$ssh_host" "sudo rm -f /etc/systemd/system/k3s-agent.service.d/open-iscsi.conf || true; sudo systemctl daemon-reload || true"
}

# MAIN: add/remove
INVENTORY="$(find_inventory || true)"
if [ -z "$INVENTORY" ]; then
  echowarn "No inventory auto-detected. Use --inventory to set the path. Exiting."
  exit 3
fi

case "$ACTION" in
  remove)
    echoinfo "REMOVE node: $NODE"
    # 1) k8s: cordon/drain/delete
    if [ "$NO_K8S" = false ]; then
      drain_and_delete_node "$K3S_SERVER" "$NODE"
    fi

    # 2) run ansible (optional) to remove any node-specific config (we run for remaining k3s_hosts)
    if [ "$NO_ANSIBLE" = false ]; then
      echoinfo "Running ansible on remaining k3s_hosts to reconcile config"
      run_ansible "k3s_hosts"
    fi

    # 3) remove node from inventory
    echoinfo "Removing $NODE from inventory $INVENTORY"
    if grep -q -E "^[[:space:]]*\[k3s_hosts\]" "$INVENTORY" 2>/dev/null; then
      remove_node_from_ini_group "$INVENTORY" "$NODE"
      echoinfo "Inventory updated; commit changes if desired:"
      echoinfo "  git add \"$INVENTORY\" && git commit -m \"Remove $NODE from k3s_hosts\""
    else
      echowarn "k3s_hosts group not found in $INVENTORY - manual edit required"
    fi

    # 4) optionally uninstall k3s agent on node
    if [ "$FORCE" = true ]; then
      echoinfo "FORCE requested: uninstalling k3s agent for $NODE via SSH host $SSH_HOST"
      uninstall_k3s_agent_on_node "$SSH_HOST"
    else
      echoinfo "Skipping k3s-agent uninstall on $NODE (use --force to remove agent)"
    fi

    echoinfo "REMOVE workflow finished. Verify cluster health and inventory changes."
    ;;

  add)
    echoinfo "ADD node: $NODE"
    # 1) Add node to inventory
    echoinfo "Adding $NODE to inventory $INVENTORY"
    add_node_to_ini_group "$INVENTORY" "$NODE"
    echoinfo "Inventory updated; commit changes if desired:"
    echoinfo "  git add \"$INVENTORY\" && git commit -m \"Add $NODE to k3s_hosts\""

    # 2) run ansible to provision node (only that node)
    if [ "$NO_ANSIBLE" = false ]; then
      echoinfo "Provisioning node with ansible (limit: $NODE)"
      run_ansible "$NODE"
    fi

    # 3) wait for k3s agent to register with server
    if [ "$NO_K8S" = false ]; then
      echoinfo "Waiting for node to join the k3s cluster and become Ready"
      wait_for_node_ready "$K3S_SERVER" "$NODE" "$TIMEOUT" || echowarn "Node did not become Ready within timeout"
    fi

    echoinfo "ADD workflow finished. Verify cluster health and inventory changes."
    ;;

  *)
    echofatal "Unknown action: $ACTION"
    ;;
esac