#!/bin/bash

# omarchy:summary=Show or configure the system DNS provider
# omarchy:args=[Cloudflare|Google|DHCP|Custom]
# omarchy:examples=omarchy dns | omarchy dns Cloudflare | omarchy dns Custom

set -euo pipefail

provider_from_arg() {
  case "${1:-}" in
  Cloudflare | cloudflare)
    echo "Cloudflare"
    ;;
  Google | google)
    echo "Google"
    ;;
  DHCP | dhcp)
    echo "DHCP"
    ;;
  Custom | custom)
    echo "Custom"
    ;;
  *)
    return 1
    ;;
  esac
}

current_dns_provider() {
  local dns=""
  local compact=""

  dns=$(awk -F= '
    /^[[:space:]]*#/ { next }
    /^[[:space:]]*DNS[[:space:]]*=/ {
      value=$0
      sub(/^[^=]*=/, "", value)
      print value
      exit
    }
  ' /etc/systemd/resolved.conf 2>/dev/null || true)

  compact=$(printf '%s' "$dns" | tr -d '[:space:]')

  if [[ -z $compact ]]; then
    echo "DHCP"
  elif [[ $dns == *"cloudflare-dns.com"* || $dns == *"1.1.1.1"* || $dns == *"2606:4700:4700::1111"* ]]; then
    echo "Cloudflare"
  elif [[ $dns == *"dns.google"* || $dns == *"8.8.8.8"* || $dns == *"2001:4860:4860::8888"* ]]; then
    echo "Google"
  else
    echo "Custom"
  fi
}

self_path() {
  local self="${BASH_SOURCE[0]}"

  if [[ $self != /* ]]; then
    self=$(command -v -- "$self")
  fi

  readlink -f "$self"
}

require_root() {
  if (( EUID == 0 )); then
    return
  fi

  if ! command -v pkexec >/dev/null; then
    echo "Error: pkexec is required to change DNS settings." >&2
    exit 1
  fi

  exec pkexec "$(self_path)" "$@"
}

lock_dns_to_resolved() {
  local file=""

  for file in /etc/systemd/network/*.network; do
    [[ -f $file ]] || continue
    if ! grep -q "^\[DHCPv4\]" "$file"; then continue; fi

    if ! sed -n '/^\[DHCPv4\]/,/^\[/p' "$file" | grep -q "^UseDNS="; then
      sed -i '/^\[DHCPv4\]/a UseDNS=no' "$file"
    fi

    if grep -q "^\[IPv6AcceptRA\]" "$file" && ! sed -n '/^\[IPv6AcceptRA\]/,/^\[/p' "$file" | grep -q "^UseDNS="; then
      sed -i '/^\[IPv6AcceptRA\]/a UseDNS=no' "$file"
    fi
  done
}

unlock_dns_to_dhcp() {
  local file=""

  for file in /etc/systemd/network/*.network; do
    [[ -f $file ]] || continue
    sed -i '/^\[DHCPv4\]/{n;/^UseDNS=no$/d}' "$file"
    sed -i '/^\[IPv6AcceptRA\]/{n;/^UseDNS=no$/d}' "$file"
  done
}

usage() {
  echo "Usage: omarchy-dns [Cloudflare|Google|DHCP|Custom]" >&2
}

if (( $# == 0 )); then
  current_dns_provider
  exit 0
fi

if (( $# > 1 )); then
  usage
  exit 1
fi

if ! provider=$(provider_from_arg "$1"); then
  usage
  exit 1
fi

require_root "$provider"

case "$provider" in
Cloudflare)
  tee /etc/systemd/resolved.conf >/dev/null <<'EOF'
[Resolve]
DNS=1.1.1.1#cloudflare-dns.com 1.0.0.1#cloudflare-dns.com 2606:4700:4700::1111#cloudflare-dns.com 2606:4700:4700::1001#cloudflare-dns.com
FallbackDNS=9.9.9.9#dns.quad9.net 149.112.112.112#dns.quad9.net 2620:fe::fe#dns.quad9.net 2620:fe::9#dns.quad9.net
DNSOverTLS=opportunistic
EOF
  lock_dns_to_resolved
  ;;

Google)
  tee /etc/systemd/resolved.conf >/dev/null <<'EOF'
[Resolve]
DNS=8.8.8.8#dns.google 8.8.4.4#dns.google 2001:4860:4860::8888#dns.google 2001:4860:4860::8844#dns.google
FallbackDNS=9.9.9.9#dns.quad9.net 149.112.112.112#dns.quad9.net 2620:fe::fe#dns.quad9.net 2620:fe::9#dns.quad9.net
DNSOverTLS=opportunistic
EOF
  lock_dns_to_resolved
  ;;

DHCP)
  tee /etc/systemd/resolved.conf >/dev/null <<'EOF'
[Resolve]
DNSOverTLS=no
EOF
  unlock_dns_to_dhcp
  ;;

Custom)
  echo "Enter your DNS servers (space-separated, e.g. '192.168.1.1 1.1.1.1'):"
  if ! read -r dns_servers; then
    dns_servers=""
  fi

  if [[ -z $dns_servers ]]; then
    echo "Error: No DNS servers provided." >&2
    exit 1
  fi

  tee /etc/systemd/resolved.conf >/dev/null <<EOF
[Resolve]
DNS=$dns_servers
FallbackDNS=9.9.9.9#dns.quad9.net 149.112.112.112#dns.quad9.net 2620:fe::fe#dns.quad9.net 2620:fe::9#dns.quad9.net
EOF
  lock_dns_to_resolved
  ;;
esac

systemctl restart systemd-networkd systemd-resolved