#!/bin/bash
set -e

# ─────────────────────────────────────────────

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

if ! command -v unbound-anchor &>/dev/null; then
    error "unbound-anchor not found. Make sure the full unbound package is installed."
fi

# ── Variables ───────────────────────────────
UNBOUND_DIR="/var/lib/unbound"
CONF_DIR="/etc/unbound/unbound.conf.d"
CONF_FILE="$CONF_DIR/personal.conf"
ROOT_HINTS="$UNBOUND_DIR/root.hints"
ROOT_KEY="$UNBOUND_DIR/root.key"
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"

# ── Initialize DNSSEC anchor ─────────────────
info "Initializing DNSSEC trust anchor..."
unbound-anchor -a "$ROOT_KEY" || warn "unbound-anchor returned non-zero (this can be normal if key already exists)"
success "DNSSEC anchor initialized"

# ── 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)"

# ── Write config ──────────────────────────────
info "Writing config to $CONF_FILE..."
cat > "$CONF_FILE" <<'EOF'
server:
    # Listen on localhost only
    interface: 127.0.0.1
    interface: ::1
    port: 53

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

    # 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
    aggressive-nsec: yes

    # DNSSEC
    auto-trust-anchor-file: "/var/lib/unbound/root.key"

    # Root hints
    root-hints: "/var/lib/unbound/root.hints"

    # 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 "$CONF_FILE" &>/dev/null; then
    success "Config is valid"
else
    warn "Config check failed, trying full config..."
    if unbound-checkconf /etc/unbound/unbound.conf; then
        success "Full config is valid"
    else
        error "Config validation failed. Run: unbound-checkconf /etc/unbound/unbound.conf"
    fi
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
    error "unbound failed to start. Check: journalctl -u unbound -n 30"
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 — try manually: dig @127.0.0.1 google.com"
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 +dnssec   # look for 'ad' flag"
echo "    dig @127.0.0.1 sigfail.verteiltesysteme.net  # should be SERVFAIL"
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 ""