#!/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

[ -f "$CONF" ] || {
    echo "[wifi-bssid] Missing config: $CONF"
    exit 1
}

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")"

[ -n "${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.service
}

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    bgscan=\"\"", block)
                }

                sub(/ssid="[^"]*"/, "&\n    bssid=" 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
