#!/bin/sh
set -eu

CONFIG="/etc/default/wifi-bssid-failsafe"

IFACE="wlan0"
CONF="/etc/wpa_supplicant/wpa_supplicant-wlan0.conf"
PING_HOST="8.8.8.8"
STATE_DIR="/run/wifi-bssid-failsafe"
FAIL_FILE="$STATE_DIR/fail-count"
BAD_DIR="$STATE_DIR/bad-bssid"
MAX_FAILS="3"
MAX_BSSID_TRIES="3"
BAD_BSSID_TTL="300"
CONNECT_WAIT="10"

[ -f "$CONFIG" ] && . "$CONFIG"

EVENT="${2:-manual}"

case "$EVENT" in
	CONNECTED|DISCONNECTED|manual)
		echo "[wifi-bssid] Handling event: $EVENT"
		;;
	*)
		exit 0
		;;
esac

SSID="$(awk '
/^[ \t]*network=\{/ { in_block=1 }
/^[ \t]*ssid="/ && in_block {
	line=$0
	sub(/^[ \t]*ssid="/, "", line)
	sub(/".*/, "", line)
	print line
	exit
}
/^[ \t]*\}/ { in_block=0 }
' "$CONF")"

[ -z "${SSID:-}" ] && {
	echo "[wifi-bssid] No SSID found"
	exit 1
}

mkdir -p "$STATE_DIR" "$BAD_DIR"

now() { date +%s; }

is_online() {
	iw dev "$IFACE" link 2>/dev/null | grep -q "Connected to" &&
	ping -I "$IFACE" -c 1 -W 3 "$PING_HOST" >/dev/null 2>&1
}

get_current_bssid() {
	iw dev "$IFACE" link 2>/dev/null | awk '/Connected to/ {print $3}'
}

safe_bssid_name() {
	echo "$1" | tr ':A-F' '-a-f'
}

mark_bad_bssid() {
	BSSID="$1"
	FILE="$BAD_DIR/$(safe_bssid_name "$BSSID")"
	echo "$(now)" > "$FILE"
	echo "[wifi-bssid] Marked bad: $BSSID"
}

is_bad_bssid() {
	BSSID="$1"
	FILE="$BAD_DIR/$(safe_bssid_name "$BSSID")"

	[ -f "$FILE" ] || return 1

	TS="$(cat "$FILE" 2>/dev/null || echo 0)"
	AGE=$(( $(now) - TS ))

	if [ "$AGE" -lt "$BAD_BSSID_TTL" ]; then
		return 0
	fi

	rm -f "$FILE"
	return 1
}

clear_bad_bssid() {
	BSSID="$1"
	FILE="$BAD_DIR/$(safe_bssid_name "$BSSID")"
	rm -f "$FILE"
}

pick_bssids() {
	iw dev "$IFACE" scan 2>/dev/null | awk -v ssid="$SSID" '
	/^BSS / {
		bssid=$2
		sub(/\(.*/, "", bssid)
		signal=""
	}
	/signal:/ { signal=$2 }
	/SSID:/ {
		name=$0
		sub(/^[ \t]*SSID: /, "", name)
		if (name == ssid && signal != "") {
			print signal, bssid
		}
	}' | sort -nr | awk '{print $2}'
}

get_config_bssid() {
	awk '
	/^[ \t]*network=\{/ { in_block=1 }
	in_block && /^[ \t]*bssid=/ {
		sub(/^[ \t]*bssid=/, "", $0)
		print
		exit
	}
	/^[ \t]*\}/ { in_block=0 }
	' "$CONF"
}

reset_fail_count() { echo 0 > "$FAIL_FILE"; }

inc_fail_count() {
	COUNT="$(cat "$FAIL_FILE" 2>/dev/null || echo 0)"
	COUNT=$((COUNT + 1))
	echo "$COUNT" > "$FAIL_FILE"
	echo "$COUNT"
}

reconfigure_wifi() {
	wpa_cli -i "$IFACE" reconfigure >/dev/null 2>&1 || systemctl restart wpa_supplicant
}

update_bssid() {
	NEW_BSSID="$1"
	CURRENT_CFG_BSSID="$(get_config_bssid || true)"

	[ "$CURRENT_CFG_BSSID" = "$NEW_BSSID" ] && return

	cp "$CONF" "$CONF.bssid-backup"

	awk -v ssid="$SSID" -v bssid="$NEW_BSSID" '
	BEGIN { in_block=0; matched=0 }

	/^[ \t]*network=\{/ { in_block=1; block="" }

	in_block {
		block = block $0 "\n"

		if ($0 ~ /^[ \t]*ssid="/) {
			line=$0
			sub(/^[ \t]*ssid="/, "", line)
			sub(/".*/, "", line)
			if (line == ssid) matched=1
		}

		if ($0 ~ /^[ \t]*\}/) {
			if (matched) {
				gsub(/\n[ \t]*bssid=.*\n/, "\n", block)

				if (block !~ /bgscan=""/) {
					sub(/ssid="[^"]*"/, "&\n\tbgscan=\"\"", block)
				}

				sub(/ssid="[^"]*"/, "&\n\tbssid=" bssid, block)
			}

			printf "%s", block
			in_block=0
			matched=0
			next
		}

		next
	}

	{ print }
	' "$CONF" > "$CONF.tmp"

	mv "$CONF.tmp" "$CONF"

	reconfigure_wifi
}

remove_bssid() {
	cp "$CONF" "$CONF.bssid-backup"

	awk -v ssid="$SSID" '
	BEGIN { in_block=0; matched=0 }

	/^[ \t]*network=\{/ { in_block=1; block="" }

	in_block {
		block = block $0 "\n"

		if ($0 ~ /^[ \t]*ssid="/) {
			line=$0
			sub(/^[ \t]*ssid="/, "", line)
			sub(/".*/, "", line)
			if (line == ssid) matched=1
		}

		if ($0 ~ /^[ \t]*\}/) {
			if (matched) {
				gsub(/\n[ \t]*bssid=.*\n/, "\n", block)
			}

			printf "%s", block
			in_block=0
			matched=0
			next
		}

		next
	}

	{ print }
	' "$CONF" > "$CONF.tmp"

	mv "$CONF.tmp" "$CONF"

	reconfigure_wifi
}

CURRENT_BSSID="$(get_current_bssid || true)"

if [ -n "$CURRENT_BSSID" ] && is_online; then
	update_bssid "$CURRENT_BSSID"
	clear_bad_bssid "$CURRENT_BSSID"
	reset_fail_count
	exit 0
fi

FAILS="$(inc_fail_count)"

if [ "$FAILS" -ge "$MAX_FAILS" ]; then
	remove_bssid
	reset_fail_count
fi

BSSID_LIST="$STATE_DIR/list"
pick_bssids > "$BSSID_LIST"

TRIED=0

while read -r BSSID; do
	[ -n "$BSSID" ] || continue

	if is_bad_bssid "$BSSID"; then
		continue
	fi

	TRIED=$((TRIED + 1))
	[ "$TRIED" -gt "$MAX_BSSID_TRIES" ] && break

	update_bssid "$BSSID"
	sleep "$CONNECT_WAIT"

	if is_online; then
		clear_bad_bssid "$BSSID"
		reset_fail_count
		exit 0
	fi

	mark_bad_bssid "$BSSID"

done < "$BSSID_LIST"

exit 1
