#!/usr/bin/env bash

LAYEROPS_UPDATE_CHECK_SCRIPT_PATH=/usr/local/bin/layerops-image-check-update.sh
LAYEROPS_ROOT_DIR=/data/layerops
LAYEROPS_CONFIG_DIR=${LAYEROPS_ROOT_DIR}/config
LAYEROPS_IMAGES_DIR=${LAYEROPS_ROOT_DIR}/images
LAYEROPS_STORAGE_DIR=${LAYEROPS_ROOT_DIR}/storage
LAYEROPS_STORAGE_NAME=layerops
DOCKER_DATA_PATH=/data/docker
LAYEROPS_DATA_PATH=/data/layerops
WIREGUARD_UDP_PORT=51820
LXC_BRIDGE_DEVICE=lxdbr0
LXC_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 lxc_init() {
  if [[ "$(echo $OSTYPE)" == "darwin"* ]];
  then
    msg "${RED}***Please check REMOTE LXC is configured.***"
    exit 1
  fi

  if ! [ -x "$(command -v lxd)" ]; then
    msg "${RED}***Please install lxd (see: https://linuxcontainers.org/lxd/getting-started-cli/ )"
    exit 1
  fi

  # 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 jq ipcalc-ng lsof || exit 1

  # Init LXD
  DNS_TCP_LISTEN=$(lsof -nP -i4TCP:53 -sTCP:LISTEN | grep 53 | grep -v 127.0) #  | awk '{print $9}'
  if [ "$DNS_TCP_LISTEN" == "" ]; then
    lxd init --minimal
  else
    cat > /tmp/preseed.yaml <<EOF
config: {}
networks:
- config:
    ipv4.address: auto
    ipv6.address: none
    raw.dnsmasq: |-
      port=0
      dhcp-option=option:dns-server,$ALTERNATIVE_DNS_SERVER
  name: $LXC_BRIDGE_DEVICE
  type: bridge
  project: default
storage_pools:
- name: default
  driver: dir
profiles:
- config: {}
  description: Default LXD profile
  devices:
    eth0:
      name: eth0
      network: $LXC_BRIDGE_DEVICE
      type: nic
    root:
      path: /
      pool: default
      type: disk
  name: default
projects:
- config:
    features.images: "true"
    features.networks: "true"
    features.profiles: "true"
    features.storage.volumes: "true"
  description: Default LXD project
  name: default
EOF
    lxd init --preseed < /tmp/preseed.yaml
    rm -f /tmp/preseed.yaml
  fi

  # Create storage pool
  mkdir -p $LAYEROPS_STORAGE_DIR
  lxc storage create $LAYEROPS_STORAGE_NAME dir source=$LAYEROPS_STORAGE_DIR

  echo "Check netwok connectivity"
  echo "==== LXC Networks found: ===="
  lxc network ls

  lxc remote ls |grep -q ubuntu-minimal || lxc remote add ubuntu-minimal https://cloud-images.ubuntu.com/minimal/releases/ --protocol simplestreams --public

  CURL_RESULT=1
  MAX_RETRIES=10
  RETRY_COUNT=0
  lxc launch ubuntu-minimal:24.04 test-curl
  while [[ $CURL_RESULT -ne 0 && $RETRY_COUNT -lt $MAX_RETRIES ]]
  do
    (( RETRY_COUNT++ ))
    sleep 1
    lxc exec test-curl -- curl --connect-timeout 1 -I https://example.org
    CURL_RESULT=$?
  done
  lxc stop -f test-curl
  lxc 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 LXD bridge network (mainly lxdbr0)"
    msg "${RED}(see: https://documentation.ubuntu.com/lxd/en/stable-5.0/howto/network_bridge_firewalld )"
    exit 1
  fi

}

function check_local_image() {
  LAYEROPS_IMAGE_NAME=layerops-client-lxc-${LXC_IMAGE_VERSION}
  LAYEROPS_IMAGE_FILE=${LAYEROPS_IMAGES_DIR}/${LAYEROPS_IMAGE_NAME}.tar.gz
  LAYEROPS_IMAGE_CHECKSUM_FILE=${LAYEROPS_IMAGE_FILE}.sha256
  LXC_IMAGE_CHECKSUM_URL=${LXC_IMAGE_URL}.sha256
  IMAGE_FOUND=false
  for IMAGE in $(lxc image ls -f compact | grep "$LAYEROPS_IMAGE_NAME" | awk '{print $1}')
  do
    if [ "$IMAGE" == "$LAYEROPS_IMAGE_NAME" ] ; then
      IMAGE_FOUND=true
      msg "${GREEN}${LAYEROPS_IMAGE_NAME} Image found"
      break
    fi
  done
  if [ "$IMAGE_FOUND" == "false" ]; then
    msg "${YELLOW}No ${LAYEROPS_IMAGE_NAME} image found"
    rm -f ${LAYEROPS_IMAGE_FILE} ${LAYEROPS_IMAGE_FILE}.sha256
    msg "${YELLOW}Download ${LAYEROPS_IMAGE_NAME} image"
    mkdir -p $LAYEROPS_IMAGES_DIR
    wget -O ${LAYEROPS_IMAGE_FILE} ${LXC_IMAGE_URL}
    wget -O ${LAYEROPS_IMAGE_CHECKSUM_FILE} ${LXC_IMAGE_CHECKSUM_URL}

    msg "${YELLOW}Verify ${LAYEROPS_IMAGE_FILE} checksum"
    if [ ! -f "${LAYEROPS_IMAGE_FILE}" ]; then
      msg "${RED}ERROR: Missing image file ${LAYEROPS_IMAGE_FILE}"
      exit 1
    fi
    EXPECTED_CHECKSUM=$(cat ${LAYEROPS_IMAGE_CHECKSUM_FILE} | awk '{print $1}')
    ACTUAL_CHECKSUM=$(sha256sum ${LAYEROPS_IMAGE_FILE} | awk '{print $1}')
    if [ "$EXPECTED_CHECKSUM" == "$ACTUAL_CHECKSUM" ];then
      msg "${GREEN}Checksum OK."
    else
      msg "${RED}ERROR: Bad checksum for image file ${LAYEROPS_IMAGE_FILE}:"
      msg "${RED}    Expected: ${EXPECTED_CHECKSUM}"
      msg "${RED}    Found: ${ACTUAL_CHECKSUM}"
      exit 1
    fi

    msg "${YELLOW}Import ${LAYEROPS_IMAGE_NAME} image"
    lxc image import ${LAYEROPS_IMAGE_FILE} --alias $LAYEROPS_IMAGE_NAME

    rm -f ${LAYEROPS_IMAGE_FILE} ${LAYEROPS_IMAGE_FILE}.sha256
  fi
}

function lxc_clean() {
  msg "delete"

  LAYEROPS_IMAGES=$(lxc image ls -f compact|grep layerops-client|awk '{print $1}')
  if [ "$LAYEROPS_IMAGES" == "" ]
  then
    msg "${GREEN}No 'layerops-client' Image found: nothing do delete"
  else
    for LAYEROPS_IMAGE_NAME in $LAYEROPS_IMAGES
    do
      msg "${YELLOW} Delete image ${LAYEROPS_IMAGE_NAME}"
      lxc image rm ${LAYEROPS_IMAGE_NAME}
    done
  fi

  if [[ "$(lxc storage ls -f json |jq -r ".[] | select(.name == \"$LAYEROPS_STORAGE_NAME\") | .name")" == "$LAYEROPS_STORAGE_NAME" ]]
  then
    msg "${YELLOW}${LAYEROPS_STORAGE_NAME} storage pool found: delete it"
    lxc storage rm $LAYEROPS_STORAGE_NAME
  else
    msg "${GREEN}No ${LAYEROPS_STORAGE_NAME} storage pool found: nothing do delete"
  fi

  rm -fR $LAYEROPS_CONFIG_DIR $LAYEROPS_IMAGES_DIR $LAYEROPS_STORAGE_DIR
}

function lxc_run () {
  LXC_CONTAINER_ID=${1}

  # Setup user-data
  mkdir -p ${LAYEROPS_CONFIG_DIR}/${LXC_CONTAINER_ID}
  test -f ${LAYEROPS_CONFIG_DIR}/${LXC_CONTAINER_ID}/user-data || \
    cat > ${LAYEROPS_CONFIG_DIR}/${LXC_CONTAINER_ID}/user-data <<EOF
#cloud-config
write_files:
- path: /etc/profile.d/layerops.sh
  content: |
    export INSTANCE_API_URL=${INSTANCE_API_URL}
    export INSTANCE_POOL_UUID=${INSTANCE_POOL_UUID}
    export INSTANCE_UUID=${INSTANCE_UUID}
    export INSTANCE_TOKEN=${INSTANCE_TOKEN}
    export LOCAL_TIMEZONE=${LOCAL_TIMEZONE}
    export LXC_USE_NVIDIA_GPU=${LXC_USE_NVIDIA_GPU:-false}
    export INIT_JOB_NAME=instance_init
  append: false
- path: /usr/local/bin/startup.sh
  permissions: 0755
  content: |
    #!/bin/bash
    . /etc/profile.d/layerops.sh
    envsubst < /etc/layerops/api-access.json.tpl > /etc/layerops/api-access.json
    chmod 600 /etc/layerops/api-access.json
    echo $LOCAL_TIMEZONE > /etc/layerops/timezone
    /etc/layerops/scripts/launch.sh
runcmd:
  - bash /usr/local/bin/startup.sh
EOF
  chmod 600 ${LAYEROPS_CONFIG_DIR}/${LXC_CONTAINER_ID}/user-data
  [ -z "$INIT_JOB_NAME" ] || sed "s-^    export INIT_JOB_NAME=.*-    export INIT_JOB_NAME=${INIT_JOB_NAME}-" -i ${LAYEROPS_CONFIG_DIR}/${LXC_CONTAINER_ID}/user-data

  # Create Container
  LAYEROPS_IMAGE_NAME=layerops-client-lxc-${LXC_IMAGE_VERSION}
  lxc init $LAYEROPS_IMAGE_NAME ${LXC_CONTAINER_ID} -s $LAYEROPS_STORAGE_NAME
  lxc config set ${LXC_CONTAINER_ID} security.nesting=true security.syscalls.intercept.mknod=true security.syscalls.intercept.setxattr=true
  lxc config set ${LXC_CONTAINER_ID} user.user-data - < ${LAYEROPS_CONFIG_DIR}/${LXC_CONTAINER_ID}/user-data

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

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

  # Create storage volumes
  lxc storage volume create ${LAYEROPS_STORAGE_NAME} "${LXC_CONTAINER_ID}-docker"
  lxc config device add ${LXC_CONTAINER_ID} docker disk pool=${LAYEROPS_STORAGE_NAME} source="${LXC_CONTAINER_ID}-docker" path=${DOCKER_DATA_PATH}

  lxc storage volume info ${LAYEROPS_STORAGE_NAME} "${LXC_CONTAINER_ID}-data" > /dev/null 2>&1 \
    || lxc storage volume create ${LAYEROPS_STORAGE_NAME} "${LXC_CONTAINER_ID}-data"
  lxc config device add ${LXC_CONTAINER_ID} layerops disk pool=${LAYEROPS_STORAGE_NAME} source="${LXC_CONTAINER_ID}-data" path=${LAYEROPS_DATA_PATH}

  if [ ! -z "$PROVIDER_UUID" ]
  then
    # Configure Network interface
    LXC_BRIDGE_CIDR=$(lxc network list -f json | jq -r ".[] | select (.name == \"$LXC_BRIDGE_DEVICE\").config.\"ipv4.address\"")
    LXC_CONTAINER_IP="$(ipcalc-ng -n $LXC_BRIDGE_CIDR | cut -d '=' -f 2 | cut -d '.' -f 1,2,3).100"
    lxc config device add ${LXC_CONTAINER_ID} $LXC_CONTAINER_NIC nic nictype=bridged parent=$LXC_BRIDGE_DEVICE ipv4.address=$LXC_CONTAINER_IP

    # Configure NAT proxy for incoming traffic
    [ -z "$HOST_NET_INTERFACE" ] && HOST_NET_INTERFACE=$(ip route list |grep "${HOST_NET_ROUTE_FILTER:-default}"  | awk -F 'dev' '{print $2}' | awk '{print $1}')
    HOST_IP=$(ip -brief address show $HOST_NET_INTERFACE | awk '{print $3}' | cut -d '/' -f 1)
    lxc config device add ${LXC_CONTAINER_ID} wireguard proxy listen=udp:${HOST_IP}:${WIREGUARD_UDP_PORT} connect=udp:${LXC_CONTAINER_IP}:${WIREGUARD_UDP_PORT} nat=true
  fi

  # Start Container
  lxc start ${LXC_CONTAINER_ID}
}

function lxc_remove () {
  lxc stop $1
  lxc rm $1
}

function lxc_remove_volume () {
  lxc storage volume rm $LAYEROPS_STORAGE_NAME "$1"
}

function node_upgrade () {
  msg "${PURPLE}LayerOps:"
  msg "- ${RED}Stop and Delete node $1"
  lxc_remove $1
  lxc_remove_volume "${1}-docker"

  msg "${PURPLE}LayerOps:"
  msg "- ${BLUE}Setup and Deploy node $1"
  check_local_image
  INIT_JOB_NAME=instance_upgrade lxc_run $1
}

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

function node_delete() {
  msg "${PURPLE}LayerOps:"
  msg "- ${RED}Stop and Delete node $1"
  lxc_remove $1
  lxc_remove_volume "${1}-docker"
  lxc_remove_volume "${1}-data"
  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 \
      INSTANCE_API_URL \
      LXC_IMAGE_URL \
      LXC_IMAGE_VERSION \
      INSTANCE_TOKEN
    [ "$LXC_USE_NVIDIA_GPU" == "true" ] && check_nvidia
    lxc_init
    node_create layerops-${INSTANCE_UUID}
    ;;
  stop)
    check_envvars INSTANCE_UUID
    node_delete layerops-${INSTANCE_UUID}
    ;;
  upgrade)
    check_envvars \
      INSTANCE_UUID \
      LXC_IMAGE_URL \
      LXC_IMAGE_VERSION
    node_upgrade layerops-${INSTANCE_UUID}
    ;;
  clean)
    lxc_clean;;
  *)
    usage;;
esac