#!/usr/bin/env bash

LAYEROPS_ROOT_DIR=/data/layerops
LAYEROPS_CONFIG_DIR=${LAYEROPS_ROOT_DIR}/config
LAYEROPS_MANAGE_UFW_RULES_FILE=${LAYEROPS_CONFIG_DIR}/manage-ufw-rules
LAYEROPS_STORAGE_NAME=layerops
INCUS_DATA_DIR=/data/layerops/incus
WIREGUARD_UDP_PORT=51820
INCUS_BRIDGE_DEVICE=incusbr0
INCUS_CONTAINER_NIC=eth0
#ALTERNATIVE_DNS_SERVER=9.9.9.9
LAYEROPS_UPDATE_INCUS_CONFIG_SCRIPT_PATH=/usr/local/bin/layerops-lb-incus-config-update
HTTP_PORT=80
HTTPS_PORT=443
SSH_CONTAINERS_PORT=2222

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 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) ufw 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 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 setup_lb_fw_rules() {
  mkdir -p ${LAYEROPS_CONFIG_DIR}
  echo "  Do you want to setup auto firewall rules updates ?"
  echo "  if answer 'yes', this will add/remove automatically ufw rules for added/deleted custom ports"
  echo "  if answer 'no', you have to be sure that your local firewall allow incoming traffic to ports"
  echo "                  80/tcp, 443/tcp, 443/udp, 51820/udp,"
  echo "                  and any custom ports you will add to your layerops environment."
  read -rp " [y/N] " answer </dev/tty
  if [[ "${answer}" =~ ^[Yy]$ ]]; then
    # Init firewall Config
    sed -i 's/IPV6=yes/IPV6=no/' /etc/default/ufw
    echo y | ufw enable
    ufw allow 22/tcp
    ufw route allow out on $INCUS_BRIDGE_DEVICE from any to any proto udp port $WIREGUARD_UDP_PORT
    ufw route allow out on $INCUS_BRIDGE_DEVICE from any to any proto tcp port $HTTP_PORT
    ufw route allow out on $INCUS_BRIDGE_DEVICE from any to any proto tcp port $HTTPS_PORT
    ufw route allow out on $INCUS_BRIDGE_DEVICE from any to any proto udp port $HTTPS_PORT
    ufw allow in on $INCUS_BRIDGE_DEVICE
    ufw route allow in on $INCUS_BRIDGE_DEVICE
    ufw route allow out on $INCUS_BRIDGE_DEVICE from any to any proto tcp port $SSH_CONTAINERS_PORT
    echo 1 > $LAYEROPS_MANAGE_UFW_RULES_FILE
  else
    info "Skipping UFW configuration."
    echo 0 > $LAYEROPS_MANAGE_UFW_RULES_FILE
  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 INSTANCE_UUID=${INSTANCE_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}
  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}/load_balancer/layerops-ssh.sh
  - chmod 755 /usr/local/bin/layerops-ssh
  - curl -o /usr/local/bin/layerops-install ${LAYEROPS_INSTALL_SCRIPT_URL}/load_balancer/layerops-install.sh
  - chmod 755 /usr/local/bin/layerops-install
  - 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
  incus start ${INCUS_CONTAINER_ID}

  # Configure Network interface
  INCUS_CONTAINER_IP=""
  MAX_RETRIES=30
  RETRY_COUNT=0
  while [[ -z "$INCUS_CONTAINER_IP" && $RETRY_COUNT -lt $MAX_RETRIES ]]
  do
    (( RETRY_COUNT++ ))
    sleep 1
    INCUS_CONTAINER_IP=$(incus exec ${INCUS_CONTAINER_ID} -- ip -4 addr show $INCUS_CONTAINER_NIC | grep -oP '(?<=inet\s)\S+' | awk -F '/' '{print $1}')
  done
  if [[ -z "$INCUS_CONTAINER_IP" ]]
  then
    msg "${RED}*** Unable to find LB container's internal IP, cannot setup the load-balancer instance"
    msg "${RED}Check Incus logs for more details."
    exit 1
  fi
  incus config device add ${INCUS_CONTAINER_ID} $INCUS_CONTAINER_NIC nic nictype=bridged parent=$INCUS_BRIDGE_DEVICE ipv4.address=$INCUS_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)
  incus config device add ${INCUS_CONTAINER_ID} wireguard proxy listen=udp:${HOST_IP}:${WIREGUARD_UDP_PORT} connect=udp:${INCUS_CONTAINER_IP}:${WIREGUARD_UDP_PORT} nat=true
  for tcp_port in 80 443 2222
  do
    incus config device add ${INCUS_CONTAINER_ID} tcp_${tcp_port} proxy listen=tcp:${HOST_IP}:${tcp_port} connect=tcp:${INCUS_CONTAINER_IP}:${tcp_port} nat=true
  done

  # Add HTTP/3 (QUIC) support
  incus config device add ${INCUS_CONTAINER_ID} udp_443 proxy listen=udp:${HOST_IP}:443 connect=udp:${INCUS_CONTAINER_IP}:443 nat=true

  # Add update-incus-config script
  cat > $LAYEROPS_UPDATE_INCUS_CONFIG_SCRIPT_PATH <<EOF
#!/bin/bash
set -e

LAYEROPS_MANAGE_UFW_RULES_FILE=${LAYEROPS_MANAGE_UFW_RULES_FILE}
LAYEROPS_API_URL=${LAYEROPS_API_URL}
ENVIRONMENT_GROUP_UUID=${ENVIRONMENT_GROUP_UUID}
INSTANCE_UUID=${INSTANCE_UUID}
INCUS_BRIDGE_DEVICE=${INCUS_BRIDGE_DEVICE}
INCUS_CONTAINER_IP=${INCUS_CONTAINER_IP}
LB_CONTAINER=${INCUS_CONTAINER_ID}
HOST_IP=${HOST_IP}

EOF

  cat >> $LAYEROPS_UPDATE_INCUS_CONFIG_SCRIPT_PATH <<'EOF'
# Check if the container exists
if ! incus info $LB_CONTAINER &> /dev/null; then
  echo "Container $LB_CONTAINER does not exist. Exiting."
  exit 0
fi

echo "Found LB container: $LB_CONTAINER"

# Get the list of custom public ports from the API
echo "Fetching custom public ports from API..."
CUSTOM_PORTS=$(curl -s -A "Nimeops-UA" "$LAYEROPS_API_URL/v1/environments/$ENVIRONMENT_GROUP_UUID/loadBalancers/customPublicPorts")

# Check if the response is valid JSON
if ! echo "$CUSTOM_PORTS" | jq -e . >/dev/null 2>&1; then
  echo "ERROR: API returned invalid JSON response. Exiting for safety."
  exit 1
fi

if [ -z "$CUSTOM_PORTS" ] || [ "$CUSTOM_PORTS" == "null" ] || [ "$CUSTOM_PORTS" == "[]" ]; then
  echo "No custom ports defined or API call failed."
  CUSTOM_PORTS="[]"
  # Initialize an empty API_PORTS variable to avoid processing errors
  API_PORTS=""
else
  # Get the complete list of ports from the API only if we have valid data
  API_PORTS=$(echo "$CUSTOM_PORTS" | jq -c '.[]')
fi

# Parse the current INCUS configuration
echo "Getting current INCUS configuration..."
CURRENT_CONFIG=$(incus config device show $LB_CONTAINER)

# Find all custom port devices
CUSTOM_DEVICES=$(echo "$CURRENT_CONFIG" | grep -o "custom_[a-z0-9_]*" || echo "")

# Process each port from the API
echo "Processing custom ports from API..."
if [ ! -z "$API_PORTS" ]; then
  echo "$API_PORTS" | while read -r PORT_ENTRY; do
    PROTOCOL=$(echo "$PORT_ENTRY" | jq -r '.protocol')
    PORT=$(echo "$PORT_ENTRY" | jq -r '.port')

    ALLOWED_IPS=$(echo "$PORT_ENTRY" | jq '.allowedIpRanges')
    TMPIPFILE=$(mktemp)
    if [ -z "$ALLOWED_IPS" ] || [ "$ALLOWED_IPS" == "null" ] || [ "$ALLOWED_IPS" == "[]" ]; then
      echo "any" > $TMPIPFILE
    else
      echo "$ALLOWED_IPS" | jq -r -c '.[]' > $TMPIPFILE
    fi

    # Update Firewall rules if needed
    if [ "$(cat $LAYEROPS_MANAGE_UFW_RULES_FILE 2> /dev/null || echo 0)" == "1" ]; then
      IFS=$'\n'
      for RULE in $(ufw status | grep "ALLOW FWD" |grep "$PORT/$PROTOCOL")
      do
        IP=$(echo $RULE | awk '{print $NF}')
        [ "$IP" == "Anywhere" ] && IP=any
        grep $IP $TMPIPFILE >/dev/null || ufw route delete allow out on $INCUS_BRIDGE_DEVICE from $IP to any proto $PROTOCOL port $PORT
      done
      for IP in $(cat $TMPIPFILE); do
        ufw route allow out on $INCUS_BRIDGE_DEVICE from $IP to any proto $PROTOCOL port $PORT
      done
      unset IFS
    fi
    rm -f $TMPIPFILE

    # Check if the device already exists
    DEVICE_NAME="custom_${PROTOCOL}_${PORT}"
    if [[ " 80 443 2222 " =~ " $PORT " ]]; then
      echo "Port $PROTOCOL/$PORT already configured"
    elif echo "$CURRENT_CONFIG" | grep -q "$DEVICE_NAME"; then
      echo "Port $PROTOCOL/$PORT already configured"
    else
      echo "Adding new port configuration: $PROTOCOL/$PORT"
      incus config device add $LB_CONTAINER "$DEVICE_NAME" proxy\
        connect=$PROTOCOL:${INCUS_CONTAINER_IP}:$PORT \
        listen=$PROTOCOL:${HOST_IP}:$PORT \
        nat=true
    fi
  done
else
  echo "No ports to add - API returned empty list"
fi

# Check and remove obsolete ports
echo "Checking for obsolete port configurations..."
for DEVICE in $CUSTOM_DEVICES; do
  # Only process devices that start with custom_
  if [[ "$DEVICE" =~ ^custom_([a-z]+)_([0-9]+)$ ]]; then
    PROTOCOL="${BASH_REMATCH[1]}"
    PORT="${BASH_REMATCH[2]}"

    # Check if this port is present in the API list
     if [ -z "$(echo "$CUSTOM_PORTS" | jq ".[] | select(.port == ($PORT|tonumber) and .protocol == \"$PROTOCOL\" )")" ]; then
      echo "Removing obsolete port: $PROTOCOL/$PORT ($device)"
      incus config device remove $LB_CONTAINER "$DEVICE"
      if [ "$(cat $LAYEROPS_MANAGE_UFW_RULES_FILE 2> /dev/null || echo 0)" == "1" ]; then
        IFS=$'\n'
        for RULE in $(ufw status | grep "ALLOW FWD" |grep "$PORT/$PROTOCOL")
        do
          IP=$(echo $RULE | awk '{print $NF}')
          [ "$IP" == "Anywhere" ] && IP=any
          ufw route delete allow out on $INCUS_BRIDGE_DEVICE from $IP to any proto $PROTOCOL port $PORT
        done
        unset IFS
      fi
    fi
  fi
done

echo "INCUS configuration update completed."
EOF
  chmod 755 $LAYEROPS_UPDATE_INCUS_CONFIG_SCRIPT_PATH

  cat > /etc/systemd/system/layerops-lb-incus-config-update.service <<EOF
[Unit]
Description=Update INCUS LB config

[Service]
Type=oneshot
ExecStart=$LAYEROPS_UPDATE_INCUS_CONFIG_SCRIPT_PATH
StandardOutput=journal

[Install]
WantedBy=multi-user.target
EOF

    cat > /etc/systemd/system/layerops-lb-incus-config-update.timer <<EOF
[Unit]
Description=Update INCUS LB config

[Timer]
OnBootSec=1m
OnUnitActiveSec=1m

[Install]
WantedBy=timers.target
EOF

  systemctl daemon-reload
  systemctl enable --now layerops-lb-incus-config-update.timer
}

function incus_remove () {
  systemctl disable --now layerops-lb-incus-config-update.timer
  rm -f /etc/systemd/system/layerops-lb-incus-config-update.*
  systemctl daemon-reload
  rm -f $LAYEROPS_UPDATE_INCUS_CONFIG_SCRIPT_PATH
  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 \
      LAYEROPS_API_URL \
      LAYEROPS_ORCHESTRATOR_URL \
      LAYEROPS_WORKER_URL \
      LAYEROPS_INSTALL_SCRIPT_URL \
      ENVIRONMENT_GROUP_UUID \
      INSTANCE_TOKEN \
      WIREGUARD_PRIVATE_IP
    setup_incus
    setup_lb_fw_rules
    node_create layerops-lb-${INSTANCE_UUID}
    ;;
  stop)
    check_envvars INSTANCE_UUID
    node_delete layerops-lb-${INSTANCE_UUID}
    ;;
  clean)
    incus_clean;;
  *)
    usage;;
esac