#!/usr/bin/env bash
export DEBIAN_FRONTEND=noninteractive
export CURL_USER_AGENT=NimeOps-UA

LOCAL_BIN_DIR=/usr/local/bin
LAYEROPS_USER=layerops
LAYEROPS_GROUP=$LAYEROPS_USER
LAYEROPS_HOME_DIR=/opt/layerops
LAYEROPS_ADMIN_USER=layerops-admin
LAYEROPS_ADMIN_GROUP=$LAYEROPS_ADMIN_USER
LAYEROPS_ADMIN_HOME_DIR=/home/layerops-admin
LAYEROPS_DATA_DIR=/data/layerops
LAYEROPS_BIN_DIR=${LAYEROPS_HOME_DIR}/bin
LAYEROPS_WORKER_PATH=${LAYEROPS_BIN_DIR}/worker
LAYEROPS_WORKER_UPDATE_SCRIPT=${LOCAL_BIN_DIR}/layerops-worker-update
LAYEROPS_ETC_DIR=${LAYEROPS_HOME_DIR}/etc
LAYEROPS_ETC_DEFAULT_FILE=${LAYEROPS_ETC_DIR}/default
LAYEROPS_WIREGUARD_INIT_SYSTEMD_FILE=/etc/systemd/system/layerops-wireguard-init.service
LAYEROPS_WORKER_SYSTEMD_FILE=/etc/systemd/system/layerops-worker@.service
LAYEROPS_WORKER_UPDATE_SYSTEMD_FILE=/etc/systemd/system/layerops-worker-update.service
LAYEROPS_WORKER_UPDATE_SYSTEMD_TIMER_FILE=/etc/systemd/system/layerops-worker-update.timer
LAYEROPS_WORKER_VERSION_CHECK=${LAYEROPS_API_URL}/v1/environments/${ENVIRONMENT_GROUP_UUID}/workerVersion
LAYEROPS_WORKER_SIGNATURE_PUBLIC_KEY=${LAYEROPS_API_URL}/v1/environments/${ENVIRONMENT_GROUP_UUID}/workerSignature

SUDO_FILE=/etc/sudoers.d/99-layerops-users

# === Wireguard ===
WIREGUARD_CONFIG_DIR=/etc/wireguard
WIREGUARD_MTU=1420
WIREGUARD_SUBNET=172.24.0.0/13
WIREGUARD_QUICK_RELOAD_PATH=${LAYEROPS_BIN_DIR}/wg-quick-up-reload
WIREGUARD_RELOAD_PATH=${LAYEROPS_BIN_DIR}/wireguard_start_or_reload
LAYEROPS_CHOWN_PATH=${LAYEROPS_BIN_DIR}/layerops-chown
WIREGUARD_ECMP_SYSTEMD_UNIT=/etc/systemd/system/layerops-wireguard-ecmp.service
WIREGUARD_ECMP_SCRIPT=/usr/local/bin/wg-ecmp

# === Docker ===
DOCKER_DATA_DIR=${LAYEROPS_DATA_DIR}/docker
DOCKER_GROUP=docker
DOCKER_DAEMON_FILE=/etc/docker/daemon.json
DOCKER_NETWORK_BASE=172.16.0.0/13
DOCKER_NETWORK_SIZE=26
DOCKER_NETWORK_MTU=1400

# === Incus ===
INCUS_DATA_DIR=${LAYEROPS_DATA_DIR}/incus
INCUS_ADMIN_GROUP=incus-admin
INCUS_NETWORK_MTU=1400

# === ZFS backup script ===
ZFS_BACKUP_SCRIPT=${LAYEROPS_BIN_DIR}/zfs-s3-backup
ZFS_RESTORE_SCRIPT=${LAYEROPS_BIN_DIR}/zfs-s3-restore

# === Stats ===
STATS_PORT=9889
STATS_CACHE_TTL="15s"

# === Node Exporter ===
NODE_EXPORTER_VERSION=1.10.2
NODE_EXPORTER_BIND_PORT=9888
NODE_EXPORTER_SYSTEMD_FILE=/etc/systemd/system/node_exporter.service
NODE_EXPORTER_USERGROUP=node_exporter

# === SPIRE ===
SPIRE_CONFIG_DIR=${LAYEROPS_ETC_DIR}/spire
SPIRE_AGENT_CONFIG_FILE=${SPIRE_CONFIG_DIR}/agent.conf
SPIRE_AGENT_DATA_DIR=${LAYEROPS_DATA_DIR}/spire/agent
SPIRE_VERSION=1.14.0
SPIRE_URL=https://github.com/spiffe/spire/releases/download/v${SPIRE_VERSION}/spire-${SPIRE_VERSION}-linux-amd64-musl.tar.gz
SPIRE_AGENT_PATH=${LAYEROPS_BIN_DIR}/spire-agent
SPIRE_AGENT_LAUNCH_SCRIPT=${LAYEROPS_BIN_DIR}/start-spire-agent
SPIRE_SERVER_IP=172.24.0.1
SPIRE_BIND_PORT=8081
SPIRE_HTTP_CHALLENGE_PORT=1020
SPIRE_TRUST_DOMAIN=${ENVIRONMENT_GROUP_UUID:0:8}.layerops.io
SPIRE_AGENT_SYSTEMD_FILE=/etc/systemd/system/spire-agent.service
SPIRE_TRUST_BUNDLE_URL=${LAYEROPS_ORCHESTRATOR_URL}/environments/spire-trust-bundle?environmentGroupUuid=${ENVIRONMENT_GROUP_UUID}

# =============================================================================
# UTILITY FUNCTIONS
# =============================================================================

# Check that all required environment variables are set
# Requires REQUIRED_ENVVARS to be defined before calling
function check_envvars() {
  MISSING=""
  for envvar in $REQUIRED_ENVVARS
  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 curl_with_retry() {
  local RETRY_COUNT=0
  local RETRY_DELAY=5
  local MAX_RETRY_DELAY=60

  while true; do
    if curl "$@"; then
      return 0
    fi

    ((RETRY_COUNT++))
    echo "curl command failed (attempt #$RETRY_COUNT). Retrying in ${RETRY_DELAY}s..." >&2
    sleep $RETRY_DELAY

    # Exponential backoff with max delay cap
    RETRY_DELAY=$((RETRY_DELAY * 2))
    if [ $RETRY_DELAY -gt $MAX_RETRY_DELAY ]; then
      RETRY_DELAY=$MAX_RETRY_DELAY
    fi
  done
}

function _get_worker_version() {
  REMOTE_VERSION=$(curl_with_retry -s --max-time 10 "$LAYEROPS_WORKER_VERSION_CHECK")
  if [ -z "$REMOTE_VERSION" ]; then
    echo "1.0.0"
    return
  fi

  echo "$REMOTE_VERSION"
}

# Verify worker binary GPG signature
# Arguments: $1 = path to worker binary, $2 = download URL of worker
function verify_worker_signature() {
  local WORKER_PATH="$1"
  local WORKER_URL="$2"

  # Check if LAYEROPS_WORKER_SIGNATURE_PUBLIC_KEY is set
  if [ -z "$LAYEROPS_WORKER_SIGNATURE_PUBLIC_KEY" ]; then
    echo "No GPG public key URL configured, skipping signature verification"
    return 0
  fi

  # Get the public key
  PUBLIC_KEY=$(curl -s --max-time 10 "$LAYEROPS_WORKER_SIGNATURE_PUBLIC_KEY")

  # If no public key returned, skip verification
  if [ -z "$PUBLIC_KEY" ]; then
    echo "No GPG public key returned, skipping signature verification"
    return 0
  fi

  # Try to download the signature file
  SIGNATURE_URL="${WORKER_URL}.asc"
  SIGNATURE_FILE="${WORKER_PATH}.asc"

  HTTP_CODE=$(curl -s -w "%{http_code}" -o "$SIGNATURE_FILE" --max-time 10 "$SIGNATURE_URL")

  # If signature file doesn't exist (404 or other error), skip verification
  if [ "$HTTP_CODE" != "200" ] || [ ! -s "$SIGNATURE_FILE" ]; then
    echo "No signature file found at $SIGNATURE_URL, skipping signature verification"
    rm -f "$SIGNATURE_FILE"
    return 0
  fi

  # Create temporary GPG home directory
  GPG_HOME=$(mktemp -d)
  local GPG_STATUS_FILE=$(mktemp)

  # Import the public key
  if ! echo "$PUBLIC_KEY" | gpg --homedir "$GPG_HOME" --batch --quiet --import 2>/dev/null; then
    echo "Error: Failed to import GPG public key"
    rm -f "$SIGNATURE_FILE" "$GPG_STATUS_FILE"
    rm -rf "$GPG_HOME"
    return 1
  fi

  # Check if the imported key is revoked
  local KEY_STATUS
  KEY_STATUS=$(gpg --homedir "$GPG_HOME" --batch --list-keys --with-colons 2>/dev/null)
  if echo "$KEY_STATUS" | grep -q "^pub:r:"; then
    echo "Error: GPG public key has been revoked"
    rm -f "$SIGNATURE_FILE" "$GPG_STATUS_FILE"
    rm -rf "$GPG_HOME"
    return 1
  fi

  # Verify the signature and capture status output
  local VERIFY_OUTPUT
  VERIFY_OUTPUT=$(gpg --homedir "$GPG_HOME" --batch --status-fd 1 --verify "$SIGNATURE_FILE" "$WORKER_PATH" 2>&1)
  local VERIFY_EXIT_CODE=$?

  # Check for revoked key in signature verification output
  if echo "$VERIFY_OUTPUT" | grep -qE "(REVKEYSIG|KEYREVOKED)"; then
    echo "Error: Signature was made with a revoked key"
    rm -f "$SIGNATURE_FILE" "$GPG_STATUS_FILE"
    rm -rf "$GPG_HOME"
    return 1
  fi

  # Check for expired key
  if echo "$VERIFY_OUTPUT" | grep -qE "(EXPKEYSIG|KEYEXPIRED)"; then
    echo "Error: Signature was made with an expired key"
    rm -f "$SIGNATURE_FILE" "$GPG_STATUS_FILE"
    rm -rf "$GPG_HOME"
    return 1
  fi

  # Check if verification succeeded (GOODSIG must be present)
  if [ $VERIFY_EXIT_CODE -ne 0 ] || ! echo "$VERIFY_OUTPUT" | grep -q "GOODSIG"; then
    echo "Error: GPG signature verification failed!"
    echo "GPG output: $VERIFY_OUTPUT"
    rm -f "$SIGNATURE_FILE" "$GPG_STATUS_FILE"
    rm -rf "$GPG_HOME"
    return 1
  fi

  echo "GPG signature verification successful"
  rm -f "$SIGNATURE_FILE" "$GPG_STATUS_FILE"
  rm -rf "$GPG_HOME"
  return 0
}

# =============================================================================
# INSTALLATION HELPER FUNCTIONS
# =============================================================================

# Wait for apt locks and install packages
# Usage: install_base_packages pkg1 pkg2 pkg3 ...
function install_base_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 "$@" || exit 1
}

# Create the layerops system user
function create_layerops_user() {
  id $LAYEROPS_USER > /dev/null 2>&1 || adduser -q --gecos "" --disabled-password --home $LAYEROPS_HOME_DIR $LAYEROPS_USER
  mkdir -p $LAYEROPS_HOME_DIR/.ssh
  touch $LAYEROPS_HOME_DIR/.ssh/authorized_keys
  [ -z "$ADDITIONAL_SSH_KEYS_URL" ] || curl_with_retry -H "User-Agent: NimeOps-UA" $ADDITIONAL_SSH_KEYS_URL >> $LAYEROPS_HOME_DIR/.ssh/authorized_keys
}

# Create layerops-admin user (for DEV mode)
function create_layerops_admin_user() {
  if [[ "${ENVIRONMENT_MODE:-DEFAULT}" == "DEV" ]]; then
    id $LAYEROPS_ADMIN_USER > /dev/null 2>&1 || adduser -q --gecos "" --disabled-password --home $LAYEROPS_ADMIN_HOME_DIR $LAYEROPS_ADMIN_USER
    mkdir -p $LAYEROPS_ADMIN_HOME_DIR/.ssh
    touch $LAYEROPS_ADMIN_HOME_DIR/.ssh/authorized_keys
    [ -z "$ADDITIONAL_SSH_KEYS_URL" ] || curl_with_retry -H "User-Agent: NimeOps-UA" $ADDITIONAL_SSH_KEYS_URL >> $LAYEROPS_ADMIN_HOME_DIR/.ssh/authorized_keys
    cat >> $SUDO_FILE <<EOF
# Group rules for layerops-admin
$LAYEROPS_ADMIN_USER ALL=(ALL) NOPASSWD: ALL
EOF
  fi
}

# Create mandatory directories
function create_layerops_directories() {
  mkdir -p \
    $WIREGUARD_CONFIG_DIR \
    $LAYEROPS_ETC_DIR \
    $LAYEROPS_BIN_DIR \
    $DOCKER_DATA_DIR \
    $INCUS_DATA_DIR \
    $SPIRE_CONFIG_DIR \
    $SPIRE_AGENT_DATA_DIR \
    $LAYEROPS_DATA_DIR
  touch \
    $LAYEROPS_ETC_DEFAULT_FILE \
    $SPIRE_AGENT_CONFIG_FILE
  chmod 700 $WIREGUARD_CONFIG_DIR $LAYEROPS_ETC_DIR $SPIRE_CONFIG_DIR
  chmod 710 $INCUS_DATA_DIR
  chown root:$INCUS_ADMIN_GROUP $INCUS_DATA_DIR
  chmod 600 $LAYEROPS_ETC_DEFAULT_FILE $SPIRE_AGENT_CONFIG_FILE
  ln -s $INCUS_DATA_DIR /var/lib/incus
}

# Install Docker
function install_docker() {
  mkdir -p /etc/docker
  cat <<EOF > $DOCKER_DAEMON_FILE
{
  "ip-forward-no-drop": true,
  "data-root": "$DOCKER_DATA_DIR",
  "log-driver": "journald",
  "mtu": $DOCKER_NETWORK_MTU,
  "default-network-opts": {
    "bridge": {
      "com.docker.network.driver.mtu": "$DOCKER_NETWORK_MTU"
    }
  },
  "default-address-pools": [
    {"base":"$DOCKER_NETWORK_BASE","size":$DOCKER_NETWORK_SIZE}
  ]
}
EOF

  mkdir -p /etc/apt/keyrings
  curl_with_retry -fsSL https://download.docker.com/linux/debian/gpg -o /etc/apt/keyrings/docker.asc
  chmod a+r /etc/apt/keyrings/docker.asc
  echo \
  "deb [arch=$(dpkg --print-architecture) signed-by=/etc/apt/keyrings/docker.asc] https://download.docker.com/linux/$(. /etc/os-release && echo $ID) \
  $(. /etc/os-release && echo "$VERSION_CODENAME") stable" \
    > /etc/apt/sources.list.d/docker.list
  apt-get update && apt-get install -y --no-install-recommends docker-ce docker-ce-cli
  adduser $LAYEROPS_USER $DOCKER_GROUP
}

# Install Incus
function install_incus() {
  ## Create Incus directory
  mkdir -p $INCUS_DATA_DIR
  echo "INCUS_DIR=$INCUS_DATA_DIR" >> /etc/environment

  ## Install incus from zabbly packages (see. https://github.com/zabbly/incus )
  mkdir -p /etc/apt/keyrings
  curl_with_retry -fsSL https://pkgs.zabbly.com/key.asc -o /etc/apt/keyrings/zabbly.asc
  chmod a+r /etc/apt/keyrings/zabbly.asc
  cat <<EOF > /etc/apt/sources.list.d/zabbly-incus-stable.sources
Enabled: yes
Types: deb
URIs: https://pkgs.zabbly.com/incus/stable
Suites: $(. /etc/os-release && echo ${VERSION_CODENAME})
Components: main
Architectures: $(dpkg --print-architecture)
Signed-By: /etc/apt/keyrings/zabbly.asc

EOF

  apt-get update && apt-get install -y --no-install-recommends incus zfsutils-linux

  ## Init incus server
  cat <<EOF | incus admin init --preseed
config:
  images.auto_update_interval: 15
networks:
- name: incusbr0
  type: bridge
  config:
    ipv4.address: auto
    ipv4.nat: "true"
    ipv6.address: none
    bridge.mtu: $INCUS_NETWORK_MTU
storage_pools:
- name: default
  driver: zfs
  config:
    size: $(df -k $INCUS_DATA_DIR | tail -n1 | awk '{printf("%.f\n", $4 * 0.8 * 1024)}')
profiles:
- name: default
  devices:
    root:
      path: /
      pool: default
      type: disk
    eth0:
      name: eth0
      network: incusbr0
      type: nic
EOF

  ## Setup Network ACLS
  incus network acl create bridge-acl
  incus network set incusbr0 security.acls=bridge-acl
  incus network acl create instance-acl
  incus profile device set default eth0 security.acls=instance-acl
  # Allow egress traffic from incus instances
  incus network acl rule add bridge-acl egress destination=0.0.0.0/0 action=allow
  incus network acl rule add instance-acl egress destination=0.0.0.0/0 action=allow
  # Allow ingress traffic to incus instances only from bridge interface
  INCUS_BRIDGE_IP=$(ip -4 addr show incusbr0  | grep -oP "(?<=inet ).*(?=/)")
  incus network acl rule add instance-acl ingress source=$INCUS_BRIDGE_IP action=allow

  ## Allow layerops user to launch incus instances
  adduser $LAYEROPS_USER $INCUS_ADMIN_GROUP
}

function setup_incus_backup_scripts() {
  cat <<EOF > $ZFS_BACKUP_SCRIPT
#!/bin/bash

NOW=\$(date +%Y-%m-%d_%H%M)
S3_REMOTE=\$1
INSTANCE_UUID=\$1
ZFS_DATASET=default/containers/\$2
SNAPNAME=\${3:-full}

LATEST_FULL_BACKUP=""
if [ "\$SNAPNAME" != "full" ]
then
  # If no full backup found in S3: force a full backup
  LATEST_FULL_BACKUP=\$(rclone cat \${S3_REMOTE}:\${INSTANCE_UUID}/.latest.full)
fi

ARCHIVE=\${INSTANCE_UUID}-\${NOW}.gz
if [ -z "\$LATEST_FULL_BACKUP" ]
then
  ARCHIVE=\${INSTANCE_UUID}-\${NOW}-full.gz
  sudo /sbin/zfs send -c -L \${ZFS_DATASET}@snapshot-\${SNAPNAME} | rclone rcat \${S3_REMOTE}:\${INSTANCE_UUID}/\${ARCHIVE} >&2
  echo "{\"full\": \"\${ARCHIVE}\"}" | rclone rcat \${S3_REMOTE}:\${INSTANCE_UUID}/\${ARCHIVE}.meta.json >&2

  [ "\$SNAPNAME" == "full" ] && echo \$ARCHIVE | rclone rcat \${S3_REMOTE}:\${INSTANCE_UUID}/.latest.full >&2
else
  sudo /sbin/zfs send -c -L -I \${ZFS_DATASET}@snapshot-full \${ZFS_DATASET}@snapshot-\${SNAPNAME} | rclone rcat \${S3_REMOTE}:\${INSTANCE_UUID}/\${ARCHIVE} >&2
  echo "{\"full\": \"\${LATEST_FULL_BACKUP}\", \"diff\": \"\${ARCHIVE}\"}" | rclone rcat \${S3_REMOTE}:\${INSTANCE_UUID}/\${ARCHIVE}.meta.json >&2
fi

echo \$ARCHIVE | rclone rcat \${S3_REMOTE}:\${INSTANCE_UUID}/.latest >&2

echo "{\"sizeBytes\": \$(rclone ls \${S3_REMOTE}:\${INSTANCE_UUID}/\${ARCHIVE} | awk '{print \$1}'), \"archiveName\": \"\${ARCHIVE}\"}"

EOF

  cat <<EOF > $ZFS_RESTORE_SCRIPT
#!/bin/bash

S3_REMOTE=\$1
INSTANCE_UUID=\$1
CONTAINER_NAME=\$2
ZFS_DATASET=default/containers/\$2
ARCHIVE=\$3

if [ -z "\$ARCHIVE" ]
then
  ARCHIVE=\$(rclone cat \${S3_REMOTE}:\${INSTANCE_UUID}/.latest)
fi
# Archive not found: nothing to restore
if [ -z "\$ARCHIVE" ]
then
  echo "No backup archive found for instance \${INSTANCE_UUID}" >&2
  exit 1
fi

META=\$(rclone cat \${S3_REMOTE}:\${INSTANCE_UUID}/\${ARCHIVE}.meta.json)
if [ -z "\$META" ]
then
  echo "Backup metadata not found: \${ARCHIVE}.meta.json" >&2
  exit 1
fi

FULL_BACKUP=\$(echo \$META | jq -r '.full // empty')
if [ -z "\$FULL_BACKUP" ]
then
  echo "Missing 'full' attribute in backup metadata \${ARCHIVE}.meta.json" >&2
  exit 1
fi

DATASET_FOUND=\$(sudo /sbin/zfs list \${ZFS_DATASET} >/dev/null 2>&1 && echo \${ZFS_DATASET} || echo '')
[ -z "\$DATASET_FOUND" ] || sudo /sbin/zfs destroy -r \${ZFS_DATASET}

rclone cat \${S3_REMOTE}:\${INSTANCE_UUID}/\${FULL_BACKUP} | sudo /sbin/zfs recv -Fv \${ZFS_DATASET}

DIFF_BACKUP=\$(echo \$META | jq -r '.diff // empty')
[ -z "\$DIFF_BACKUP" ] || rclone cat \${S3_REMOTE}:\${INSTANCE_UUID}/\${DIFF_BACKUP} | sudo /sbin/zfs recv -Fv \${ZFS_DATASET}

# Clean incus and zfs snapshots to avoid later conflicts...
for SNAP in \$(sudo /sbin/zfs list -t snapshot | grep \${ZFS_DATASET} | awk '{print \$1}')
do
  sudo /sbin/zfs destroy \$SNAP
done
for SNAP in \$(incus snapshot ls -c n -f compact,noheader \$CONTAINER_NAME )
do
  incus snapshot rm \$CONTAINER_NAME \$SNAP
done

EOF

  chmod a+x $ZFS_BACKUP_SCRIPT $ZFS_RESTORE_SCRIPT
}


# Install Node Exporter
function install_node_exporter() {
  id $NODE_EXPORTER_USERGROUP > /dev/null 2>&1 || adduser --system --group --shell /bin/false --no-create-home $NODE_EXPORTER_USERGROUP
  curl_with_retry -fsSL https://github.com/prometheus/node_exporter/releases/download/v${NODE_EXPORTER_VERSION}/node_exporter-${NODE_EXPORTER_VERSION}.linux-amd64.tar.gz \
    | tar -zxvf - -C ${LOCAL_BIN_DIR} --strip-components=1 node_exporter-${NODE_EXPORTER_VERSION}.linux-amd64/node_exporter
  chown $NODE_EXPORTER_USERGROUP:$NODE_EXPORTER_USERGROUP ${LOCAL_BIN_DIR}/node_exporter
  cat <<EOF > $NODE_EXPORTER_SYSTEMD_FILE
[Unit]
Description=Node Exporter
Requires=network-online.target layerops-wireguard-init.service
After=network-online.target layerops-wireguard-init.service

[Service]
User=node_exporter
Group=node_exporter
ExecStart=${LOCAL_BIN_DIR}/node_exporter --web.listen-address=${WIREGUARD_PRIVATE_IP}:${NODE_EXPORTER_BIND_PORT}

[Install]
WantedBy=multi-user.target
EOF
}

# Enable IP forwarding
function enable_ip_forwarding() {
  touch /etc/sysctl.conf
  grep -qxF 'net.ipv4.ip_forward = 1' /etc/sysctl.conf || echo 'net.ipv4.ip_forward = 1' >> /etc/sysctl.conf
  sysctl -p /etc/sysctl.conf
}

function _launch_wireguard_ecmp_service() {
  cat > $WIREGUARD_ECMP_SYSTEMD_UNIT <<EOF
[Unit]
Description=Update wireguard routing rules to use the less-latency gateway
Requires=network-online.target layerops-wireguard-init.service
After=network-online.target layerops-wireguard-init.service

[Service]
Type=simple
EnvironmentFile=-$LAYEROPS_ETC_DEFAULT_FILE
ExecStart=$WIREGUARD_ECMP_SCRIPT

Restart=on-failure
RestartSec=5s

[Install]
WantedBy=multi-user.target
EOF
}

# Initialize Wireguard base configuration
# Sets WIREGUARD_PRIVATE_KEY, WIREGUARD_PUBLIC_KEY, WIREGUARD_PORT as globals
function init_wireguard_base() {
  WIREGUARD_PRIVATE_KEY=$(wg genkey)
  WIREGUARD_PUBLIC_KEY=$(echo "$WIREGUARD_PRIVATE_KEY" | wg pubkey)
  WIREGUARD_PORT=51820

  # Set instance wireguard private ip
  cat <<EOF > $LAYEROPS_WIREGUARD_INIT_SYSTEMD_FILE
[Unit]
Description=Init wireguard private ip
Requires=network-online.target
After=network-online.target

[Service]
Type=oneshot
RemainAfterExit=yes
ExecStart=ip addr add $WIREGUARD_PRIVATE_IP/32 dev lo
ExecStop=ip addr del $WIREGUARD_PRIVATE_IP/32 dev lo

[Install]
WantedBy=multi-user.target
EOF

  cat > $WIREGUARD_QUICK_RELOAD_PATH <<EOF
#!/bin/bash
set -e

WG_INTERFACE=\$1

wg show \$WG_INTERFACE 2> /dev/null && systemctl reload wg-quick@\$WG_INTERFACE || systemctl start wg-quick@\$WG_INTERFACE

# Update routes
ACTUAL_WG_ROUTE_IPS=\$(ip route ls |grep \$WG_INTERFACE | awk '{print \$1}')
UPDATED_WG_ROUTE_IPS=""
for IP in \$(wg show \$WG_INTERFACE allowed-ips)
do
  if [[ "\$IP" == */32 ]]
  then
    ip route replace \$IP proto static scope global src $WIREGUARD_PRIVATE_IP nexthop dev \$WG_INTERFACE
    UPDATED_WG_ROUTE_IPS="\$UPDATED_WG_ROUTE_IPS \${IP%/32}"
  fi
done

# Delete obsolete routes
for IP in \$ACTUAL_WG_ROUTE_IPS
do
  [[ "\$UPDATED_WG_ROUTE_IPS" =~ "\$IP" ]] || ip route del \$IP
done
EOF
  chmod 755 $WIREGUARD_QUICK_RELOAD_PATH

  cat > $WIREGUARD_RELOAD_PATH <<EOF
#!/bin/bash
sudo $WIREGUARD_QUICK_RELOAD_PATH \$1
EOF
  chmod 755 $WIREGUARD_RELOAD_PATH
}

# Create layerops-chown helper script (used by client)
function create_layerops_chown_script() {
  cat > $LAYEROPS_CHOWN_PATH <<EOF
#!/bin/bash
set -e
[ ! -e "\$1" ] && exit 1
cp "\$1" "\$2"
chown "\$3:\$4" "\$2"
EOF
  chmod 755 $LAYEROPS_CHOWN_PATH
}

# Install SPIRE agent binary
function install_spire_agent_binary() {
  cd /tmp
  curl_with_retry -s -N -L $SPIRE_URL | tar xz || exit 1
  mv spire*/bin/spire-agent $SPIRE_AGENT_PATH
  chmod 755 $SPIRE_AGENT_PATH
}

# Configure SPIRE agent (common configuration)
# Argument: $1 = bootstrap settings, $2 = rebootstrap mode settings (optional, can be empty)
function configure_spire_agent() {
  local BOOTSTRAP_SETTINGS="$1"
  local REBOOSTRAP_MODE_SETTINGS="$2"

  cat > $SPIRE_AGENT_CONFIG_FILE <<EOF
agent {
    data_dir = "$SPIRE_AGENT_DATA_DIR"
    log_level = "DEBUG"
    trust_domain = "$SPIRE_TRUST_DOMAIN"
    server_address = "$SPIRE_SERVER_IP"
    server_port = $SPIRE_BIND_PORT
    $BOOTSTRAP_SETTINGS
    $REBOOSTRAP_MODE_SETTINGS
}

plugins {
   KeyManager "disk" {
        plugin_data {
            directory = "$SPIRE_AGENT_DATA_DIR"
        }
    }

    NodeAttestor "http_challenge" {
        plugin_data {
            hostname = "$WIREGUARD_PRIVATE_IP"
            port = $SPIRE_HTTP_CHALLENGE_PORT
        }
    }

    WorkloadAttestor "unix" {
        plugin_data {}
    }
    WorkloadAttestor "docker" {
        plugin_data {}
    }
}
EOF
}

# Create SPIRE agent systemd service
function create_spire_agent_systemd() {
  cat > $SPIRE_AGENT_SYSTEMD_FILE <<EOF
[Unit]
Description=Spire Agent
Requires=network-online.target layerops-wireguard-init.service
After=network-online.target layerops-wireguard-init.service

[Service]
User=$LAYEROPS_USER
Group=$LAYEROPS_GROUP

TimeoutStartSec=0
ExecStartPre=+setcap 'cap_net_bind_service=+ep' $SPIRE_AGENT_PATH
ExecStart=$SPIRE_AGENT_PATH run -config $SPIRE_AGENT_CONFIG_FILE
ExecStop=/bin/kill -9 \$MAINPID

Restart=on-failure
RestartSec=5s

[Install]
WantedBy=multi-user.target
EOF

  cat > $SPIRE_AGENT_LAUNCH_SCRIPT <<EOF
#!/bin/bash
if [ "\$1" == "reset" ]
then
  systemctl stop spire-agent
  rm -f $SPIRE_AGENT_DATA_DIR/*
fi
systemctl enable --now spire-agent
EOF
  chmod 755 $SPIRE_AGENT_LAUNCH_SCRIPT
}

# Download and install layerops worker
# Sets REMOTE_VERSION as global
function install_layerops_worker() {
  REMOTE_VERSION=$(_get_worker_version)
  DOWNLOAD_URL="${LAYEROPS_WORKER_URL}/worker-${REMOTE_VERSION}"
  echo "Downloading worker from $DOWNLOAD_URL..."
  curl_with_retry -s -N -L -o $LAYEROPS_WORKER_PATH $DOWNLOAD_URL
  chmod 755 $LAYEROPS_WORKER_PATH

  # Verify GPG signature if available
  if ! verify_worker_signature "$LAYEROPS_WORKER_PATH" "$DOWNLOAD_URL"; then
    echo "Error: Signature verification failed. Aborting installation."
    rm -f "$LAYEROPS_WORKER_PATH"
    exit 1
  fi
}

# Create worker systemd services
function create_worker_systemd_services() {
  cat > $LAYEROPS_WORKER_SYSTEMD_FILE <<EOF
[Unit]
Description=Launch layerops worker v%i
After=network.target

[Service]
Type=simple
User=$LAYEROPS_USER
Group=$LAYEROPS_GROUP
EnvironmentFile=-$LAYEROPS_ETC_DEFAULT_FILE
ExecStart=$LAYEROPS_WORKER_PATH --version %i

KillMode=mixed
TimeoutStopSec=1800

Restart=on-failure
RestartSec=5s

[Install]
WantedBy=multi-user.target
EOF

  cat > $LAYEROPS_WORKER_UPDATE_SYSTEMD_FILE <<EOF
[Unit]
Description=Check for LayerOps Worker updates

[Service]
Type=oneshot
ExecStart=$LAYEROPS_WORKER_UPDATE_SCRIPT
EnvironmentFile=-$LAYEROPS_ETC_DEFAULT_FILE
EOF

  cat > $LAYEROPS_WORKER_UPDATE_SYSTEMD_TIMER_FILE <<EOF
[Unit]
Description=Run LayerOps updater every 5 minutes

[Timer]
OnBootSec=1min
OnUnitActiveSec=5min

[Install]
WantedBy=timers.target
EOF
}

# Setup SSH user and configuration
# Usage: setup_ssh [--with-docker] [--with-incus]
#   --with-docker: also add the SSH user to the docker group (for client instances)
#   --with-incus: also add the SSH user to the incus-admin group (for system containers)
function setup_ssh() {
  local ADD_TO_DOCKER=false
  local ADD_TO_INCUS=false
  for arg in "$@"; do
    case "$arg" in
      --with-docker) ADD_TO_DOCKER=true ;;
      --with-incus) ADD_TO_INCUS=true ;;
    esac
  done

  # Set CUSTOM_SSH_USER if not already defined
  [ -z "$CUSTOM_SSH_USER" ] && CUSTOM_SSH_USER=dev

  # Create custom ssh user
  id $CUSTOM_SSH_USER > /dev/null 2>&1 || adduser -q --gecos "" --disabled-password --home /home/dev $CUSTOM_SSH_USER

  # Add custom ssh user to docker group if requested (for "exec" command)
  if [ "$ADD_TO_DOCKER" == "true" ]; then
    adduser $CUSTOM_SSH_USER $DOCKER_GROUP
  fi

  # Add custom ssh user to incus-admin group if requested (for system containers)
  if [ "$ADD_TO_INCUS" == "true" ]; then
    adduser $CUSTOM_SSH_USER $INCUS_ADMIN_GROUP
  fi

  # Add specific config to sshd_config for user dev
  SSHD_CONFIG=/etc/ssh/sshd_config

  # Add port 22 if it doesn't already exist (and not commented)
  if ! grep -q "^Port 22" $SSHD_CONFIG; then
    # Add Port 22 before any Match blocks
    if grep -q "^Match" $SSHD_CONFIG; then
      sed -i '/^Match/i Port 22' $SSHD_CONFIG
    else
      echo "Port 22" >> $SSHD_CONFIG
    fi
  fi

  # Add port 2222 if it doesn't already exist (and not commented)
  if ! grep -q "^Port 2222" $SSHD_CONFIG; then
    # Add Port 2222 before any Match blocks
    if grep -q "^Match" $SSHD_CONFIG; then
      sed -i '/^Match/i Port 2222' $SSHD_CONFIG
    else
      echo "Port 2222" >> $SSHD_CONFIG
    fi
  fi

  # Add SSH config if it doesn't already exist
  if ! grep -q "Match User dev" $SSHD_CONFIG; then
    cat >> $SSHD_CONFIG <<EOF

Match User dev
    AuthorizedKeysCommand /bin/cat $LAYEROPS_DATA_DIR/ssh/dev_authorized_keys
    AuthorizedKeysCommandUser root

Match User dev LocalPort 22
    DenyUsers dev
EOF
    # Reload sshd to apply changes
    systemctl reload sshd || systemctl reload ssh
  fi
}

# Set ownership for layerops directories
function set_layerops_ownership() {
  chown -R $LAYEROPS_USER:$LAYEROPS_GROUP $LAYEROPS_HOME_DIR $WIREGUARD_CONFIG_DIR $LAYEROPS_ETC_DIR $SPIRE_CONFIG_DIR $LAYEROPS_DATA_DIR
  # Docker and Incus manage their own file permissions — restore ownership recursively
  chown -R root:root $DOCKER_DATA_DIR
  chown -R root:$INCUS_ADMIN_GROUP $INCUS_DATA_DIR
}

# Common clean function
function clean_layerops_base() {
  # Remove layerops sudo config
  rm -f $SUDO_FILE

  # Delete layerops user
  id $LAYEROPS_USER > /dev/null 2>&1 && userdel $LAYEROPS_USER

  # Uninstall Wireguard
  apt-get remove -y wireguard wireguard-tools

  # Clean files
  rm -fR \
    $LAYEROPS_HOME_DIR \
    $LAYEROPS_WORKER_PATH
    $LAYEROPS_ETC_DIR
}