#!/bin/sh

# mwan3-lb-test - Load balancing distribution verifier
#
# Verifies that a mwan3 load-balancing policy distributes traffic according
# to its configured member weights. Inserts temporary nft counter rules,
# generates a test command using mwan3 tracking IPs, waits for the user to
# run it, then reports per-member actual vs expected hit counts.
#
# Usage: mwan3-lb-test [-6] -c <client_ip> <policy_name> [ip1 ip2 ...]
#        mwan3-lb-test cleanup
#
# -6            Test IPv6 traffic distribution instead of IPv4 (default).
#
# -c <client_ip>
#               IP address of the LAN client that will run the test pings
#               (required). Used to add isolation rules that block pings to
#               the test destination IPs from all other LAN clients and from
#               the router itself, preventing counter contamination.
#
# cleanup       Remove any stale mwan3_lb_test_* sets and associated rules
#               left behind by a run that was killed before it could clean up.
#
# Optional IP arguments override the default set of well-known public IPs
# used as ping destinations. Provide at least two distinct IPs. Use this
# when the well-known defaults are not reachable from your WAN links.
#
# Run on the router as root.

. /lib/functions.sh

TABLE="inet mwan3"
FORWARD_TABLE="inet fw4"
MWAN3TRACK_STATUS_DIR="/var/run/mwan3track"
MMX_MASK=$(uci -q get mwan3.globals.mmx_mask 2>/dev/null || echo "0x3f00")
HANDLES=""
TEST_RULE_HANDLE=""
FORWARD_BLOCK_HANDLE=""
OUTPUT_ISOLATION_HANDLE=""
TEST_SET="mwan3_lb_test_$$"
IPV6=0
CLIENT_IP=""
POLICY=""
MEMBERS=""
MEMBER_MARKS=""
MEMBER_WEIGHTS=""
MEMBER_IFACES=""

# ---------------------------------------------------------------------------
# Utilities
# ---------------------------------------------------------------------------

die() {
	echo "ERROR: $*" >&2
	cleanup
	exit 1
}

warn() {
	echo "WARNING: $*" >&2
}

cleanup() {
	# Remove isolation rules
	if [ -n "$OUTPUT_ISOLATION_HANDLE" ]; then
		nft delete rule $TABLE mwan3_output handle "$OUTPUT_ISOLATION_HANDLE" 2>/dev/null
		OUTPUT_ISOLATION_HANDLE=""
	fi
	if [ -n "$FORWARD_BLOCK_HANDLE" ]; then
		nft delete rule $FORWARD_TABLE forward handle "$FORWARD_BLOCK_HANDLE" 2>/dev/null
		FORWARD_BLOCK_HANDLE=""
	fi
	# Remove temporary test rule from mwan3_rules
	if [ -n "$TEST_RULE_HANDLE" ]; then
		nft delete rule $TABLE mwan3_rules handle "$TEST_RULE_HANDLE" 2>/dev/null
		TEST_RULE_HANDLE=""
	fi
	# Remove temporary destination sets (one per table)
	nft delete set $TABLE "$TEST_SET" 2>/dev/null
	nft delete set $FORWARD_TABLE "$TEST_SET" 2>/dev/null
	# Remove diagnostic counter rules from policy chain
	[ -z "$HANDLES" ] && return
	for h in $HANDLES; do
		nft delete rule $TABLE "mwan3_policy_${POLICY}" handle "$h" 2>/dev/null
	done
	HANDLES=""
}

trap 'echo; echo "Interrupted - cleaning up..."; cleanup; exit 1' INT TERM PIPE

# Append handle to HANDLES list
add_handle() {
	if [ -z "$HANDLES" ]; then
		HANDLES="$1"
	else
		HANDLES="$HANDLES $1"
	fi
}

# ---------------------------------------------------------------------------
# Argument check
# ---------------------------------------------------------------------------

# ---------------------------------------------------------------------------
# Cleanup-only invocation: mwan3-lb-test cleanup
#
# Removes any stale mwan3_lb_test_* sets and associated rules left behind by
# an aborted run (e.g. killed with SIGKILL before cleanup() could run).
# ---------------------------------------------------------------------------

if [ "$1" = "cleanup" ]; then
	found=0
	for stale_set in $(nft list sets inet 2>/dev/null \
			| awk '/mwan3_lb_test_/{gsub(/.*mwan3_lb_test_/,"mwan3_lb_test_"); gsub(/ \{.*/,""); print}'); do
		found=1
		echo "Removing stale set: $stale_set"
		stale_fwd=$(nft -a list chain $FORWARD_TABLE forward 2>/dev/null \
			| awk "/@${stale_set}.*drop/{print \$NF; exit}")
		[ -n "$stale_fwd" ] && {
			echo "  Removing forward isolation rule (handle $stale_fwd)"
			nft delete rule $FORWARD_TABLE forward handle "$stale_fwd" 2>/dev/null
		}
		stale_out=$(nft -a list chain $TABLE mwan3_output 2>/dev/null \
			| awk "/@${stale_set}.*return/{print \$NF; exit}")
		[ -n "$stale_out" ] && {
			echo "  Removing output isolation rule (handle $stale_out)"
			nft delete rule $TABLE mwan3_output handle "$stale_out" 2>/dev/null
		}
		stale_handle=$(nft -a list chain $TABLE mwan3_rules 2>/dev/null \
			| awk "/@${stale_set}/{print \$NF; exit}")
		[ -n "$stale_handle" ] && {
			echo "  Removing mwan3_rules test rule (handle $stale_handle)"
			nft delete rule $TABLE mwan3_rules handle "$stale_handle" 2>/dev/null
		}
		nft delete set $TABLE "$stale_set" 2>/dev/null
		nft delete set $FORWARD_TABLE "$stale_set" 2>/dev/null
	done
	[ "$found" -eq 0 ] && echo "Nothing to clean up."
	exit 0
fi

while true; do
	case "$1" in
		-6) IPV6=1; shift ;;
		-c)
			[ -z "$2" ] && { echo "ERROR: -c requires an IP address argument" >&2; exit 1; }
			CLIENT_IP="$2"; shift 2 ;;
		*) break ;;
	esac
done

if [ -z "$1" ]; then
	echo "Usage: $0 [-6] -c <client_ip> <policy_name> [ip1 ip2 ...]"
	echo "       $0 cleanup"
	echo
	echo "  -6              Test IPv6 traffic distribution (default: IPv4)"
	echo "  -c <client_ip>  IP address of the LAN client running the test pings (required)"
	echo "  cleanup         Remove stale rules/sets from an aborted run"
	echo
	echo "  Optional IP arguments override the default well-known public IPs."
	echo
	echo "Available load-balancing policies:"
	nft list chains inet 2>/dev/null \
		| awk '/mwan3_policy_/{gsub(/.*mwan3_policy_/,""); gsub(/ \{.*/,""); print}' \
		| sort -u \
		| while read -r name; do
			nft list chain inet mwan3 "mwan3_policy_${name}" 2>/dev/null \
				| grep -q "numgen" && echo "  $name"
		done
	exit 1
fi

[ -z "$CLIENT_IP" ] && die "client IP is required - use -c <ip>"

POLICY="$1"
shift
CHAIN="mwan3_policy_${POLICY}"
CUSTOM_IPS="$*"

# ---------------------------------------------------------------------------
# Pre-flight: chain exists
# ---------------------------------------------------------------------------

nft list chain $TABLE "$CHAIN" >/dev/null 2>&1 \
	|| die "chain $CHAIN not found - is mwan3 running?"

# ---------------------------------------------------------------------------
# Pre-flight: chain contains a numgen rule (load-balancing policy)
# ---------------------------------------------------------------------------

NUMGEN_LINE=$(nft list chain $TABLE "$CHAIN" 2>/dev/null | grep "numgen")
[ -n "$NUMGEN_LINE" ] \
	|| die "policy $POLICY has no numgen rule - it is not a load-balancing policy (single active member or all members offline)"

# ---------------------------------------------------------------------------
# Pre-flight: read members and weights from UCI
# ---------------------------------------------------------------------------

config_load mwan3

collect_member() {
	local member="$1"
	local iface weight
	config_get iface  "$member" interface ""
	config_get weight "$member" weight    1
	[ -z "$iface" ] && return
	MEMBERS="$MEMBERS $member"
	MEMBER_IFACES="$MEMBER_IFACES $iface"
	MEMBER_WEIGHTS="$MEMBER_WEIGHTS $weight"
}

config_list_foreach "$POLICY" use_member collect_member

MEMBERS="${MEMBERS# }"
MEMBER_IFACES="${MEMBER_IFACES# }"
MEMBER_WEIGHTS="${MEMBER_WEIGHTS# }"

[ -z "$MEMBERS" ] && die "policy $POLICY has no members in UCI config"

MEMBER_COUNT=$(echo "$MEMBER_WEIGHTS" | wc -w)
[ "$MEMBER_COUNT" -ge 2 ] || die "policy $POLICY has fewer than two members in UCI config"

# ---------------------------------------------------------------------------
# Pre-flight: parse member marks from live numgen vmap
#
# The vmap entries are in member order, matching UCI use_member list order.
# Each entry jumps to mwan3_or_meta_0xMARK; extract MARK from chain name.
# ---------------------------------------------------------------------------

MEMBER_MARKS=$(echo "$NUMGEN_LINE" \
	| grep -oE 'mwan3_or_meta_0x[0-9a-f]+' \
	| sed 's/mwan3_or_meta_//')

MARK_COUNT=$(echo "$MEMBER_MARKS" | wc -w)

# For mixed IPv4/IPv6 policies there may be two numgen lines; marks may repeat.
# Deduplicate while preserving order.
MEMBER_MARKS=$(echo "$MEMBER_MARKS" | tr ' ' '\n' | awk '!seen[$0]++' | tr '\n' ' ')
MEMBER_MARKS="${MEMBER_MARKS% }"
MARK_COUNT=$(echo "$MEMBER_MARKS" | wc -w)

[ "$MARK_COUNT" -lt 2 ] \
	&& die "could not parse at least two member marks from numgen vmap"

[ "$MARK_COUNT" -ne "$MEMBER_COUNT" ] \
	&& warn "mark count ($MARK_COUNT) differs from UCI member count ($MEMBER_COUNT) - mixed IPv4/IPv6 policy; per-member results may be approximate"

# Normalise marks to consistent hex format (no leading zeros beyond 0x prefix)
# so our inserted rule syntax and the awk extraction both use the same values.
MEMBER_MARKS_NORM=""
for mark in $MEMBER_MARKS; do
	norm=$(printf "0x%x" "$((mark))")
	MEMBER_MARKS_NORM="$MEMBER_MARKS_NORM $norm"
done
MEMBER_MARKS="${MEMBER_MARKS_NORM# }"

# ---------------------------------------------------------------------------
# Startup cleanup: remove stale mwan3_lb_test_* artifacts from aborted runs
# ---------------------------------------------------------------------------

for stale_set in $(nft list sets inet 2>/dev/null \
		| awk '/mwan3_lb_test_/{gsub(/.*mwan3_lb_test_/,"mwan3_lb_test_"); gsub(/ \{.*/,""); print}'); do
	stale_fwd=$(nft -a list chain $FORWARD_TABLE forward 2>/dev/null \
		| awk "/@${stale_set}.*drop/{print \$NF; exit}")
	[ -n "$stale_fwd" ] && \
		nft delete rule $FORWARD_TABLE forward handle "$stale_fwd" 2>/dev/null
	stale_out=$(nft -a list chain $TABLE mwan3_output 2>/dev/null \
		| awk "/@${stale_set}.*return/{print \$NF; exit}")
	[ -n "$stale_out" ] && \
		nft delete rule $TABLE mwan3_output handle "$stale_out" 2>/dev/null
	stale_handle=$(nft -a list chain $TABLE mwan3_rules 2>/dev/null \
		| awk "/@${stale_set}/{print \$NF; exit}")
	[ -n "$stale_handle" ] && \
		nft delete rule $TABLE mwan3_rules handle "$stale_handle" 2>/dev/null
	nft delete set $TABLE "$stale_set" 2>/dev/null
	nft delete set $FORWARD_TABLE "$stale_set" 2>/dev/null
done

# ---------------------------------------------------------------------------
# Pre-flight: warn if other rules use this policy
# ---------------------------------------------------------------------------

OTHER_RULES=$(nft list chain $TABLE mwan3_rules 2>/dev/null \
	| grep "jump ${CHAIN}" \
	| grep -v "^[[:space:]]*#")

RULE_COUNT=$(echo "$OTHER_RULES" | grep -c "jump ${CHAIN}" 2>/dev/null)
if [ "$RULE_COUNT" -gt 1 ]; then
	warn "$RULE_COUNT rules send traffic to policy $POLICY."
	warn "flows from rules other than your test rule will also increment the counter."
	warn "consider temporarily disabling other rules using this policy before testing."
	echo
fi

# ---------------------------------------------------------------------------
# Compute N
#
# NITER = (total_weight / GCD(all_weights)) * multiplier
# where multiplier = ceil(30 / base_N), giving NITER >= 30.
# This ensures NITER * weight_i / total_weight is an integer for every member.
# ---------------------------------------------------------------------------

TOTAL_WEIGHT=0
for w in $MEMBER_WEIGHTS; do
	TOTAL_WEIGHT=$((TOTAL_WEIGHT + w))
done

GCD_WEIGHTS=$(echo "$MEMBER_WEIGHTS" | tr ' ' '\n' | \
	awk 'BEGIN{g=0} {
		a=$1; b=g;
		while(b!=0){t=b; b=a%b; a=t}
		g=a
	} END{print g}')

BASE_N=$((TOTAL_WEIGHT / GCD_WEIGHTS))
MULTIPLIER=$(( (30 + BASE_N - 1) / BASE_N ))
NITER=$((BASE_N * MULTIPLIER))

# ---------------------------------------------------------------------------
# Collect test destination IPs
#
# Use well-known public IPs that should be reachable from any WAN link.
# Exclude any IPs that mwan3track is already pinging: those pings go through
# mwan3_output -> mwan3_rules and would match the test rule, contaminating
# the counter with router-originated traffic.
# ---------------------------------------------------------------------------

if [ -n "$CUSTOM_IPS" ]; then
	TRACK_IPS="$CUSTOM_IPS"
	CUSTOM_IP_COUNT=$(echo "$TRACK_IPS" | wc -w)
	[ "$CUSTOM_IP_COUNT" -lt 2 ] && die "at least two destination IPs are required (got $CUSTOM_IP_COUNT)"
else
	# Collect all track_ip values configured across all mwan3 interfaces.
	MWAN3_TRACKED_IPS=""
	collect_iface_track_ips() {
		local iface_track_ips
		config_get iface_track_ips "$1" track_ip ""
		MWAN3_TRACKED_IPS="$MWAN3_TRACKED_IPS $iface_track_ips"
	}
	config_foreach collect_iface_track_ips interface

	# Default pools - exclude IPs that mwan3track is already pinging
	# (those go through mwan3_output -> mwan3_rules and contaminate the counter).
	if [ "$IPV6" -eq 1 ]; then
		DEFAULT_POOL="2001:4860:4860::8888 2001:4860:4860::8844 2606:4700:4700::1111 2606:4700:4700::1001 2620:fe::fe 2620:fe::9 2620:119:35::35 2620:119:53::53"
	else
		DEFAULT_POOL="8.8.8.8 8.8.4.4 1.1.1.1 1.0.0.1 9.9.9.9 149.112.112.112 208.67.222.222 208.67.220.220 4.2.2.1 4.2.2.2 185.228.168.168 185.228.169.168"
	fi

	TRACK_IPS=""
	for ip in $DEFAULT_POOL; do
		case " $MWAN3_TRACKED_IPS " in
			*" $ip "*) ;;
			*) TRACK_IPS="$TRACK_IPS $ip" ;;
		esac
	done
	TRACK_IPS="${TRACK_IPS# }"

	[ -z "$TRACK_IPS" ] && die "all default destination IPs are configured as mwan3 tracking IPs - use the IP override argument to specify test destinations not being tracked"
fi
TRACK_COUNT=$(echo "$TRACK_IPS" | wc -w)

# Compute Windows cmd.exe delay between pings.
# Windows ping uses a fixed ICMP identifier (id=1) so repeated pings to the
# same destination reuse the same conntrack entry. The ct mark from the first
# ping is restored on the second, bypassing the test rule (mark != 0 guard).
# The full cycle through TRACK_COUNT IPs (each taking up to 2s for the ping
# plus WIN_DELAY seconds) must exceed the 30s ICMP conntrack timeout so the
# entry expires before we revisit each IP.
# Formula: WIN_DELAY such that TRACK_COUNT * (2 + WIN_DELAY) > 30.
# Using WIN_DELAY = 30/TRACK_COUNT + 3 gives comfortable margin.
WIN_DELAY=$(( 30 / TRACK_COUNT + 3 ))

# Build the rotated destination array for the generated command.
# Each iteration pings TRACK_IPS[i % TRACK_COUNT].
# Expand to exactly N destinations, rotating through the pool.
DEST_LIST=""
i=0
for seq in $(seq 1 "$NITER"); do
	idx=$((i % TRACK_COUNT + 1))
	dst=$(echo "$TRACK_IPS" | tr ' ' '\n' | sed -n "${idx}p")
	DEST_LIST="$DEST_LIST $dst"
	i=$((i + 1))
done
DEST_LIST="${DEST_LIST# }"

# ---------------------------------------------------------------------------
# Compute expected per-member hits
# ---------------------------------------------------------------------------

EXPECTED_LIST=""
for w in $MEMBER_WEIGHTS; do
	expected=$((NITER * w / TOTAL_WEIGHT))
	EXPECTED_LIST="$EXPECTED_LIST $expected"
done
EXPECTED_LIST="${EXPECTED_LIST# }"

# ---------------------------------------------------------------------------
# Insert temporary test rule into mwan3_rules
#
# Creates a temporary nft set containing the test destination IPs and inserts
# a rule at the head of mwan3_rules that directs matching traffic to the
# policy under test. This ensures the test pings hit the policy chain without
# requiring the user to add a permanent rule. Both the rule and the set are
# removed by cleanup().
# ---------------------------------------------------------------------------

echo "Adding temporary test set, isolation rules, and test rule..."
echo

if [ "$IPV6" -eq 1 ]; then
	SET_TYPE="ipv6_addr"
	PROTO_MATCH="meta l4proto ipv6-icmp ip6 daddr"
	PING_CMD="ping6"
	FORWARD_BLOCK="ip6 saddr != $CLIENT_IP ip6 daddr @$TEST_SET meta l4proto ipv6-icmp icmpv6 type echo-request"
	OUTPUT_ISOLATION="meta l4proto ipv6-icmp icmpv6 type echo-request ip6 daddr @$TEST_SET"
else
	SET_TYPE="ipv4_addr"
	PROTO_MATCH="meta l4proto icmp ip daddr"
	PING_CMD="ping"
	FORWARD_BLOCK="ip saddr != $CLIENT_IP ip daddr @$TEST_SET meta l4proto icmp icmp type echo-request"
	OUTPUT_ISOLATION="meta l4proto icmp icmp type echo-request ip daddr @$TEST_SET"
fi

# The rule matches ICMP/ICMPv6 only - not DNS, TCP, or other traffic to these
# IPs. This prevents contamination from DNS forwarders and rogue clients that
# bypass the local resolver, without needing to avoid popular DNS server IPs.
#
# nft sets are table-scoped. The mwan3_output and mwan3_rules rules live in
# table inet mwan3; the forward isolation rule lives in table inet fw4.
# Create the set in both tables and populate both with the same IPs.
nft add set $TABLE "$TEST_SET" { type "$SET_TYPE" \; } \
	|| die "failed to create temporary destination set $TEST_SET in $TABLE"
nft add set $FORWARD_TABLE "$TEST_SET" { type "$SET_TYPE" \; } \
	|| die "failed to create temporary destination set $TEST_SET in $FORWARD_TABLE"

for ip in $TRACK_IPS; do
	nft add element $TABLE "$TEST_SET" { "$ip" } 2>/dev/null
	nft add element $FORWARD_TABLE "$TEST_SET" { "$ip" } 2>/dev/null
done

# Insert forward isolation rule: drop pings to test IPs from all LAN clients
# except the nominated test client. Scoped to the test set so other LAN traffic
# is unaffected.
nft insert rule $FORWARD_TABLE forward $FORWARD_BLOCK drop \
	|| die "failed to insert forward isolation rule"
FORWARD_BLOCK_HANDLE=$(nft -a list chain $FORWARD_TABLE forward 2>/dev/null \
	| awk "/@${TEST_SET}.*drop/{print \$NF; exit}")
[ -n "$FORWARD_BLOCK_HANDLE" ] || die "could not read forward isolation rule handle"

# Insert output isolation rule: return (skip mwan3 marking) for any router
# process pinging test IPs. mwan3track IPs are already excluded from the test
# set, so this only affects unrelated processes. 'return' in a base chain is
# equivalent to accept - the pings still work, they just don't hit the policy
# chain.
nft insert rule $TABLE mwan3_output $OUTPUT_ISOLATION return \
	|| die "failed to insert output isolation rule"
OUTPUT_ISOLATION_HANDLE=$(nft -a list chain $TABLE mwan3_output 2>/dev/null \
	| awk "/@${TEST_SET}.*return/{print \$NF; exit}")
[ -n "$OUTPUT_ISOLATION_HANDLE" ] || die "could not read output isolation rule handle"

nft insert rule $TABLE mwan3_rules \
	$PROTO_MATCH @"$TEST_SET" meta mark \& "$MMX_MASK" == 0 \
	jump "mwan3_policy_${POLICY}" \
	|| die "failed to insert temporary test rule into mwan3_rules"

TEST_RULE_HANDLE=$(nft -a list chain $TABLE mwan3_rules 2>/dev/null \
	| awk "/@${TEST_SET}/{print \$NF; exit}")
[ -n "$TEST_RULE_HANDLE" ] || die "could not read temporary test rule handle"

# ---------------------------------------------------------------------------
# Print summary
# ---------------------------------------------------------------------------

echo "=== mwan3 load balancing distribution test ==="
echo
echo "Policy          : $POLICY"
echo "MMX mask        : $MMX_MASK"
echo "IP family       : $([ "$IPV6" -eq 1 ] && echo IPv6 || echo IPv4)"
echo "Test iterations : $NITER  (base=$BASE_N x${MULTIPLIER}, total_weight=$TOTAL_WEIGHT, GCD=$GCD_WEIGHTS)"
if [ -n "$CUSTOM_IPS" ]; then
	echo "Destinations    : $TRACK_IPS  (user override)"
elif [ -n "$MWAN3_TRACKED_IPS" ]; then
	echo "Destinations    : $TRACK_IPS  (well-known IPs, excluding mwan3 tracking IPs)"
else
	echo "Destinations    : $TRACK_IPS  (well-known IPs)"
fi
echo
printf "%-20s %-10s %-10s %-16s %s\n" "Member" "Interface" "Weight" "Mark" "Expected hits"
printf "%-20s %-10s %-10s %-16s %s\n" "------" "---------" "------" "----" "-------------"

i=1
for member in $MEMBERS; do
	iface=$(echo "$MEMBER_IFACES" | tr ' ' '\n' | sed -n "${i}p")
	weight=$(echo "$MEMBER_WEIGHTS" | tr ' ' '\n' | sed -n "${i}p")
	mark=$(echo "$MEMBER_MARKS" | tr ' ' '\n' | sed -n "${i}p")
	expected=$(echo "$EXPECTED_LIST" | tr ' ' '\n' | sed -n "${i}p")
	printf "%-20s %-10s %-10s %-16s %s\n" "$member" "$iface" "$weight" "${mark:--}" "${expected:--}"
	i=$((i + 1))
done
echo

# ---------------------------------------------------------------------------
# Insert diagnostic rules
#
# Insertion order (each insert prepends, so insert in reverse of desired order):
#   desired: [unguarded] [guarded] [numgen] [member counters...] [last_resort]
#   insert:  guarded first (becomes head), then unguarded (pushes guarded to #2)
#   member counters inserted before last_resort handle using 'add rule position'
# ---------------------------------------------------------------------------

echo "Adding diagnostic rules..."

# Guarded counter (mark==0) - inserted at head
nft insert rule $TABLE "$CHAIN" meta mark \& "$MMX_MASK" == 0 counter \
	|| die "failed to add guarded counter"

GUARDED_HANDLE=$(nft -a list chain $TABLE "$CHAIN" 2>/dev/null \
	| awk '/meta mark.*== 0x.*counter/{print $NF; exit}')
[ -n "$GUARDED_HANDLE" ] || die "could not read guarded counter handle"
add_handle "$GUARDED_HANDLE"

# Unguarded counter - inserted at head (pushes guarded to position 2)
nft insert rule $TABLE "$CHAIN" counter \
	|| die "failed to add unguarded counter"

UNGUARDED_HANDLE=$(nft -a list chain $TABLE "$CHAIN" 2>/dev/null \
	| awk '/^\t\tcounter packets/{print $NF; exit}')
[ -n "$UNGUARDED_HANDLE" ] || die "could not read unguarded counter handle"
add_handle "$UNGUARDED_HANDLE"

# Per-member counters - inserted before last_resort (last rule in chain)
# Use 'nft add rule position <handle>' which inserts AFTER the given handle.
# We insert after the last numgen rule handle so counters sit between numgen
# and last_resort regardless of how many numgen lines the chain has.
NUMGEN_HANDLE=$(nft -a list chain $TABLE "$CHAIN" 2>/dev/null \
	| awk '/numgen/{last=$NF} END{print last}')
[ -n "$NUMGEN_HANDLE" ] || die "could not read numgen rule handle"

PREV_HANDLE="$NUMGEN_HANDLE"
for mark in $MEMBER_MARKS; do
	nft add rule $TABLE "$CHAIN" position "$PREV_HANDLE" \
		meta mark \& "$MMX_MASK" == "$mark" counter \
		|| die "failed to add per-member counter for mark $mark"

	# Read handle of the rule just inserted (it follows PREV_HANDLE)
	NEW_HANDLE=$(nft -a list chain $TABLE "$CHAIN" 2>/dev/null \
		| awk -v prev="$PREV_HANDLE" '
			/handle/{
				if (found) { print $NF; exit }
				if ($NF == prev) found=1
			}')
	[ -n "$NEW_HANDLE" ] || die "could not read per-member counter handle for mark $mark"
	add_handle "$NEW_HANDLE"
	PREV_HANDLE="$NEW_HANDLE"
done

echo "Done. Current chain state:"
nft list chain $TABLE "$CHAIN"
echo

# ---------------------------------------------------------------------------
# Print test command
# ---------------------------------------------------------------------------

echo "----------------------------------------------------------------------"
echo
echo "Run the following on a LAN client (or on the router itself to test"
echo "the output path):"
echo

# Print DEST_LIST wrapped at 4 IPs per line with \ continuation
printf "  for dst in"
col=0
for dst in $DEST_LIST; do
	if [ "$col" -eq 0 ]; then
		printf " \\\\\n    %s" "$dst"
	else
		printf " %s" "$dst"
	fi
	col=$(( (col + 1) % 4 ))
done
echo "; do"
echo "    echo -n \"\$dst: \""
echo "    $PING_CMD -c1 -W2 \"\$dst\" >/dev/null 2>&1 && echo ok || echo timeout"
echo "    sleep 2"
echo "  done"
echo
echo "If testing from a Windows client cmd.exe (uses fixed ICMP id - requires"
echo "longer inter-ping delay so conntrack entries expire between revisits):"
echo
if [ "$IPV6" -eq 1 ]; then
	WIN_PING="ping -6 -n 1 -w 2000"
else
	WIN_PING="ping -n 1 -w 2000"
fi
# Print Windows FOR loop with ^ continuation at 4 IPs per line
printf "  for %%i in (^\n"
col=0
total_dsts=$(echo "$DEST_LIST" | wc -w)
cur=0
for dst in $DEST_LIST; do
	cur=$((cur + 1))
	if [ "$col" -eq 0 ]; then
		printf "    %s" "$dst"
	else
		printf " %s" "$dst"
	fi
	col=$(( (col + 1) % 4 ))
	if [ "$col" -eq 0 ] && [ "$cur" -lt "$total_dsts" ]; then
		printf " ^\n"
	fi
done
printf ") do (%s %%i & timeout /t %d /nobreak >nul)\n" "$WIN_PING" "$WIN_DELAY"
echo
echo "Press Enter here when the test is complete..."
echo
echo "----------------------------------------------------------------------"
read _

# ---------------------------------------------------------------------------
# Read results
# ---------------------------------------------------------------------------

echo
echo "=== Results ==="
echo
CHAIN_DUMP=$(nft list chain $TABLE "$CHAIN" 2>/dev/null)

TOTAL=$(echo "$CHAIN_DUMP" \
	| awk '/^\t\tcounter packets/{for(i=1;i<=NF;i++) if($i=="packets"){print $(i+1); exit}}')
NEW=$(echo "$CHAIN_DUMP" \
	| awk '/meta mark.*== 0x.*counter/{for(i=1;i<=NF;i++) if($i=="packets"){print $(i+1); exit}}')

# Read per-member counts. The member counter rules match 'meta mark & X == MARK counter'.
# They appear after the numgen rule. Extract in mark order.
MEMBER_ACTUALS=""
for mark in $MEMBER_MARKS; do
	count=$(echo "$CHAIN_DUMP" \
		| awk -v mark="$mark" '
			/meta mark.*== .*counter/{
				for(i=1;i<=NF;i++){
					if($i=="==" && strtonum($(i+1))==strtonum(mark) && strtonum(mark)!=0){
						for(j=i;j<=NF;j++){
							if($j=="packets"){print $(j+1); exit}
						}
					}
				}
			}')
	MEMBER_ACTUALS="$MEMBER_ACTUALS ${count:-0}"
done
MEMBER_ACTUALS="${MEMBER_ACTUALS# }"

# ---------------------------------------------------------------------------
# Report
# ---------------------------------------------------------------------------

echo "--- Summary ---"
echo
printf "%-40s %s\n" "Completed:" "$(date '+%Y-%m-%d %H:%M:%S')"
printf "%-40s %s\n" "Total packets entering chain:" "${TOTAL:-unknown}"
printf "%-40s %s\n" "New-connection packets (mark=0):" "${NEW:-unknown}"
printf "%-40s %s\n" "Test iterations:" "$NITER"
echo

FW4_RELOADED=0
if [ -z "$TOTAL" ]; then
	FW4_RELOADED=1
	warn "Counter rules are absent from $CHAIN."
	warn "fw4 reload likely occurred during the test, wiping all dynamic nft rules."
	warn "Re-run the test and return before fw4 triggers (e.g. within a few minutes)."
	echo
fi

CONTAMINATED=0
if [ -n "$TOTAL" ] && [ -n "$NEW" ]; then
	if [ "$TOTAL" -gt "$NEW" ] 2>/dev/null; then
		STALE=$((TOTAL - NEW))
		warn "$STALE packets had mark != 0 when entering the chain."
		warn "Established-connection packets should not reach this chain."
		warn "Per-member counts below may be inflated - check for prerouting early-exit bugs."
		echo
		CONTAMINATED=1
	fi
fi

echo "--- Per-member distribution ---"
echo
printf "%-20s %-10s %-10s %-10s %-10s\n" "Member" "Interface" "Expected" "Actual" "Deviation"
printf "%-20s %-10s %-10s %-10s %-10s\n" "------" "---------" "--------" "------" "---------"

PASS=1
i=1
for member in $MEMBERS; do
	iface=$(echo "$MEMBER_IFACES" | tr ' ' '\n' | sed -n "${i}p")
	expected=$(echo "$EXPECTED_LIST" | tr ' ' '\n' | sed -n "${i}p")
	actual=$(echo "$MEMBER_ACTUALS" | tr ' ' '\n' | sed -n "${i}p")
	[ -z "$actual" ] && actual=0
	deviation=$((actual - expected))
	dev_str="$deviation"
	[ "$deviation" -gt 0 ] && dev_str="+$deviation"
	abs_dev=$deviation
	[ "$abs_dev" -lt 0 ] && abs_dev=$((-abs_dev))
	[ "$abs_dev" -gt 2 ] && PASS=0
	printf "%-20s %-10s %-10s %-10s %-10s\n" "$member" "$iface" "$expected" "$actual" "$dev_str"
	i=$((i + 1))
done
echo

if [ -n "$NEW" ] && [ "$NEW" -ne "$NITER" ] 2>/dev/null; then
	EXTRA=$((NEW - NITER))
	if [ "$EXTRA" -gt 0 ]; then
		echo "FAIL: $EXTRA extra new-connection hits beyond $NITER test iterations."
		echo "      Something other than your test flows is reaching this policy chain."
		echo "      Check whether another rule also sends traffic to this policy, or"
		echo "      whether a process on the router is generating connections to the test IPs."
		PASS=0
	else
		echo "NOTE: Fewer hits than expected ($NEW/$NITER)."
		echo "      Some test flows may have timed out or been routed by a higher-priority rule."
	fi
	echo
fi

if [ "$FW4_RELOADED" -eq 1 ]; then
	echo "FAIL: test invalidated - fw4 reload wiped counter rules during the wait."
	echo "      Re-run the test and do not leave it unattended for more than a few minutes."
elif [ "$PASS" -eq 1 ] && [ "$CONTAMINATED" -eq 0 ]; then
	echo "PASS: distribution matches expected weights within tolerance (+/-2 per member)."
	echo "      Load balancing is working according to configured policy."
else
	echo "FAIL: distribution is outside tolerance."
	echo "      Actual counts deviate from expected by more than 2 for at least one member."
fi

echo
echo "Cleaning up diagnostic rules..."
cleanup
echo "Chain restored."
