#!/usr/bin/env bash

set -euo pipefail

readonly SCRIPT_NAME="dns-override.sh"
readonly SCRIPT_VERSION="1.0.0"

readonly NC='\033[0m'
readonly RED='\033[1;31m'
readonly GREEN='\033[1;32m'
readonly YELLOW='\033[1;33m'
readonly BLUE='\033[1;34m'
readonly GREY='\033[0;37m'

readonly EMOJI_INFO='ℹ️'
readonly EMOJI_OK='✅'
readonly EMOJI_WARN='⚠️'
readonly EMOJI_ERROR='❌'
readonly EMOJI_GEAR='🛠️'
readonly EMOJI_ROCKET='🚀'
readonly EMOJI_NET='🌐'
readonly EMOJI_LOCK='🔒'

SUDO_BIN=""
declare -a DNS_SERVERS=()
declare -a DETECTED_MANAGERS=()
declare -a ACTIONS_ATTEMPTED=()
declare -a ACTIONS_SUCCEEDED=()
declare -a ACTIONS_FAILED=()
declare -a ACTIONS_SKIPPED=()

log_info() { printf "%b%s %s%b\n" "$GREY" "$EMOJI_INFO" "$1" "$NC"; }
log_success() { printf "%b%s %s%b\n" "$GREEN" "$EMOJI_OK" "$1" "$NC"; }
log_warning() { printf "%b%s %s%b\n" "$YELLOW" "$EMOJI_WARN" "$1" "$NC"; }
log_error() { printf "%b%s %s%b\n" "$RED" "$EMOJI_ERROR" "$1" "$NC" >&2; }
log_step() { printf "%b%s %s%b\n" "$BLUE" "$EMOJI_GEAR" "$1" "$NC"; }

die() {
  log_error "$1"
  exit 1
}

command_exists() {
  command -v "$1" >/dev/null 2>&1
}

trim() {
  local value="$1"
  value="${value#"${value%%[![:space:]]*}"}"
  value="${value%"${value##*[![:space:]]}"}"
  printf '%s' "$value"
}

contains_value() {
  local needle="$1"
  shift
  local item
  for item in "$@"; do
    [[ "$item" == "$needle" ]] && return 0
  done
  return 1
}

is_valid_ipv4() {
  local ip="$1"
  [[ "$ip" =~ ^([0-9]{1,3}\.){3}[0-9]{1,3}$ ]] || return 1

  local o1 o2 o3 o4
  IFS='.' read -r o1 o2 o3 o4 <<<"$ip"
  local octet
  for octet in "$o1" "$o2" "$o3" "$o4"; do
    [[ "$octet" =~ ^[0-9]+$ ]] || return 1
    ((octet >= 0 && octet <= 255)) || return 1
  done
  return 0
}

is_valid_ipv6() {
  local ip="$1"
  local ipv6_re='^(([0-9A-Fa-f]{1,4}:){7}([0-9A-Fa-f]{1,4}|:)|([0-9A-Fa-f]{1,4}:){1,7}:|([0-9A-Fa-f]{1,4}:){1,6}:[0-9A-Fa-f]{1,4}|([0-9A-Fa-f]{1,4}:){1,5}(:[0-9A-Fa-f]{1,4}){1,2}|([0-9A-Fa-f]{1,4}:){1,4}(:[0-9A-Fa-f]{1,4}){1,3}|([0-9A-Fa-f]{1,4}:){1,3}(:[0-9A-Fa-f]{1,4}){1,4}|([0-9A-Fa-f]{1,4}:){1,2}(:[0-9A-Fa-f]{1,4}){1,5}|[0-9A-Fa-f]{1,4}(:[0-9A-Fa-f]{1,4}){1,6}|:((:[0-9A-Fa-f]{1,4}){1,7}|:))(%.+)?$'
  [[ "$ip" =~ $ipv6_re ]]
}

is_valid_ip() {
  local ip="$1"
  is_valid_ipv4 "$ip" || is_valid_ipv6 "$ip"
}

require_tty() {
  [[ -r /dev/tty && -w /dev/tty ]] || die "No interactive TTY found. This script needs a terminal for prompts (even when using curl | bash)."
}

read_tty_input() {
  local prompt="$1"
  local value
  printf "%b" "$prompt" >/dev/tty
  IFS= read -r value </dev/tty || return 1
  printf '%s' "$(trim "$value")"
}

run_root() {
  if [[ -n "$SUDO_BIN" ]]; then
    "$SUDO_BIN" "$@"
  else
    "$@"
  fi
}

write_root_file() {
  local target="$1"
  local content="$2"
  if [[ -n "$SUDO_BIN" ]]; then
    printf '%s' "$content" | "$SUDO_BIN" tee "$target" >/dev/null
  else
    printf '%s' "$content" >"$target"
  fi
}

append_root_file() {
  local target="$1"
  local content="$2"
  if [[ -n "$SUDO_BIN" ]]; then
    printf '%s' "$content" | "$SUDO_BIN" tee -a "$target" >/dev/null
  else
    printf '%s' "$content" >>"$target"
  fi
}

ensure_root_context() {
  if [[ "$EUID" -eq 0 ]]; then
    return 0
  fi

  command_exists sudo || die "This script must run as root or with sudo available."

  local source_path="${BASH_SOURCE[0]:-}"
  if [[ -n "$source_path" && -f "$source_path" ]]; then
    log_info "Re-launching with sudo for full root context..."
    exec sudo -E bash "$source_path" "$@"
  fi

  log_warning "Running from stdin (likely curl | bash). Falling back to sudo for privileged operations."
  sudo -v || die "Could not acquire sudo privileges."
  SUDO_BIN="sudo"
}

add_manager() {
  local manager="$1"
  if ! contains_value "$manager" "${DETECTED_MANAGERS[@]}"; then
    DETECTED_MANAGERS+=("$manager")
  fi
}

detect_managers() {
  DETECTED_MANAGERS=()

  if command_exists systemctl && systemctl is-active --quiet systemd-resolved 2>/dev/null; then
    add_manager "systemd-resolved"
  elif [[ -L /etc/resolv.conf ]] && readlink /etc/resolv.conf 2>/dev/null | grep -Eq 'systemd/resolve'; then
    add_manager "systemd-resolved"
  fi

  if command_exists nmcli || [[ -d /etc/NetworkManager ]]; then
    add_manager "NetworkManager"
  fi

  if command_exists resolvconf || [[ -d /etc/resolvconf ]] || [[ -f /etc/resolvconf.conf ]]; then
    add_manager "resolvconf/openresolv"
  fi

  if command_exists dhclient || [[ -f /etc/dhcp/dhclient.conf ]]; then
    add_manager "dhclient"
  fi

  if command_exists systemctl && systemctl is-active --quiet systemd-networkd 2>/dev/null; then
    add_manager "systemd-networkd"
  elif [[ -d /etc/systemd/network ]]; then
    add_manager "systemd-networkd"
  fi

  if [[ -f /etc/network/interfaces ]] || [[ -d /etc/network/interfaces.d ]]; then
    add_manager "ifupdown"
  fi
}

prompt_dns_servers() {
  DNS_SERVERS=()
  local input
  local prompt
  local index

  for index in 1 2; do
    while true; do
      prompt="${BLUE}${EMOJI_NET} Enter DNS server ${index} (required): ${NC}"
      input="$(read_tty_input "$prompt")" || die "Failed to read DNS input from TTY."

      [[ -n "$input" ]] || {
        log_warning "DNS server ${index} is required."
        continue
      }

      if ! is_valid_ip "$input"; then
        log_warning "Invalid IP literal: '$input'. Please enter a valid IPv4 or IPv6 address."
        continue
      fi

      if contains_value "$input" "${DNS_SERVERS[@]}"; then
        log_warning "Duplicate DNS server '$input' is not allowed."
        continue
      fi

      DNS_SERVERS+=("$input")
      break
    done
  done

  while true; do
    prompt="${BLUE}${EMOJI_NET} Enter DNS server 3 (optional, press Enter to skip): ${NC}"
    input="$(read_tty_input "$prompt")" || die "Failed to read DNS input from TTY."

    if [[ -z "$input" ]]; then
      break
    fi

    if ! is_valid_ip "$input"; then
      log_warning "Invalid IP literal: '$input'. Please enter a valid IPv4 or IPv6 address."
      continue
    fi

    if contains_value "$input" "${DNS_SERVERS[@]}"; then
      log_warning "Duplicate DNS server '$input' is not allowed."
      continue
    fi

    DNS_SERVERS+=("$input")
    break
  done

  log_success "DNS server list accepted: ${DNS_SERVERS[*]}"
}

check_dns_connectivity() {
  local dns="$1"

  if command_exists dig; then
    dig @"$dns" example.com A +time=2 +tries=1 +short >/dev/null 2>&1 && return 0
    dig @"$dns" . NS +time=2 +tries=1 +short >/dev/null 2>&1 && return 0
    return 1
  fi

  if command_exists drill; then
    drill @"$dns" example.com A >/dev/null 2>&1 && return 0
    drill @"$dns" . NS >/dev/null 2>&1 && return 0
    return 1
  fi

  if command_exists nslookup; then
    nslookup -timeout=2 example.com "$dns" >/dev/null 2>&1 && return 0
    nslookup example.com "$dns" >/dev/null 2>&1 && return 0
    return 1
  fi

  if command_exists nc; then
    nc -z -w2 "$dns" 53 >/dev/null 2>&1 && return 0
    return 1
  fi

  if [[ "$dns" != *:* ]] && command_exists timeout; then
    timeout 2 bash -c "exec 3<>/dev/tcp/$dns/53" >/dev/null 2>&1 && return 0
    return 1
  fi

  return 2
}

run_connectivity_checks() {
  log_step "Running best-effort connectivity checks for supplied DNS servers..."

  local dns
  local rc
  for dns in "${DNS_SERVERS[@]}"; do
    if check_dns_connectivity "$dns"; then
      log_success "Reachability check passed for $dns"
      continue
    fi

    rc=$?
    if [[ "$rc" -eq 2 ]]; then
      log_warning "No available tool for probing $dns on this host. Proceeding anyway."
    else
      log_warning "Could not reach DNS server $dns during probe. Proceeding as requested."
    fi
  done
}

dns_nameserver_lines() {
  local lines=""
  local dns
  for dns in "${DNS_SERVERS[@]}"; do
    lines+="nameserver ${dns}"$'\n'
  done
  printf '%s' "$lines"
}

dns_space_list() {
  printf '%s' "${DNS_SERVERS[*]}"
}

dns_comma_list() {
  local out=""
  local dns
  for dns in "${DNS_SERVERS[@]}"; do
    if [[ -z "$out" ]]; then
      out="$dns"
    else
      out+=",${dns}"
    fi
  done
  printf '%s' "$out"
}

run_action() {
  local action="$1"
  shift

  ACTIONS_ATTEMPTED+=("$action")
  log_step "$action"

  local rc=0
  if "$@"; then
    ACTIONS_SUCCEEDED+=("$action")
    log_success "$action completed"
    return 0
  fi

  rc=$?
  if [[ "$rc" -eq 2 ]]; then
    ACTIONS_SKIPPED+=("$action")
    log_info "$action not applicable on this system"
    return 0
  fi

  ACTIONS_FAILED+=("$action")
  log_warning "$action failed (best-effort mode: continuing)"
  return 0
}

ensure_resolv_conf_mutable_if_needed() {
  if command_exists lsattr && command_exists chattr; then
    local attr_line attr_field
    attr_line="$(run_root lsattr /etc/resolv.conf 2>/dev/null || true)"
    attr_field="${attr_line%% *}"
    if [[ "$attr_field" == *i* ]]; then
      log_warning "/etc/resolv.conf is immutable. Temporarily removing immutable flag for update."
      run_root chattr -i /etc/resolv.conf || return 1
    fi
  fi
  return 0
}

apply_static_resolv_conf() {
  ensure_resolv_conf_mutable_if_needed || return 1

  local now lines content
  now="$(date -u '+%Y-%m-%dT%H:%M:%SZ')"
  lines="$(dns_nameserver_lines)"
  content="# Managed by ${SCRIPT_NAME} v${SCRIPT_VERSION} at ${now}"$'\n'
  content+="$lines"
  content+="options timeout:2 attempts:2 rotate"$'\n'

  write_root_file "/etc/resolv.conf" "$content" || return 1
  run_root chmod 644 /etc/resolv.conf || true
  return 0
}

apply_systemd_resolved() {
  if ! command_exists systemctl && ! command_exists resolvectl && [[ ! -d /etc/systemd ]]; then
    return 2
  fi

  local dns_line content
  dns_line="$(dns_space_list)"
  content="[Resolve]"$'\n'
  content+="DNS=${dns_line}"$'\n'
  content+="FallbackDNS="$'\n'
  content+="Domains=~."$'\n'

  run_root mkdir -p /etc/systemd/resolved.conf.d || return 1
  write_root_file "/etc/systemd/resolved.conf.d/99-dns-override.conf" "$content" || return 1

  if command_exists systemctl; then
    run_root systemctl daemon-reload >/dev/null 2>&1 || true
    run_root systemctl restart systemd-resolved >/dev/null 2>&1 || run_root systemctl try-restart systemd-resolved >/dev/null 2>&1 || true
  fi

  if command_exists resolvectl; then
    if command_exists ip; then
      local link
      while IFS= read -r link; do
        [[ -n "$link" && "$link" != "lo" ]] || continue
        run_root resolvectl dns "$link" "${DNS_SERVERS[@]}" >/dev/null 2>&1 || true
        run_root resolvectl domain "$link" "~." >/dev/null 2>&1 || true
      done < <(ip -o link show | awk -F': ' '{print $2}' | cut -d'@' -f1)
    fi
    run_root resolvectl flush-caches >/dev/null 2>&1 || true
  fi

  return 0
}

apply_networkmanager() {
  if ! command_exists nmcli && [[ ! -d /etc/NetworkManager ]]; then
    return 2
  fi

  local dns_csv nm_conf
  dns_csv="$(dns_comma_list)"
  nm_conf="[main]"$'\n'
  nm_conf+="dns=default"$'\n'
  nm_conf+="rc-manager=symlink"$'\n'

  run_root mkdir -p /etc/NetworkManager/conf.d || return 1
  write_root_file "/etc/NetworkManager/conf.d/90-dns-override.conf" "$nm_conf" || return 1

  if command_exists nmcli; then
    local conn
    while IFS= read -r conn; do
      [[ -n "$conn" ]] || continue
      run_root nmcli connection modify "$conn" ipv4.ignore-auto-dns yes >/dev/null 2>&1 || true
      run_root nmcli connection modify "$conn" ipv6.ignore-auto-dns yes >/dev/null 2>&1 || true
      run_root nmcli connection modify "$conn" ipv4.dns "$dns_csv" >/dev/null 2>&1 || true
      run_root nmcli connection modify "$conn" ipv6.dns "$dns_csv" >/dev/null 2>&1 || true
      run_root nmcli connection up "$conn" >/dev/null 2>&1 || true
    done < <(nmcli -g NAME connection show --active 2>/dev/null || true)

    run_root nmcli general reload >/dev/null 2>&1 || true
  fi

  if command_exists systemctl; then
    run_root systemctl reload NetworkManager >/dev/null 2>&1 || run_root systemctl restart NetworkManager >/dev/null 2>&1 || true
  fi

  return 0
}

apply_resolvconf_openresolv() {
  if ! command_exists resolvconf && [[ ! -d /etc/resolvconf ]] && [[ ! -f /etc/resolvconf.conf ]]; then
    return 2
  fi

  local lines
  lines="$(dns_nameserver_lines)"

  run_root mkdir -p /etc/resolvconf/resolv.conf.d || true
  write_root_file "/etc/resolvconf/resolv.conf.d/head" "$lines" || true
  write_root_file "/etc/resolvconf/resolv.conf.d/base" "$lines" || true

  if command_exists resolvconf; then
    printf '%s' "$lines" | run_root resolvconf -a dns-override >/dev/null 2>&1 || true
    run_root resolvconf -u >/dev/null 2>&1 || true
  fi

  return 0
}

apply_dhclient() {
  if ! command_exists dhclient && [[ ! -f /etc/dhcp/dhclient.conf ]] && [[ ! -d /etc/dhcp ]]; then
    return 2
  fi

  local conf_path="/etc/dhcp/dhclient.conf"
  local dns_csv block
  dns_csv="$(dns_comma_list)"
  block=$'\n'"# dns-override.sh managed begin"$'\n'
  block+="supersede domain-name-servers ${dns_csv};"$'\n'
  block+="prepend domain-name-servers ${dns_csv};"$'\n'
  block+="# dns-override.sh managed end"$'\n'

  run_root mkdir -p /etc/dhcp || true
  run_root touch "$conf_path" || return 1
  run_root sed -i '/# dns-override.sh managed begin/,/# dns-override.sh managed end/d' "$conf_path" || true
  append_root_file "$conf_path" "$block" || return 1

  return 0
}

apply_systemd_networkd() {
  if ! command_exists systemctl && [[ ! -d /etc/systemd/network ]]; then
    return 2
  fi

  local dns_lines dropin_content
  dns_lines=""
  local dns
  for dns in "${DNS_SERVERS[@]}"; do
    dns_lines+="DNS=${dns}"$'\n'
  done

  run_root mkdir -p /etc/systemd/network || return 1

  local has_network_file=0
  local net_file
  for net_file in /etc/systemd/network/*.network; do
    [[ -e "$net_file" ]] || continue
    has_network_file=1
    run_root mkdir -p "${net_file}.d" || true
    dropin_content="[Network]"$'\n'"${dns_lines}Domains=~."$'\n'
    write_root_file "${net_file}.d/99-dns-override.conf" "$dropin_content" || true
  done

  if [[ "$has_network_file" -eq 0 ]]; then
    dropin_content="[Match]"$'\n'"Name=*"$'\n\n'"[Network]"$'\n'"${dns_lines}Domains=~."$'\n'
    write_root_file "/etc/systemd/network/99-dns-override.network" "$dropin_content" || true
  fi

  if command_exists systemctl; then
    run_root systemctl daemon-reload >/dev/null 2>&1 || true
    run_root systemctl restart systemd-networkd >/dev/null 2>&1 || run_root systemctl try-restart systemd-networkd >/dev/null 2>&1 || true
  fi

  return 0
}

patch_ifupdown_file() {
  local file_path="$1"
  [[ -f "$file_path" ]] || return 1

  local dns_line
  dns_line="$(dns_space_list)"
  local temp_file
  temp_file="$(mktemp)"

  set +e
  awk -v dns="$dns_line" '
    BEGIN { iface_count = 0 }
    /^[[:space:]]*dns-nameservers[[:space:]]+/ { next }
    {
      print
      if ($0 ~ /^[[:space:]]*iface[[:space:]]+/ && $0 !~ /[[:space:]]lo[[:space:]]/) {
        print "    dns-nameservers " dns
        iface_count++
      }
    }
    END {
      if (iface_count == 0) {
        exit 2
      }
    }
  ' "$file_path" >"$temp_file"
  local rc=$?
  set -e
  if [[ "$rc" -eq 2 ]]; then
    rm -f "$temp_file"
    return 2
  fi
  [[ "$rc" -eq 0 ]] || {
    rm -f "$temp_file"
    return 1
  }

  if [[ -n "$SUDO_BIN" ]]; then
    "$SUDO_BIN" mv "$temp_file" "$file_path"
  else
    mv "$temp_file" "$file_path"
  fi
  return 0
}

apply_ifupdown() {
  if [[ ! -f /etc/network/interfaces ]] && [[ ! -d /etc/network/interfaces.d ]]; then
    return 2
  fi

  local touched=0
  local patch_rc=0

  if [[ -f /etc/network/interfaces ]]; then
    if patch_ifupdown_file /etc/network/interfaces; then
      touched=1
    else
      patch_rc=$?
      [[ "$patch_rc" -eq 2 ]] || return 1
    fi
  fi

  if [[ -d /etc/network/interfaces.d ]]; then
    local iface_file rc
    for iface_file in /etc/network/interfaces.d/*; do
      [[ -f "$iface_file" ]] || continue
      if patch_ifupdown_file "$iface_file"; then
        touched=1
      else
        rc=$?
        [[ "$rc" -eq 2 ]] || return 1
      fi
    done
  fi

  if [[ "$touched" -eq 0 ]]; then
    log_warning "ifupdown detected, but no writable iface stanzas were found to patch automatically."
    log_info "Actionable hint: add 'dns-nameservers $(dns_space_list)' inside each target iface block."
  fi

  return 0
}

prompt_immutable_lock() {
  command_exists chattr || {
    log_info "chattr not available; skipping immutable lock prompt."
    return 0
  }

  local answer
  answer="$(read_tty_input "${YELLOW}${EMOJI_LOCK} Make /etc/resolv.conf immutable with chattr +i? [y/N]: ${NC}")" || return 1
  case "$answer" in
    y|Y|yes|YES)
      if run_root chattr +i /etc/resolv.conf; then
        log_success "Immutable lock applied to /etc/resolv.conf."
        log_warning "To edit DNS later, run: sudo chattr -i /etc/resolv.conf"
      else
        log_warning "Failed to apply immutable lock."
      fi
      ;;
    *)
      log_info "Immutable lock skipped."
      ;;
  esac
}

print_summary() {
  echo
  log_step "Execution summary"
  log_info "Detected managers: ${DETECTED_MANAGERS[*]:-(none detected)}"
  log_info "Actions attempted: ${#ACTIONS_ATTEMPTED[@]}"
  log_success "Actions succeeded: ${#ACTIONS_SUCCEEDED[@]}"
  log_warning "Actions failed: ${#ACTIONS_FAILED[@]}"
  log_info "Actions skipped: ${#ACTIONS_SKIPPED[@]}"

  if [[ "${#ACTIONS_FAILED[@]}" -gt 0 ]]; then
    local failed
    for failed in "${ACTIONS_FAILED[@]}"; do
      log_warning "Failed action: $failed"
    done
  fi

  echo
  log_step "Final /etc/resolv.conf"
  run_root cat /etc/resolv.conf 2>/dev/null || cat /etc/resolv.conf 2>/dev/null || log_warning "Could not read /etc/resolv.conf"

  if command_exists resolvectl; then
    echo
    log_step "resolvectl status snapshot"
    run_root resolvectl status 2>/dev/null || resolvectl status 2>/dev/null || log_warning "Could not get resolvectl status."
  fi

  if command_exists nmcli; then
    echo
    log_step "NetworkManager DNS snapshot"
    run_root nmcli -f GENERAL.DEVICE,IP4.DNS,IP6.DNS device show 2>/dev/null || nmcli -f GENERAL.DEVICE,IP4.DNS,IP6.DNS device show 2>/dev/null || log_warning "Could not get nmcli DNS snapshot."
  fi
}

main() {
  echo
  printf "%b%s %s v%s%b\n" "$BLUE" "$EMOJI_ROCKET" "$SCRIPT_NAME" "$SCRIPT_VERSION" "$NC"
  log_info "Aggressive best-effort DNS override for Linux hosts."
  log_warning "This script edits multiple DNS manager configs and may restart networking services."
  echo

  require_tty
  ensure_root_context "$@"
  prompt_dns_servers
  run_connectivity_checks
  detect_managers

  run_action "Apply systemd-resolved override" apply_systemd_resolved
  run_action "Apply NetworkManager override" apply_networkmanager
  run_action "Apply resolvconf/openresolv override" apply_resolvconf_openresolv
  run_action "Apply dhclient override" apply_dhclient
  run_action "Apply systemd-networkd override" apply_systemd_networkd
  run_action "Apply ifupdown override" apply_ifupdown
  run_action "Write /etc/resolv.conf directly" apply_static_resolv_conf

  prompt_immutable_lock || log_warning "Immutable lock prompt failed; continuing."
  print_summary

  echo
  log_success "DNS override flow completed."
}

main "$@"