#!/usr/bin/env bash

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
LAYEROPS_UPDATE_LXC_CONFIG_SCRIPT_PATH=/usr/local/bin/ly-lxc-config-update
LXC_BRIDGE_DEVICE=lxdbr0
WIREGUARD_UDP_PORT=51820
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 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 ufw 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: 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

  # 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 $LXC_BRIDGE_DEVICE from any to any proto udp port $WIREGUARD_UDP_PORT
  ufw route allow out on $LXC_BRIDGE_DEVICE from any to any proto tcp port $HTTP_PORT
  ufw route allow out on $LXC_BRIDGE_DEVICE from any to any proto tcp port $HTTPS_PORT
  ufw route allow out on $LXC_BRIDGE_DEVICE from any to any proto udp port $HTTPS_PORT
  ufw allow in on $LXC_BRIDGE_DEVICE
  ufw route allow in on $LXC_BRIDGE_DEVICE
  ufw route allow out on $LXC_BRIDGE_DEVICE from any to any proto tcp port $SSH_CONTAINERS_PORT
}

function lxc_clean() {
  msg "delete"

  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 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:
  - wget -O /usr/local/bin/layerops-worker-update ${LAYEROPS_INSTALL_SCRIPT_URL}/worker-update.sh
  - chmod 755 /usr/local/bin/layerops-worker-update
  - wget -O /usr/local/bin/layerops-ssh ${LAYEROPS_INSTALL_SCRIPT_URL}/load_balancer/layerops-ssh.sh
  - chmod 755 /usr/local/bin/layerops-ssh
  - wget -O /usr/local/bin/layerops-install ${LAYEROPS_INSTALL_SCRIPT_URL}/load_balancer/layerops-install.sh
  - chmod 755 /usr/local/bin/layerops-install
runcmd:
  - /usr/local/bin/startup.sh
EOF
  chmod 600 ${LAYEROPS_CONFIG_DIR}/${LXC_CONTAINER_ID}/user-data

  # Create Container
  LAYEROPS_IMAGE_NAME=ubuntu-minimal:
  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

  # 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}

  # 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
  for tcp_port in 80 443 2222
  do
    lxc config device add ${LXC_CONTAINER_ID} tcp_${tcp_port} proxy listen=tcp:${HOST_IP}:${tcp_port} connect=tcp:${LXC_CONTAINER_IP}:${tcp_port} nat=true
  done

  # Add QUIC
  lxc config device add ${LXC_CONTAINER_ID} udp_443 proxy listen=udp:${HOST_IP}:443 connect=udp:${LXC_CONTAINER_IP}:443 nat=true

  # Add update-lxc-config script
  cat > $LAYEROPS_UPDATE_LXC_CONFIG_SCRIPT_PATH-$LXC_CONTAINER_ID <<EOF
#!/bin/bash
set -e

LAYEROPS_API_URL=${LAYEROPS_API_URL}
ENVIRONMENT_GROUP_UUID=${ENVIRONMENT_GROUP_UUID}
INSTANCE_UUID=${INSTANCE_UUID}
LXC_BRIDGE_DEVICE=${LXC_BRIDGE_DEVICE}
LXC_BRIDGE_CIDR=${LXC_BRIDGE_CIDR}
LXC_CONTAINER_IP=${LXC_CONTAINER_IP}
LB_CONTAINER=${LXC_CONTAINER_ID}
HOST_IP=${HOST_IP}
EOF

  cat >> $LAYEROPS_UPDATE_LXC_CONFIG_SCRIPT_PATH-$LXC_CONTAINER_ID <<'EOF'
# Check if the container exists
if ! lxc 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 LXC configuration
echo "Getting current LXC configuration..."
CURRENT_CONFIG=$(lxc 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
    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 $LXC_BRIDGE_DEVICE from $IP to any proto $PROTOCOL port $PORT
    done
    for IP in $(cat $TMPIPFILE); do
      ufw route allow out on $LXC_BRIDGE_DEVICE from $IP to any proto $PROTOCOL port $PORT
    done
    unset IFS
    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"
      lxc config device add $LB_CONTAINER "$DEVICE_NAME" proxy\
        connect=$PROTOCOL:$LXC_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)"
      lxc config device remove $LB_CONTAINER "$DEVICE"
      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 $LXC_BRIDGE_DEVICE from $IP to any proto $PROTOCOL port $PORT
      done
      unset IFS
    fi
  fi
done

echo "LXC configuration update completed."
EOF
  chmod 755 $LAYEROPS_UPDATE_LXC_CONFIG_SCRIPT_PATH-$LXC_CONTAINER_ID

  cat > /etc/systemd/system/ly-update-lxc-config-${LXC_CONTAINER_ID}.service <<EOF
[Unit]
Description=Update LXC config

[Service]
Type=oneshot
ExecStart=$LAYEROPS_UPDATE_LXC_CONFIG_SCRIPT_PATH-$LXC_CONTAINER_ID
StandardOutput=journal

[Install]
WantedBy=multi-user.target
EOF

    cat > /etc/systemd/system/ly-update-lxc-config-${LXC_CONTAINER_ID}.timer <<EOF
[Unit]
Description=Update LXC config

[Timer]
OnBootSec=1m
OnUnitActiveSec=1m

[Install]
WantedBy=timers.target
EOF

  systemctl daemon-reload
  systemctl enable --now ly-update-lxc-config-${LXC_CONTAINER_ID}.timer

  # Start Container
  lxc start ${LXC_CONTAINER_ID}
}

function lxc_remove () {
  systemctl disable --now ly-update-lxc-config-$1.timer
  rm -f /etc/systemd/system/ly-update-lxc-config-$1.*
  systemctl daemon-reload
  rm -f $LAYEROPS_UPDATE_LXC_CONFIG_SCRIPT_PATH-$1
  lxc stop $1
  lxc rm $1
}

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

function node_create() {
  msg "${PURPLE}LayerOps:"
  msg "- ${BLUE}Setup and Deploy node $1"
  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 \
      LAYEROPS_API_URL \
      LAYEROPS_ORCHESTRATOR_URL \
      LAYEROPS_WORKER_URL \
      LAYEROPS_INSTALL_SCRIPT_URL \
      ENVIRONMENT_GROUP_UUID \
      INSTANCE_TOKEN \
      WIREGUARD_PRIVATE_IP
    lxc_init
    node_create layerops-lb-${INSTANCE_UUID}
    ;;
  stop)
    check_envvars INSTANCE_UUID
    node_delete layerops-lb-${INSTANCE_UUID}
    ;;
  clean)
    lxc_clean;;
  *)
    usage;;
esac