#!/bin/bash
set -e

# 颜色
RED='\033[0;31m'; GREEN='\033[0;32m'; YELLOW='\033[1;33m'; BLUE='\033[0;34m'; NC='\033[0m'

# 日志函数
log_info() { echo -e "${BLUE}[INFO]${NC} $1" >&2; }
log_success() { echo -e "${GREEN}[SUCCESS]${NC} $1" >&2; }
log_warning() { echo -e "${YELLOW}[WARNING]${NC} $1" >&2; }
log_error() { echo -e "${RED}[ERROR]${NC} $1" >&2; }

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
LOCAL_DNS_CONF="/opt/argus-metric/dns.conf"
RESOLV_CONF="/etc/resolv.conf"
ALT_RESOLV_CONF="/run/resolv.conf"
LOG_FILE="/opt/argus-metric/.dns_sync.log"
REMOTE_DNS_CONF_URL=""

# 获取 FTP 配置
get_ftp_config() {
    log_info "获取 FTP 配置信息..."
    if [[ -z "$FTP_SERVER" || -z "$FTP_USER" || -z "$FTP_PASSWORD" ]]; then
        [[ -f "$SCRIPT_DIR/config.env" ]] && source "$SCRIPT_DIR/config.env"
    fi
    FTP_SERVER="${FTP_SERVER:-localhost}"
    FTP_USER="${FTP_USER:-ftpuser}"
    FTP_PASSWORD="${FTP_PASSWORD:-ZGClab1234!}"
    REMOTE_DNS_CONF_URL="ftp://${FTP_USER}:${FTP_PASSWORD}@${FTP_SERVER}/dns.conf"
}

# 下载远程 dns.conf
download_remote_dns_conf() {
    local tmp="/tmp/dns.remote.$$"
    log_info "测试 FTP 连接..."
    if ! curl -u "${FTP_USER}:${FTP_PASSWORD}" -sfI "ftp://${FTP_SERVER}/" >/dev/null; then
        log_error "无法连接到 FTP 服务器: $FTP_SERVER"; return 1
    fi
    if ! curl -u "${FTP_USER}:${FTP_PASSWORD}" -sf "ftp://${FTP_SERVER}/dns.conf" -o "$tmp" 2>/dev/null; then
        log_error "下载 dns.conf 失败"; rm -f "$tmp"; return 1
    fi
    echo "$tmp"
}

# 文件比较
compare_files() { diff -q "$1" "$2" >/dev/null 2>&1; }

# 从 dns.conf 提取有效 IP
get_dns_ips() {
    grep -Eo '^[0-9]{1,3}(\.[0-9]{1,3}){3}$' "$1" | sort -u
}

# 安全更新 resolv.conf（保留符号链接）
update_resolv_conf() {
    local dns_conf="$1"
    local dns_ips
    mapfile -t dns_ips < <(get_dns_ips "$dns_conf")
    [[ ${#dns_ips[@]} -eq 0 ]] && { log_warning "未检测到有效 DNS"; return; }

    local target_file="$RESOLV_CONF"
    if [[ ! -w "$RESOLV_CONF" ]]; then
        log_warning "/etc/resolv.conf 不可写，使用兜底路径 $ALT_RESOLV_CONF"
        target_file="$ALT_RESOLV_CONF"
    fi

    local temp="/tmp/resolv.new.$$"
    cp "$target_file" "${target_file}.backup.$(date +%Y%m%d_%H%M%S)" 2>/dev/null || true
    log_info "更新 DNS 配置文件: $target_file"

    # 写入新的 nameserver 行
    for ip in "${dns_ips[@]}"; do
        echo "nameserver $ip"
    done >"$temp"

    # 追加原内容（去掉重复 nameserver）
    grep -v '^nameserver' "$target_file" >>"$temp" 2>/dev/null || true
    awk '!a[$0]++' "$temp" >"${temp}.uniq"

    # ⚙️ 使用 cat 原地覆盖，避免 mv 引发 “设备忙”
    if cat "${temp}.uniq" >"$target_file" 2>/dev/null; then
        chmod 644 "$target_file"
        log_success "DNS 更新完成: ${dns_ips[*]}"
    else
        log_error "无法写入 $target_file，可能被系统锁定"
    fi

    rm -f "$temp" "${temp}.uniq"
}

# 检查 resolv.conf 是否包含 dns.conf 内容
ensure_dns_in_resolv() {
    local dns_conf="$1"
    local dns_ips
    mapfile -t dns_ips < <(get_dns_ips "$dns_conf")
    [[ ${#dns_ips[@]} -eq 0 ]] && return

    for ip in "${dns_ips[@]}"; do
        if ! grep -q "nameserver $ip" "$RESOLV_CONF" 2>/dev/null; then
            log_warning "检测到 /etc/resolv.conf 缺少 $ip，执行兜底修复"
            update_resolv_conf "$dns_conf"
            return
        fi
    done
    log_info "/etc/resolv.conf 已包含所有 DNS"
}

log_sync() { echo "[$(date '+%F %T')] $1" >>"$LOG_FILE"; }

main() {
    log_info "开始 DNS 同步检查..."
    mkdir -p /opt/argus-metric

    get_ftp_config
    local remote_file
    if ! remote_file=$(download_remote_dns_conf); then
        log_error "下载失败"; log_sync "同步失败"; exit 1
    fi

    if [[ ! -f "$LOCAL_DNS_CONF" ]]; then
        log_info "本地 dns.conf 不存在，初始化..."
        cp "$remote_file" "$LOCAL_DNS_CONF"
        update_resolv_conf "$LOCAL_DNS_CONF"
        log_sync "首次同步完成"
    else
        if compare_files "$LOCAL_DNS_CONF" "$remote_file"; then
            log_info "dns.conf 无变化"
            ensure_dns_in_resolv "$LOCAL_DNS_CONF"
            log_sync "dns.conf 无变化，执行兜底检查"
        else
            log_info "检测到 DNS 配置更新"
            cp "$remote_file" "$LOCAL_DNS_CONF"
            update_resolv_conf "$LOCAL_DNS_CONF"
            log_sync "DNS 配置同步完成"
        fi
    fi

    rm -f "$remote_file"
    log_success "DNS 同步流程完成"
}

if [[ "${BASH_SOURCE[0]}" == "${0}" ]]; then
    main "$@"
fi