#!/bin/bash

PORT=8080
DB="/usr/share/GeoIP/GeoLite2-Country.mmdb"
LOGFILE="/dev/null"
COUNTRY_FILE="countries.txt"

declare -A COUNTRY_COUNT
declare -A SEEN
declare -A COUNTRY_NAME

# --- Parse -N flag (e.g. -10 means show top 10) ---
TOPN=0
if [[ "$1" =~ ^-([0-9]+)$ ]]; then
    TOPN="${BASH_REMATCH[1]}"
fi

# --- Load external country list ---
while IFS='=' read -r ISO NAME; do
    [[ -z "$ISO" ]] && continue
    COUNTRY_NAME["$ISO"]="$NAME"
done < "$COUNTRY_FILE"

echo "Monitoring port $PORT..."
echo "Updating table every 5 seconds."
[[ $TOPN -gt 0 ]] && echo "Showing only Top $TOPN countries."

LAST_REFRESH=0

while true; do
    # FAST LOOP: collect new connections
    IPS=$(ss -tn sport = :$PORT | awk 'NR>1 {print $5}' | cut -d: -f1)

    for IP in $IPS; do
        [[ "$IP" == "127.0.0.1" ]] && continue

        if [[ -z "${SEEN[$IP]}" ]]; then
            SEEN[$IP]=1

            RAW=$(mmdblookup --file "$DB" --ip "$IP" country iso_code 2>/dev/null)
            ISO=$(echo "$RAW" | grep -oE '[A-Z]{2}')
            [[ -z "$ISO" ]] && ISO="UNKNOWN"

            COUNTRY_COUNT["$ISO"]=$(( COUNTRY_COUNT["$ISO"] + 1 ))

            NAME="${COUNTRY_NAME[$ISO]}"
            [[ -z "$NAME" ]] && NAME="Unknown Country"

            echo "$(date '+%F %T') - $IP - $ISO ($NAME)" >> "$LOGFILE"
        fi
    done

    # SLOW LOOP: refresh display every 5 seconds
    NOW=$(date +%s)
    if (( NOW - LAST_REFRESH >= 5 )); then
        LAST_REFRESH=$NOW

        clear
        echo "=== Live GeoIP Stats (Port $PORT) ==="
        echo "(Updated: $(date '+%H:%M:%S'))"
        echo

        # Sort by count (descending)
        SORTED=$(for ISO in "${!COUNTRY_COUNT[@]}"; do
            echo "${COUNTRY_COUNT[$ISO]} $ISO"
        done | sort -rn)

        COUNT=0
        while read -r LINE; do
            NUM=$(echo "$LINE" | awk '{print $1}')
            ISO=$(echo "$LINE" | awk '{print $2}')
            NAME="${COUNTRY_NAME[$ISO]}"
            [[ -z "$NAME" ]] && NAME="Unknown Country"

            echo "$ISO ($NAME): $NUM"

            ((COUNT++))
            [[ $TOPN -gt 0 && $COUNT -ge $TOPN ]] && break
        done <<< "$SORTED"
    fi

    sleep 0.1
done