#!/usr/bin/env bash

# Source common variables and functions
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
source "${SCRIPT_DIR}/layerops-install-common"

# === Client-specific variables ===
VGNAME=layerops
LVNAME=data
REQUIRED_ENVVARS="ENVIRONMENT_GROUP_UUID INSTANCE_UUID INSTANCE_TOKEN LAYEROPS_WORKER_URL LAYEROPS_ORCHESTRATOR_URL SPIRE_URL WIREGUARD_PRIVATE_IP LAYEROPS_API_URL"

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

# Create LVM volumes if available
function _setup_volumes() {
  FS_TYPE=ext4
  FSTAB_OPTIONS=defaults
  FSTAB_FILE=/etc/fstab

  DEVICE=$(lsblk --json | jq -r '[.blockdevices[] | select(.type == "disk") | select(.children == null)]  | first | .name')
  if [[ ! -z "$DEVICE" && "$DEVICE" != "null" ]]
  then
    DISK=/dev/$DEVICE

    pvcreate $DISK
    vgcreate $VGNAME $DISK
    LV_SIZE=$(vgs --reportformat=json --units B $VGNAME | jq -r '.report[].vg[].vg_free' | tr '<' ' ')
    lvcreate -L $LV_SIZE -n $LVNAME $VGNAME

    PARTITION=/dev/${VGNAME}/${LVNAME}
    mkfs.$FS_TYPE -F $PARTITION

    echo "MOUNT $LAYEROPS_DATA_DIR on $PARTITION"
    FSTAB_LINE="$PARTITION $LAYEROPS_DATA_DIR $FS_TYPE $FSTAB_OPTIONS 0 2"
    grep -q "^$PARTITION" $FSTAB_FILE && sed "s=^${PARTITION}.*=$FSTAB_LINE=" -i $FSTAB_FILE || echo "$FSTAB_LINE" >>  $FSTAB_FILE

    mount $LAYEROPS_DATA_DIR
  fi
}

# Compute instance hardware specs
function _compute_instance_specs() {
  export INSTANCE_CPU_CORES=$(cat /proc/cpuinfo | grep -i "^processor" | wc -l)

  INSTANCE_CPU_MHZ_TOTAL=0
  for CORE_MHZ in $(cat /proc/cpuinfo | grep -i "^cpu MHz" | awk -F": " '{print $2}')
  do
    INSTANCE_CPU_MHZ_TOTAL=$(echo $INSTANCE_CPU_MHZ_TOTAL $CORE_MHZ | awk '{print $1 + $2}')
  done
  export INSTANCE_CPU_MHZ=$(echo $INSTANCE_CPU_MHZ_TOTAL $INSTANCE_CPU_CORES | awk '{print int(($1 / $2) * 1000 + 0.5) / 1000}')

  INSTANCE_MEMORY_KB=$(cat /proc/meminfo  |grep MemTotal: | awk '{print $2}')
  export INSTANCE_MEMORY_GB=$(echo $INSTANCE_MEMORY_KB 1048576 | awk '{print int(($1 / $2) * 1000 + 0.5) / 1000}')

  INSTANCE_STORAGE_KB=$(df /opt/ | awk 'NR>1 {print $2}')
  export INSTANCE_STORAGE_GB=$(echo $INSTANCE_STORAGE_KB 1048576 | awk '{print int(($1 / $2) * 1000 + 0.5) / 1000}')
}

# Create client-specific config file
function _create_config_file() {
  cat > $LAYEROPS_ETC_DEFAULT_FILE <<EOF
WIREGUARD_RELOAD_PATH=$WIREGUARD_RELOAD_PATH
WIREGUARD_CLEAN_PATH=$WIREGUARD_CLEAN_PATH
WIREGUARD_PRIVATE_IP=$WIREGUARD_PRIVATE_IP
WIREGUARD_PRIVATE_KEY=$WIREGUARD_PRIVATE_KEY
WIREGUARD_PUBLIC_KEY=$WIREGUARD_PUBLIC_KEY
WIREGUARD_MTU=$WIREGUARD_MTU
WIREGUARD_SUBNET=$WIREGUARD_SUBNET
HOSTFILE_RELOAD_PATH=$HOSTFILE_RELOAD_PATH
SPIRE_AGENT_LAUNCH_SCRIPT=$SPIRE_AGENT_LAUNCH_SCRIPT
SPIRE_TRUST_DOMAIN=$SPIRE_TRUST_DOMAIN
INSTANCE_CPU_CORES=$INSTANCE_CPU_CORES
INSTANCE_CPU_MHZ=$INSTANCE_CPU_MHZ
INSTANCE_MEMORY_GB=$INSTANCE_MEMORY_GB
INSTANCE_STORAGE_GB=$INSTANCE_STORAGE_GB
LAYEROPS_DATA_DIR=$LAYEROPS_DATA_DIR
STATS_PORT=$STATS_PORT
STATS_CACHE_TTL=$STATS_CACHE_TTL
ROLE=instance
API_ORCHESTRATOR=$LAYEROPS_ORCHESTRATOR_URL
LAYEROPS_WORKER_URL=$LAYEROPS_WORKER_URL
LAYEROPS_WORKER_VERSION_CHECK=$LAYEROPS_WORKER_VERSION_CHECK
LAYEROPS_WORKER_SIGNATURE_PUBLIC_KEY=$LAYEROPS_WORKER_SIGNATURE_PUBLIC_KEY
LAYEROPS_MARKETPLACE_REGISTRY=${MARKETPLACE_REGISTRY:-registry.nimeops.net/layerops-public/marketplace}
INCUS_INSTANCE_SOURCE_SERVER=${INCUS_INSTANCE_SOURCE_SERVER:-https://images.linuxcontainers.org}
INCUS_INSTANCE_SOURCE_PROTOCOL=${INCUS_INSTANCE_SOURCE_PROTOCOL:-simplestreams}
ENVIRONMENT_GROUP_UUID=$ENVIRONMENT_GROUP_UUID
PROVIDER_UUID=$PROVIDER_UUID
INSTANCE_UUID=$INSTANCE_UUID
SIGNATURE=$INSTANCE_SIGNATURE
STATS=true
EOF
}

function _init() {
  # Install required packages
  install_base_packages ca-certificates curl coreutils gnupg iptables jq openssh-server rclone rsyslog uuid-runtime wireguard wireguard-tools

  # Set instance signature
  INSTANCE_SIGNATURE=$(uuidgen)
  INSTANCE_INIT_REQUEST=$(cat << EOF
{
  "instanceUuid": "${INSTANCE_UUID}",
  "instanceAccessToken": "${INSTANCE_TOKEN}",
  "signature": "${INSTANCE_SIGNATURE}",
  "environmentGroupUuid": "${ENVIRONMENT_GROUP_UUID}"
}
EOF
)
  curl_with_retry --fail-with-body -A $CURL_USER_AGENT -X POST -H "Content-Type: application/json" -d "$INSTANCE_INIT_REQUEST" ${LAYEROPS_ORCHESTRATOR_URL}/instances/init
  [ "$?" -ne "0" ] && exit 1

  # Create layerops user
  create_layerops_user

  # Create mandatory folders
  create_layerops_directories

  # Install NVIDIA drivers if needed
  if [ "$LAYEROPS_USE_NVIDIA_GPU" == "true" ]
  then
    NVIDIA_INSTALLED=$(nvidia-smi > /dev/null 2>&1 && echo 1 || echo 0)
    if [ $NVIDIA_INSTALLED -eq 0 ]
    then
      apt-get install -y --no-install-recommends ubuntu-drivers-common
      add-apt-repository -y ppa:graphics-drivers/ppa
      apt-get install -y --no-install-recommends $(nvidia-detector)
    fi
  fi

  # Create volumes if available
  _setup_volumes

  # Install Docker
  install_docker

  # Install Node Exporter
  install_node_exporter

  # Set sudo access to layerops user
  cat > $SUDO_FILE <<EOF
# Group rules for layerops
$LAYEROPS_USER ALL=(ALL) NOPASSWD: /usr/bin/wg-quick, /usr/bin/wg, $WIREGUARD_QUICK_RELOAD_PATH, $WIREGUARD_QUICK_CLEAN_PATH
$LAYEROPS_USER ALL=(ALL) NOPASSWD: $HOSTFILE_SUDO_RELOAD_PATH
$LAYEROPS_USER ALL=(ALL) NOPASSWD: $SPIRE_AGENT_LAUNCH_SCRIPT
$LAYEROPS_USER ALL=(ALL) NOPASSWD: $LAYEROPS_CHOWN_PATH
$LAYEROPS_USER ALL=(ALL) NOPASSWD: /sbin/zfs
EOF

  # Create layerops-admin user if in DEV mode
  create_layerops_admin_user

  # Enable ip forwarding (for docker network)
  enable_ip_forwarding

  # Install Incus (@TODO: make it optional, only if user as tagged the host as "MicroVM Host" ?)
  install_incus
  setup_incus_backup_scripts

  # Setup SSH user and configuration (with docker group for exec, incus-admin for system containers)
  setup_ssh --with-docker --with-incus

  # Init Wireguard
  init_wireguard_base
  _launch_wireguard_ecmp_service
  create_layerops_chown_script

  # Install and configure Spire agent
  install_spire_agent_binary
  BOOTSTRAP_SETTINGS="insecure_bootstrap = true"
  REBOOSTRAP_MODE_SETTINGS="rebootstrap_mode = \"auto\""
  [[ $SPIRE_TRUST_BUNDLE_URL == https://* ]] && BOOTSTRAP_SETTINGS="trust_bundle_url = \"$SPIRE_TRUST_BUNDLE_URL\"" || REBOOSTRAP_MODE_SETTINGS=""
  configure_spire_agent "$BOOTSTRAP_SETTINGS" "$REBOOSTRAP_MODE_SETTINGS" "$INSTANCE_UUID"
  create_spire_agent_systemd

  # Install layerops worker
  install_layerops_worker

  # Setup layerops config
  _compute_instance_specs
  _create_config_file

  # Create worker systemd services
  create_worker_systemd_services

  # Set ownership for /opt/layerops directory
  set_layerops_ownership

  # Enable and start layerops worker
  systemctl daemon-reload
  systemctl enable --now \
    layerops-wireguard-init \
    layerops-wireguard-ecmp \
    layerops-worker@${REMOTE_VERSION} \
    layerops-worker-update.timer \
    node_exporter
  systemctl restart docker
}


function _clean() {
  # Stop services
  systemctl disable --now \
    layerops-worker@${REMOTE_VERSION} \
    layerops-worker-update.timer \
    spire-agent \
    node_exporter \
    layerops-wireguard-ecmp \
    layerops-wireguard-init
  rm -f \
    $LAYEROPS_WORKER_SYSTEMD_FILE \
    $LAYEROPS_WIREGUARD_INIT_SYSTEMD_FILE \
    $WIREGUARD_ECMP_SYSTEMD_UNIT \
    $SPIRE_AGENT_SYSTEMD_FILE
  systemctl daemon-reload

  # Clean base layerops installation
  clean_layerops_base
}

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

case "$1" in
  init)
    check_envvars
    _init
    ;;
  clean)
    _clean;;
  *)
    usage;;
esac