#!/bin/bash

# Must run as root
if [[ $EUID -ne 0 ]]; then
    echo "🚨 This script must be run as root."
    echo "👉 Please run it using: sudo $0"
    exit 1
fi

DOMAINS_FILE="domains.txt"
HOSTS_FILE="/etc/hosts"
TMP_HOSTS=$(mktemp)

cp "$HOSTS_FILE" "$TMP_HOSTS"

CURRENT_CATEGORY=""
declare -a CATEGORY_DOMAINS=()

log() {
    echo "🔧 $1"
}

add_ipv4_block() {
    local ip="127.0.0.1"
    local comment="$1"

    log "Adding IPv4 entries under category: $comment"

    if grep -Fxq "$comment" "$TMP_HOSTS"; then
        END_LINE=$(awk -v cat="$comment" '
            $0 == cat { found=1; next }
            found && $0 ~ /^#/ { exit }
            found { line=NR }
            END { print (line ? line : 0) }
        ' "$TMP_HOSTS")
        INSERT_LINE=$((END_LINE + 1))

        for domain in "${CATEGORY_DOMAINS[@]}"; do
            if ! grep -qE "^[[:space:]]*$ip[[:space:]]+$domain\$" "$TMP_HOSTS"; then
                log "  ➕ Inserting $ip $domain at line $INSERT_LINE"
                sed -i "${INSERT_LINE}i$ip $domain" "$TMP_HOSTS"
                ((INSERT_LINE++))
            else
                log "  ⚠️ Skipping duplicate: $ip $domain"
            fi
        done
    else
        log "  🔹 Category not found, appending new category: $comment"
        echo "" >> "$TMP_HOSTS"
        echo "$comment" >> "$TMP_HOSTS"
        for domain in "${CATEGORY_DOMAINS[@]}"; do
            echo "$ip $domain" >> "$TMP_HOSTS"
            log "  ➕ Appended $ip $domain"
        done
    fi
}

add_ipv6_block() {
    local ip="::1"
    local comment="$1 ipv6"

    log "Adding IPv6 entries under category: $comment"

    if grep -Fxq "$comment" "$TMP_HOSTS"; then
        END_LINE=$(awk -v cat="$comment" '
            $0 == cat { found=1; next }
            found && $0 ~ /^#/ { exit }
            found { line=NR }
            END { print (line ? line : 0) }
        ' "$TMP_HOSTS")
        INSERT_LINE=$((END_LINE + 1))

        for domain in "${CATEGORY_DOMAINS[@]}"; do
            if ! grep -qE "^[[:space:]]*$ip[[:space:]]+$domain\$" "$TMP_HOSTS"; then
                log "  ➕ Inserting $ip $domain at line $INSERT_LINE"
                sed -i "${INSERT_LINE}i$ip $domain" "$TMP_HOSTS"
                ((INSERT_LINE++))
            else
                log "  ⚠️ Skipping duplicate: $ip $domain"
            fi
        done
    else
        log "  🔹 IPv6 category not found, appending: $comment"
        echo "" >> "$TMP_HOSTS"
        echo "$comment" >> "$TMP_HOSTS"
        for domain in "${CATEGORY_DOMAINS[@]}"; do
            echo "$ip $domain" >> "$TMP_HOSTS"
            log "  ➕ Appended $ip $domain"
        done
    fi
}

flush_category() {
    [[ -z "$CURRENT_CATEGORY" ]] && return

    log "📦 Flushing category: $CURRENT_CATEGORY"
    add_ipv4_block "$CURRENT_CATEGORY"
    add_ipv6_block "$CURRENT_CATEGORY"

    CURRENT_CATEGORY=""
    CATEGORY_DOMAINS=()
}

# Read and process domains.txt
while IFS= read -r line || [ -n "$line" ]; do
    line=$(echo "$line" | sed 's/^[ \t]*//;s/[ \t]*$//')  # Trim spaces
    [[ -z "$line" ]] && continue

    if [[ "$line" == \#* ]]; then
        flush_category
        CURRENT_CATEGORY="$line"
    else
        CATEGORY_DOMAINS+=("$line")
    fi
done < "$DOMAINS_FILE"

flush_category

# Backup and replace /etc/hosts
cp "$HOSTS_FILE" "${HOSTS_FILE}.bak"
cp "$TMP_HOSTS" "$HOSTS_FILE"

log "✅ Hosts file updated with IPv4 and IPv6 domain blocks under categories."
log "🗂️ Backup created at ${HOSTS_FILE}.bak"

rm "$TMP_HOSTS"

