#!/usr/bin/env bash

LAYEROPS_ROOT_DIR=/data/layerops
LAYEROPS_CONFIG_DIR=${LAYEROPS_ROOT_DIR}/config
LAYEROPS_STORAGE_NAME=layerops
INCUS_DATA_DIR=/data/layerops/incus
INCUS_BRIDGE_DEVICE=incusbr0
INCUS_CONTAINER_NIC=eth0
# ALTERNATIVE_DNS_SERVER=9.9.9.9

function usage {
  cat << EOF
Usage: $(basename "${BASH_SOURCE[0]}") [start|stop|clean]
EOF
  exit 0
}

function check_envvars() {
  MISSING=""
  for envvar in $@
  do
    [ -z "${!envvar}" ] && MISSING="$envvar $MISSING"
  done

  if [ ! -z "$MISSING" ]
  then
    echo "Missing following environment variables:"
    for var in $MISSING
    do
      echo "  $var"
    done
    exit 1
  fi
}

function check_nvidia() {
  echo "Check NVIDIA Drivers installation"
  nvidia-smi -L
  if [ $? -ne 0 ]
  then
    msg "${RED}***Cannot instantiate GPU instance: NVIDIA drivers are missing."
    msg "${RED}Please be sure your host is GPU compatible, and run following commands:"
    msg "${RED}  \$ apt install ubuntu-drivers-common"
    msg "${RED}  \$ add-apt-repository ppa:graphics-drivers/ppa"
    msg "${RED}  \$ apt install \$(nvidia-detector)"
    exit 1
  fi
}

function setup_incus() {
  echo ""
  echo -e "${BOLD}━━━ Requirements setup ━━━${NC}"
  echo ""

  # Install required packages
  while fuser /var/lib/dpkg/lock-frontend >/dev/null 2>&1; do sleep 1; done;
  while fuser /var/lib/apt/lists/lock >/dev/null 2>&1; do sleep 1; done;
  apt-get update && apt-get install -y --no-install-recommends jq ipcalc-ng lsof linux-modules-$(uname -r) zfsutils-linux || exit 1

  # Workaround allow zfs delegate to incus container:
  chmod 0666 /dev/zfs

  echo ""
  echo -e "${BOLD}━━━ Incus setup ━━━${NC}"
  echo ""

  ## Create Incus directory
  mkdir -p $INCUS_DATA_DIR
  [ -d "/var/lib/incus" ] || ln -s $INCUS_DATA_DIR /var/lib/incus
  echo "INCUS_DIR=$INCUS_DATA_DIR" >> /etc/environment

  if command -v incus >/dev/null 2>&1; then
    local incus_version
    incus_version="$(incus version 2>/dev/null || echo "unknown")"
    echo "Incus is already installed (${incus_version})."
    echo ""
    read -rp "  Do you want to (re)initialize Incus with the default preseed? [y/N] " answer </dev/tty
    if [[ "${answer}" =~ ^[Yy]$ ]]; then
      init_incus
    else
      info "Skipping Incus initialization."
    fi
  else
    echo "Incus is not installed."
    echo ""
    echo "  Incus v6.18+ is required. Install it via zabbly:"
    echo "    https://github.com/zabbly/incus"
    echo ""
    read -rp "  Is Incus now installed? Press Enter to continue, or Ctrl+C to abort..." _ </dev/tty

    if ! command -v incus >/dev/null 2>&1; then
      echo "Incus is still not found. Install it and re-run this script."
    fi

    init_incus
  fi
}

function init_incus() {
  # Setup ZFS storage pool
  STORAGE_CONFIG='{}'
  if [ -z "$INCUS_ZFS_STORAGE_SOURCE" ]; then
    size=$(df -k $INCUS_DATA_DIR | tail -n1 | awk '{printf("%.f\n", $4 * 0.8 * 1024)}')
    STORAGE_CONFIG="{size: ${size}, volume.zfs.delegate: true}"
  else
    STORAGE_CONFIG="{source: ${INCUS_ZFS_STORAGE_SOURCE}, volume.zfs.delegate: true}"
  fi

  # Init Incus
  cat <<PRESEEDEOF | incus admin init --preseed
config:
  images.auto_update_interval: 15
networks:
- name: $INCUS_BRIDGE_DEVICE
  type: bridge
  config:
    ipv4.address: auto
    ipv4.nat: "true"
    ipv6.address: none
storage_pools:
- name: $LAYEROPS_STORAGE_NAME
  driver: zfs
  config: $STORAGE_CONFIG
profiles:
- name: default
  devices:
    root:
      path: /
      pool: $LAYEROPS_STORAGE_NAME
      type: disk
    $INCUS_CONTAINER_NIC:
      name: $INCUS_CONTAINER_NIC
      network: $INCUS_BRIDGE_DEVICE
      type: nic
PRESEEDEOF

  # Check netwok connectivity
  echo "Check netwok connectivity"
  echo "==== INCUS Networks found: ===="
  incus network ls

  CURL_RESULT=1
  MAX_RETRIES=10
  RETRY_COUNT=0
  incus launch images:centos/10-Stream/cloud test-curl
  while [[ $CURL_RESULT -ne 0 && $RETRY_COUNT -lt $MAX_RETRIES ]]
  do
    (( RETRY_COUNT++ ))
    sleep 1
    incus exec test-curl -- curl --connect-timeout 1 -k -I https://example.org
    CURL_RESULT=$?
  done
  incus stop -f test-curl
  incus rm test-curl
  if [[ $CURL_RESULT -ne 0 ]]
  then
    msg "${RED}*** Outgoing traffic is broken."
    msg "${RED}Please check your firewall rules to allow your forward rules through the Incus bridge network (mainly $INCUS_BRIDGE_DEVICE)"
    msg "${RED}(see: https://linuxcontainers.org/incus/docs/main/howto/network_bridge_firewalld/ )"
    exit 1
  fi
}

function incus_clean() {
  msg "delete"

  rm -fR $LAYEROPS_CONFIG_DIR
}

function incus_run () {
  INCUS_CONTAINER_ID=${1}

  # Setup user-data
  mkdir -p ${LAYEROPS_CONFIG_DIR}/${INCUS_CONTAINER_ID}
  test -f ${LAYEROPS_CONFIG_DIR}/${INCUS_CONTAINER_ID}/user-data || \
    cat > ${LAYEROPS_CONFIG_DIR}/${INCUS_CONTAINER_ID}/user-data <<EOF
#cloud-config
write_files:
- path: /etc/profile.d/layerops.sh
  content: |
    export LAYEROPS_API_URL=${LAYEROPS_API_URL}
    export LAYEROPS_ORCHESTRATOR_URL=${LAYEROPS_ORCHESTRATOR_URL}
    export LAYEROPS_WORKER_URL=${LAYEROPS_WORKER_URL}
    export LAYEROPS_INSTALL_SCRIPT_URL=${LAYEROPS_INSTALL_SCRIPT_URL}
    export LAYEROPS_USE_NVIDIA_GPU=${LAYEROPS_USE_NVIDIA_GPU:-false}
    export LAYEROPS_USE_NFS=${LAYEROPS_USE_NFS:-false}
    export INSTANCE_UUID=${INSTANCE_UUID}
    export INSTANCE_POOL_UUID=${INSTANCE_POOL_UUID}
    export ENVIRONMENT_GROUP_UUID=${ENVIRONMENT_GROUP_UUID}
    export INSTANCE_TOKEN=${INSTANCE_TOKEN}
    export WIREGUARD_PRIVATE_IP=${WIREGUARD_PRIVATE_IP}
    export LOCAL_TIMEZONE=${LOCAL_TIMEZONE:-Europe/Paris}
    export INCUS_ZFS_STORAGE_SOURCE=${INCUS_ZFS_STORAGE_SOURCE:-$LAYEROPS_STORAGE_NAME}/containers/${INCUS_CONTAINER_ID}/${LAYEROPS_STORAGE_NAME}
  append: false
- path: /usr/local/bin/startup.sh
  permissions: 0755
  content: |
    #!/bin/bash
    . /etc/profile.d/layerops.sh
    /usr/local/bin/layerops-install init
bootcmd:
  - apt-get update
  - apt-get install -y curl
  - curl -o /usr/local/bin/layerops-worker-update ${LAYEROPS_INSTALL_SCRIPT_URL}/worker-update.sh
  - chmod 755 /usr/local/bin/layerops-worker-update
  - curl -o /usr/local/bin/layerops-ssh ${LAYEROPS_INSTALL_SCRIPT_URL}/client/layerops-ssh.sh
  - chmod 755 /usr/local/bin/layerops-ssh
  - curl -o /usr/local/bin/layerops-install ${LAYEROPS_INSTALL_SCRIPT_URL}/client/layerops-install.sh
  - chmod 755 /usr/local/bin/layerops-install
  - curl -o /usr/local/bin/wg-ecmp ${LAYEROPS_INSTALL_SCRIPT_URL}/wg-ecmp
  - chmod 755 /usr/local/bin/wg-ecmp
  - curl -o /usr/local/bin/layerops-install-common ${LAYEROPS_INSTALL_SCRIPT_URL}/layerops-install-common.sh
  - chmod 755 /usr/local/bin/layerops-install-common
runcmd:
  - /usr/local/bin/startup.sh
EOF
  chmod 600 ${LAYEROPS_CONFIG_DIR}/${INCUS_CONTAINER_ID}/user-data

  # Create Container
  LAYEROPS_IMAGE_NAME=images:ubuntu/noble/cloud
  incus init $LAYEROPS_IMAGE_NAME ${INCUS_CONTAINER_ID}
  incus config set ${INCUS_CONTAINER_ID} security.nesting=true security.syscalls.intercept.mknod=true security.syscalls.intercept.setxattr=true
  incus config set ${INCUS_CONTAINER_ID} user.user-data - < ${LAYEROPS_CONFIG_DIR}/${INCUS_CONTAINER_ID}/user-data

  # Enable Privileged mode if NFS shared Volume is set:
  [ "$LAYEROPS_USE_NFS" == "true" ] && incus config set ${INCUS_CONTAINER_ID} security.privileged=true

  # Enable nvidia runtime if host is GPU enabled
  if [ "$LAYEROPS_USE_NVIDIA_GPU" == "true" ]
  then
    incus config set ${INCUS_CONTAINER_ID} nvidia.runtime=true
    incus config device add ${INCUS_CONTAINER_ID} gpu gpu
  fi

  # Start Container
  incus start ${INCUS_CONTAINER_ID}
}

function incus_remove () {
  incus stop $1
  incus rm $1
}

function node_create() {
  msg "${PURPLE}LayerOps:"
  msg "- ${BLUE}Setup and Deploy node $1"
  incus_run $1
}

function node_delete() {
  msg "${PURPLE}LayerOps:"
  msg "- ${RED}Stop and Delete node $1"
  incus_remove $1
  rm -fR ${LAYEROPS_CONFIG_DIR}/${1}
}

if [[ "$#" -eq 0 ]]; then
  usage
fi

setup_colors() {
  if [[ -t 2 ]] && [[ -z "${NO_COLOR-}" ]] && [[ "${TERM-}" != "dumb" ]]; then
    NOFORMAT='\033[0m' RED='\033[0;31m' GREEN='\033[0;32m' ORANGE='\033[0;33m' BLUE='\033[0;34m' PURPLE='\033[0;35m' CYAN='\033[0;36m' YELLOW='\033[1;33m'
  else
    NOFORMAT='' RED='' GREEN='' ORANGE='' BLUE='' PURPLE='' CYAN='' YELLOW=''
  fi
}

msg() {
  echo >&2 -e "${1-}${NOFORMAT}"
}

setup_colors

case "$1" in
  start)
    check_envvars \
      INSTANCE_UUID \
      INSTANCE_POOL_UUID \
      LAYEROPS_API_URL \
      LAYEROPS_ORCHESTRATOR_URL \
      LAYEROPS_WORKER_URL \
      LAYEROPS_INSTALL_SCRIPT_URL \
      ENVIRONMENT_GROUP_UUID \
      INSTANCE_TOKEN \
      WIREGUARD_PRIVATE_IP
    [ "$INCU_USE_NVIDIA_GPU" == "true" ] && check_nvidia
    setup_incus
    node_create layerops-${INSTANCE_UUID}
    ;;
  stop)
    check_envvars INSTANCE_UUID
    node_delete layerops-${INSTANCE_UUID}
    ;;
  clean)
    incus_clean;;
  *)
    usage;;
esac