#!/bin/bash

HOSTS_FILE="/etc/hosts"
IP_ADDRESS="127.0.0.1"

if [ $# -lt 2 ]; then
    echo "Usage: $0 {add|remove} <directory_path> [extra_domain1] [extra_domain2] ..."
    exit 1
fi

ACTION="$1"
SITES_DIR="$2"

shift 2
EXTRA_DOMAINS=("$@")

ALL_DOMAINS=("${EXTRA_DOMAINS[@]}") 

if [ -d "$SITES_DIR" ]; then
    shopt -s nullglob

    for dir in "$SITES_DIR"/*/; do
        dirname=$(basename "$dir")
        ALL_DOMAINS+=("$dirname")
    done
    
    shopt -u nullglob
    echo "Scanned directory: $SITES_DIR"
else
    echo "Warning: Directory '$SITES_DIR' not found or is not a directory."
fi

add_subdomains() {
    echo "Adding entries to /etc/hosts..."
    for DOMAIN in "${ALL_DOMAINS[@]}"; do
        [ -z "$DOMAIN" ] && continue

        if ! grep -q "[[:space:]]$DOMAIN$" "$HOSTS_FILE"; then
            echo "$IP_ADDRESS $DOMAIN" | sudo tee -a "$HOSTS_FILE" > /dev/null
            echo "Added $DOMAIN"
        else
            echo "$DOMAIN is already in /etc/hosts"
        fi
    done
}

remove_subdomains() {
    echo "Removing entries from /etc/hosts..."
    for DOMAIN in "${ALL_DOMAINS[@]}"; do
        [ -z "$DOMAIN" ] && continue
        if grep -q "[[:space:]]$DOMAIN$" "$HOSTS_FILE"; then
            sudo sed -i "/[[:space:]]$DOMAIN$/d" "$HOSTS_FILE"
            echo "Removed $DOMAIN"
        else
            echo "$DOMAIN was not found in /etc/hosts"
        fi
    done
}

case "$ACTION" in
    add)
        if [ ${#ALL_DOMAINS[@]} -eq 0 ]; then
            echo "No domains found to add."
            exit 0
        fi
        add_subdomains
        ;;
    remove)
        if [ ${#ALL_DOMAINS[@]} -eq 0 ]; then
            echo "No domains found to remove."
            exit 0
        fi
        remove_subdomains
        ;;
    *)
        echo "Usage: $0 {add|remove} <directory_path> [extra_domains...]"
        exit 1
        ;;
esac