#!/usr/bin/env bash
set -euo pipefail
umask 077

[[ "$#" -eq 1 || ("$#" -eq 2 && "$2" = --check) ]] || {
  echo 'usage: apply-wireguard-mtu.sh MTU [--check]' >&2
  exit 2
}
mtu="$1"
[[ "$mtu" =~ ^[0-9]+$ && "$mtu" -ge 1280 && "$mtu" -le 1420 ]] || {
  echo 'invalid WireGuard MTU' >&2
  exit 2
}
check_only=false
[[ "${2:-}" = --check ]] && check_only=true
interface=wg0
config=/etc/wireguard/wg0.conf

fail() {
  printf '%s\n' "$*" >&2
  exit 1
}

if [[ "${WG_AUTO_MTU_TEST:-}" = 1 ]]; then
  config="${WG_AUTO_MTU_CONFIG:-$config}"
else
  [[ "$config" = /etc/wireguard/wg0.conf ]] || fail 'invalid WireGuard config path'
fi
[[ -f "$config" && ! -L "$config" ]] || fail 'invalid WireGuard config path'

validate_config() {
  awk '
    function clean(line) { sub(/[[:space:]]*[#;].*$/, "", line); return line }
    /^[[:space:]]*\[/ {
      section = tolower($0)
      in_interface = (section ~ /^[[:space:]]*\[interface\][[:space:]]*$/)
      if (in_interface) interfaces++
      next
    }
    in_interface {
      line = clean($0)
      lower = tolower(line)
      sub(/[[:space:]]+$/, "", lower)
      if (lower ~ /^[[:space:]]*saveconfig[[:space:]]*=/) {
        sub(/^[[:space:]]*saveconfig[[:space:]]*=[[:space:]]*/, "", lower)
        if (lower == "true") save_config = 1
      }
      if (lower ~ /^[[:space:]]*mtu[[:space:]]*=/) {
        mtus++
        sub(/^[[:space:]]*mtu[[:space:]]*=[[:space:]]*/, "", lower)
        line = lower
        if (line !~ /^[0-9]+[[:space:]]*$/) bad_mtu = 1
      }
    }
    END { exit !(interfaces == 1 && !save_config && mtus <= 1 && !bad_mtu) }
  ' "$config" || fail 'WireGuard config has invalid Interface, SaveConfig, or MTU directives'
}

link_info() {
  ip -o link show dev "$interface" || fail 'WireGuard interface is absent'
}

parse_link() {
  local link="$1"
  ifindex="$(awk -F: '{gsub(/^[[:space:]]+/, "", $1); print $1}' <<<"$link")"
  live_mtu="$(awk '{for (i = 1; i <= NF; i++) if ($i == "mtu") { print $(i + 1); exit }}' <<<"$link")"
  [[ "$ifindex" =~ ^[0-9]+$ && "$live_mtu" =~ ^[0-9]+$ && "$link" == *'<'*'UP'*'>'* ]] || fail 'WireGuard interface is down or malformed'
}

render_config() {
  awk -v mtu="$mtu" '
    tolower($0) ~ /^[[:space:]]*\[interface\][[:space:]]*$/ {
      print
      in_interface = 1
      next
    }
    /^[[:space:]]*\[/ {
      if (in_interface && !replaced) print "MTU = " mtu
      in_interface = 0
      print
      next
    }
    in_interface && $0 ~ /^[[:space:]]*[Mm][Tt][Uu][[:space:]]*=/ {
      print "MTU = " mtu
      replaced = 1
      next
    }
    { print }
    END { if (in_interface && !replaced) print "MTU = " mtu }
  ' "$config"
}

validate_config
link_before="$(link_info)"
parse_link "$link_before"
ifindex_before="$ifindex"
old_mtu="$live_mtu"

if [[ "$check_only" = true ]]; then
  printf '%s\n' "$mtu"
  exit 0
fi

config_dir="$(dirname "$config")"
stage="$(mktemp "$config_dir/.wg0.conf.mtu.XXXXXX")"
cp -a -- "$config" "$stage"
render_config >"$stage.render"
cat "$stage.render" >"$stage"
rm -f -- "$stage.render"

if cmp -s "$config" "$stage" && [[ "$old_mtu" = "$mtu" ]]; then
  rm -f -- "$stage"
  printf '%s\n' "$mtu"
  exit 0
fi

backup_dir="$(mktemp -d /var/tmp/WG-auto-mtu.XXXXXX)"
chmod 0700 "$backup_dir"
backup="$backup_dir/wg0.conf.before"
config_changed=false
mtu_change_attempted=false
cleanup() {
  local status="$1"
  if [[ "$status" -ne 0 ]]; then
    if [[ "$mtu_change_attempted" = true ]]; then ip link set dev "$interface" mtu "$old_mtu" || true; fi
    if [[ "$config_changed" = true ]]; then cp -a -- "$backup" "$config" || true; fi
    printf 'Rollback evidence retained at %s\n' "$backup_dir" >&2
  fi
  rm -f -- "$stage"
  exit "$status"
}
trap 'status=$?; cleanup "$status"' EXIT

cp -a -- "$config" "$backup"
chmod 0600 "$backup"
if ! cmp -s "$config" "$stage"; then
  mv -f -- "$stage" "$config"
  config_changed=true
fi
if [[ "$old_mtu" != "$mtu" ]]; then
  mtu_change_attempted=true
  ip link set dev "$interface" mtu "$mtu"
fi

link_after="$(link_info)"
parse_link "$link_after"
[[ "$ifindex" = "$ifindex_before" && "$live_mtu" = "$mtu" ]] || fail 'WireGuard interface identity, state, or MTU changed unexpectedly'
printf '%s %s\n' "$mtu" "$backup_dir"