#!/bin/bash

if [[ ${DEBUG} -gt 0 ]]; then set -x; fi

set -e

if [[ ! -v NUM_NODES ]]; then
    echo "Environment variables for nodes are missing."
    exit 1
fi

COMMAND=$1
SSH_KEY="$HOME/.ssh/id_rsa.pub"
HOSTS_FILE="/etc/hosts"
IP_EXPORT_FILE="./node_vars.env"
NAME_PREFIX="node"


create_nodes() {
  local IP_VARIABLES=""
  local NODE_IFACE=""

  if [[ ! -f $SSH_KEY ]]; then
    echo "No SSH public key found at $SSH_KEY. Please create one with 'ssh-keygen'."
    exit 1
  fi
  
  for ((i = 1; i <= NUM_NODES; i++)); do
    NAME="$NAME_PREFIX$i"
  
    echo "Creating node: $NAME"
    multipass launch --name "$NAME" \
                     --cpus 2 --memory 2G --disk 10G || {
      echo "Failed to create instance $NAME. Exiting."
      exit 1
    }

    echo "Waiting for $NAME to start."
    sleep 5

    echo "Updating PATH for $NAME"
    multipass exec $NAME -- bash -c 'echo "export PATH=\$PATH:." >> ~/.bashrc'

    echo "Installing containerd for $NAME."
    multipass exec $NAME -- sudo apt clean && sudo apt update && sudo apt upgrade -y
    multipass exec $NAME -- sudo apt install -y containerd
    
    echo "Removing $NAME from known hosts."
    ssh-keygen -R $NAME
    echo "Copying SSH public key to $NAME."
    multipass exec "$NAME" -- bash -c "mkdir -p ~/.ssh && chmod 700 ~/.ssh && echo $(cat $SSH_KEY) >> ~/.ssh/authorized_keys && chmod 600 ~/.ssh/authorized_keys"

    IP=$(multipass info "$NAME" --format json | jq -r '.info[] | .ipv4[0]')

    if [[ $i -eq 1 ]]; then
      NODE_IFACE=$(multipass exec "$NAME" -- ip -o a | grep -w "$IP" | awk '{print $2}' | sed 's/@.*$//')
      if [ -z "$NODE_IFACE" ]; then
        echo "Warning: Could not determine the network interface name for $NAME. Using default 'ens3'."
        NODE_IFACE="ens3"
      fi
    fi

    VARIABLE_NAME="NODE${i}_IP"
    IP_VARIABLES+="export ${VARIABLE_NAME}=\"$IP\"\n"

    echo "Updating /etc/hosts for $NAME..."
    if ! grep -q "$IP $NAME" $HOSTS_FILE; then
      echo "$IP $NAME" | sudo tee -a $HOSTS_FILE
    fi

    echo "$NAME setup complete."
  done

  IP_VARIABLES+="export NODE_IFACE=\"$NODE_IFACE\"\n"

  echo -e "$IP_VARIABLES" > "$IP_EXPORT_FILE"
  
  echo "$NUM_NODES nodes created."
  echo "Detected VM Interface: $NODE_IFACE"
  echo "IP addresses exported to: $IP_EXPORT_FILE (e.g., NODE1_IP, NODE2_IP, NODE_IFACE)"
  echo "To load the IPs into your current shell, run: source $IP_EXPORT_FILE"
  echo "To access nodes, use:"
  echo "  ssh ubuntu@node<number>"
  echo "  ssh ubuntu@node1"
}

status_nodes() {

  NODE_LIST=$(multipass list | grep "^${NAME_PREFIX}" || true)
  if [ -z "$NODE_LIST" ]; then
        echo "No nodes found."
  else
        echo "$NODE_LIST"
        if [ -f $IP_EXPORT_FILE ]; then
          echo
          echo "IPs/Interface: 'source $IP_EXPORT_FILE'"
          cat $IP_EXPORT_FILE
        fi
  fi
}

delete_nodes() {

  echo "Delete $NUM_NODES nodes and cleaning up DNS entries."
  for ((i = 1; i <= NUM_NODES; i++)); do
    NAME="node$i"

    echo "Deleting node: $NAME"
    multipass delete "$NAME" 2>/dev/null || echo "Instance $NAME does not exist."
    multipass purge

    echo "Removing $NAME from $HOSTS_FILE."
    sudo sed -i "/$NAME/d" $HOSTS_FILE

  done
  
  if [ -f $IP_EXPORT_FILE ]; then
    rm $IP_EXPORT_FILE
    echo "Removed IP export file: $IP_EXPORT_FILE"
  fi

  echo "Nodes deleted."
}

usage() {
  echo "Usage: $0 <command>"
  echo "Commands:"
  echo "  create    create $NUM_NODES nodes and exports IPs/Interface to $IP_EXPORT_FILE."
  echo "  status    Show the IP address and exported variables of the nodes."  
  echo "  delete    Deletes $NUM_NODES nodes and removes $IP_EXPORT_FILE."
}

if [ -z "$COMMAND" ]; then
    usage
    exit 0
fi

case $COMMAND in
    create)
        create_nodes
        ;;
    status)
        status_nodes
        ;;        
    delete)
        delete_nodes
        ;;
    *)
        echo "Invalid option: $COMMAND"
        usage
        exit 1
        ;;
esac