#!/bin/bash
set -e

get_band() {
    local freq="$1"
    local f=${freq%%.*}
    if [ -z "$f" ]; then    
        echo "未知"
    elif [ "$f" -ge 5925 ] 2>/dev/null; then
        echo "6G"
    elif [ "$f" -ge 4900 ] 2>/dev/null; then
        echo "5G"
    elif [ "$f" -ge 2400 ] 2>/dev/null && [ "$f" -lt 2500 ] 2>/dev/null; then
        echo "2.4G"
    else
        echo "未知"
    fi
}

echo "========== WiFi WPA2 一键连接脚本 =========="

# 自动检测无线网卡
WL_IFACES=$(nmcli device status | awk '$2=="wifi"{print $1}')

if [ -z "$WL_IFACES" ]; then
    echo "❌ 没有检测到无线网卡"
    echo "当前网卡列表："
    ip -br link
    exit 1
fi

WL_COUNT=$(echo "$WL_IFACES" | wc -l)

if [ "$WL_COUNT" -eq 1 ]; then
    IFACE="$WL_IFACES"
else
    echo "检测到多个无线网卡，请选择："
    i=1
    echo "$WL_IFACES" | while read -r w; do
        echo "$i) $w"
        i=$((i+1))
    done
    read -r -p "请输入无线网卡序号: " SEL
    IFACE=$(echo "$WL_IFACES" | sed -n "${SEL}p")
fi

if [ -z "$IFACE" ]; then
    echo "❌ 无线网卡选择无效"
    exit 1
fi

echo ">>> 使用无线网卡: $IFACE"

sudo ip link set "$IFACE" up || true

echo ">>> 扫描 WiFi ..."
sudo nmcli device wifi rescan ifname "$IFACE" >/dev/null 2>&1 || true
sleep 2

TMP_FILE="/tmp/wifi_list_$$.txt"

# 字段顺序：SSID,BSSID,FREQ,CHAN,SIGNAL,SECURITY
# nmcli -t 会把 BSSID 里的冒号转义为 \:，awk 分隔后 BSSID 占用第 2~7 列
# 因此：FREQ=$8, CHAN=$9, SIGNAL=$10, SECURITY=$11+
nmcli -t -f SSID,BSSID,FREQ,CHAN,SIGNAL,SECURITY device wifi list ifname "$IFACE" \
    | awk -F: '$1!="" {print $10 "\t" $0}' \
    | sort -k1,1nr \
    | cut -f2- \
    > "$TMP_FILE"

if [ ! -s "$TMP_FILE" ]; then
    echo "❌ 没有扫描到 WiFi"
    rm -f "$TMP_FILE"
    exit 1
fi

echo ""
echo "========== 可用 WiFi 列表 =========="
printf "%3s  %-28s %-17s %-5s %-10s %-6s %-8s %s\n" "序号" "SSID" "BSSID" "频段" "频率" "信道" "信号" "加密"
echo "------------------------------------------------------------------------------------------------"

awk -F: '
{
    ssid=$1
    bssid=$2":"$3":"$4":"$5":"$6":"$7
    gsub(/\\:/, ":", bssid)
    freq=$8
    chan=$9
    signal=$10
    security=$11
    for (i=12; i<=NF; i++) security=security":"$i

    band="未知"
    if (freq+0 >= 5925) band="6G"
    else if (freq+0 >= 4900) band="5G"
    else if (freq+0 >= 2400 && freq+0 < 2500) band="2.4G"

    printf "%3d. %-28s %-17s %-5s %-10s CH%-4s %-8s %s\n", NR, ssid, bssid, band, freq"MHz", chan, signal"%", security
}' "$TMP_FILE"

echo ""
echo "说明：2400MHz左右=2.4G，5000MHz左右=5G，5925MHz以上=6G"
echo ""
read -r -p "请输入要连接的 WiFi 序号: " IDX

LINE=$(sed -n "${IDX}p" "$TMP_FILE")

if [ -z "$LINE" ]; then
    echo "❌ 序号无效"
    rm -f "$TMP_FILE"
    exit 1
fi

SSID=$(echo "$LINE" | awk -F: '{print $1}')
BSSID=$(echo "$LINE" | awk -F: '{print $2":"$3":"$4":"$5":"$6":"$7}' | sed 's/\\:/:/g')
FREQ=$(echo "$LINE" | awk -F: '{print $8}')
CHAN=$(echo "$LINE" | awk -F: '{print $9}')
SIGNAL=$(echo "$LINE" | awk -F: '{print $10}')
SECURITY=$(echo "$LINE" | cut -d: -f11-)
BAND=$(get_band "$FREQ")

if [ -z "$SSID" ]; then
    echo "❌ SSID 为空，无法连接"
    rm -f "$TMP_FILE"
    exit 1
fi

echo ""
echo "你选择的是："
echo "SSID: $SSID"
echo "BSSID: $BSSID"
echo "频段: $BAND"
echo "频率: ${FREQ}MHz"
echo "信道: CH$CHAN"
echo "信号: ${SIGNAL}%"
echo "SECURITY: $SECURITY"
echo ""

read -r -s -p "请输入 WiFi 密码: " WIFI_PASS
echo ""

if [ -z "$WIFI_PASS" ]; then
    echo "❌ 密码不能为空"
    rm -f "$TMP_FILE"
    exit 1
fi

CON_NAME="$SSID"

echo ""
echo ">>> 删除旧的同名连接配置 ..."

sudo nmcli connection delete "$CON_NAME" >/dev/null 2>&1 || true

nmcli -t -f NAME,TYPE connection show | grep ":802-11-wireless" | cut -d: -f1 | while read -r c; do
    OLD_SSID=$(nmcli -g 802-11-wireless.ssid connection show "$c" 2>/dev/null || true)
    if [ "$OLD_SSID" = "$SSID" ]; then
        sudo nmcli connection delete "$c" >/dev/null 2>&1 || true
    fi
done

echo ">>> 创建 WPA2-PSK WiFi 配置 ..."

sudo nmcli connection add \
    type wifi \
    ifname "$IFACE" \
    con-name "$CON_NAME" \
    ssid "$SSID" >/dev/null

sudo nmcli connection modify "$CON_NAME" \
    802-11-wireless.mode infrastructure \
    802-11-wireless-security.key-mgmt wpa-psk \
    802-11-wireless-security.psk "$WIFI_PASS" \
    802-11-wireless-security.psk-flags 0 \
    connection.autoconnect yes \
    connection.interface-name "$IFACE"

if [ -n "$BSSID" ]; then
    echo ">>> 绑定 BSSID: $BSSID ($BAND / ${FREQ}MHz / CH$CHAN)"
    sudo nmcli connection modify "$CON_NAME" 802-11-wireless.bssid "$BSSID" || true
fi

echo ">>> 正在连接 WiFi ..."

if sudo nmcli connection up "$CON_NAME"; then
    sleep 1
    CURRENT_WIFI=$(nmcli -g GENERAL.CONNECTION device show "$IFACE")
    IP_ADDR=$(ip -4 -br addr show "$IFACE" | awk '{print $3}' | cut -d/ -f1)
    LINK_INFO=$(iw dev "$IFACE" link 2>/dev/null | grep -E "SSID|freq|signal" || true)

    echo ""
    echo "✅ WiFi 连接成功"
    echo "当前WIFI：$CURRENT_WIFI"
    echo "IP地址：$IP_ADDR"
    echo "连接频段：$BAND / ${FREQ}MHz / CH$CHAN"
    if [ -n "$LINK_INFO" ]; then
        echo ""
        echo "当前链路信息："
        echo "$LINK_INFO"
    fi
else
    echo ""
    echo "❌ WiFi 连接失败"
    echo "请检查：WiFi密码、路由器WPA模式或MAC白名单"
    rm -f "$TMP_FILE"
    exit 1
fi

rm -f "$TMP_FILE"
