#!/bin/bash
set -e

# ─────────────────────────────────────────────
# Unbound setup script (no package installation)
# ─────────────────────────────────────────────

RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
NC='\033[0m'

info()    { echo -e "${GREEN}[+]${NC} $1"; }
warn()    { echo -e "${YELLOW}[!]${NC} $1"; }
error()   { echo -e "${RED}[✗]${NC} $1"; exit 1; }
success() { echo -e "${GREEN}[✓]${NC} $1"; }

# ── Root check ──────────────────────────────
if [[ $EUID -ne 0 ]]; then
    error "Run as root: sudo bash $0"
fi

# ── Check unbound is installed ──────────────
if ! command -v unbound &>/dev/null; then
    error "unbound is not installed. Install it first:\n  apt install unbound  /  dnf install unbound  /  pacman -S unbound"
fi

# ── Variables ───────────────────────────────
UNBOUND_DIR="/var/lib/unbound"
CONF_DIR="/etc/unbound/unbound.conf.d"
CONF_FILE="$CONF_DIR/personal.conf"
MAIN_CONF="/etc/unbound/unbound.conf"
ROOT_HINTS="$UNBOUND_DIR/root.hints"
ROOT_HINTS_URL="https://www.internic.net/domain/named.root"

# ── Create directories ───────────────────────
info "Creating directories..."
mkdir -p "$UNBOUND_DIR"
mkdir -p "$CONF_DIR"

# ── Download root hints ──────────────────────
info "Downloading root hints..."
if command -v curl &>/dev/null; then
    curl -fsSL -o "$ROOT_HINTS" "$ROOT_HINTS_URL" || error "Failed to download root hints"
elif command -v wget &>/dev/null; then
    wget -q -O "$ROOT_HINTS" "$ROOT_HINTS_URL" || error "Failed to download root hints"
else
    error "Neither curl nor wget found. Cannot download root hints."
fi
success "Root hints downloaded"

# ── Fix ownership ─────────────────────────────
info "Setting permissions..."
chown -R unbound:unbound "$UNBOUND_DIR" 2>/dev/null || warn "Could not chown $UNBOUND_DIR (unbound user may not exist)"

# ── Patch main config ─────────────────────────
info "Patching main config..."

# Disable any active trust-anchor-file (causes SERVFAIL without valid key)
sed -i 's|^[[:space:]]*trust-anchor-file:|#	trust-anchor-file:|g' "$MAIN_CONF"

# Disable validator module if explicitly set
sed -i 's|^[[:space:]]*module-config:.*validator.*|# & # disabled by setup script|g' "$MAIN_CONF"

# Add include for conf.d if not already present
if ! grep -q 'unbound.conf.d' "$MAIN_CONF"; then
    echo 'include: "/etc/unbound/unbound.conf.d/*.conf"' >> "$MAIN_CONF"
    success "Added include for conf.d"
fi

# ── Write personal config ─────────────────────
info "Writing config to $CONF_FILE..."
cat > "$CONF_FILE" <<EOF
server:
    # Disable chroot so paths work correctly
    chroot: ""

    # No DNSSEC validation — iterator only
    module-config: "iterator"

    # IPv4 only (IPv6 root servers unreachable on many setups)
    do-ip6: no

    # Listen on localhost only
    interface: 127.0.0.1
    port: 53

    # Allow only localhost
    access-control: 127.0.0.0/8 allow
    access-control: 0.0.0.0/0 refuse

    # Root hints
    root-hints: "$ROOT_HINTS"

    # Performance
    num-threads: 2
    cache-min-ttl: 300
    cache-max-ttl: 86400
    msg-cache-size: 50m
    rrset-cache-size: 100m
    so-reuseport: yes

    # Privacy
    hide-identity: yes
    hide-version: yes
    qname-minimisation: yes

    # Prefetch popular records before they expire
    prefetch: yes
    prefetch-key: yes

    # Logging (minimal)
    verbosity: 1
    log-queries: no
EOF
success "Config written"

# ── Validate config ───────────────────────────
info "Validating config..."
if unbound-checkconf "$MAIN_CONF" &>/dev/null; then
    success "Config is valid"
else
    unbound-checkconf "$MAIN_CONF" || error "Config validation failed — see above"
fi

# ── Enable and restart service ────────────────
info "Enabling and restarting unbound..."
systemctl enable unbound
systemctl restart unbound
sleep 1

if systemctl is-active --quiet unbound; then
    success "unbound is running"
else
    journalctl -u unbound -n 20 --no-pager
    error "unbound failed to start — see logs above"
fi

# ── Verify DNS resolution ─────────────────────
info "Testing DNS resolution..."
sleep 1

if dig @127.0.0.1 google.com +short &>/dev/null; then
    success "DNS resolution works"
else
    warn "dig test failed — unbound may need a forwarder"
    warn "Try manually: dig @127.0.0.1 google.com"
    warn "If it fails, add a forwarder to $CONF_FILE:"
    warn "  forward-zone:"
    warn "      name: \".\""
    warn "      forward-addr: 1.1.1.1"
fi

# ── Done ──────────────────────────────────────
echo ""
echo -e "${GREEN}══════════════════════════════════════${NC}"
echo -e "${GREEN}  Unbound setup complete!${NC}"
echo -e "${GREEN}══════════════════════════════════════${NC}"
echo ""
echo "  Verify with:"
echo "    dig @127.0.0.1 google.com"
echo "    dig @127.0.0.1 google.com | grep 'Query time'  # run twice, 2nd should be 0ms"
echo ""
echo "  To use as system resolver, add to /etc/resolv.conf:"
echo "    nameserver 127.0.0.1"
echo ""
echo "  Or in WireGuard [Interface]:"
echo "    DNS = 127.0.0.1"
echo ""