import subprocess
import time
import re

class WifiManager:
    def __init__(self, interface="wlan0"):
        self.interface = interface

    def _run_wpa_cli(self, *args):
        """执行 wpa_cli 命令"""
        cmd = ["wpa_cli", "-i", self.interface] + list(args)
        try:
            result = subprocess.run(cmd, capture_output=True, text=True, check=True)
            return result.stdout.strip()
        except subprocess.CalledProcessError as e:
            print(f"执行 wpa_cli {args} 失败: {e.stderr}")
            return ""

    def list_saved_networks(self):
        """列出已保存的 wifi"""
        output = self._run_wpa_cli("list_networks")
        networks = []
        lines = output.splitlines()
        if len(lines) > 1:
            for line in lines[1:]:  # skip header
                parts = line.split()
                if len(parts) >= 4:
                    networks.append({
                        "network_id": parts[0],
                        "ssid": parts[1],
                        "bssid": parts[2],
                        "flags": parts[3]
                    })
        return networks

    def scan_networks(self):
        """扫描范围内的 wifi"""
        self._run_wpa_cli("scan")
        time.sleep(3)  # 等待扫描完成
        output = self._run_wpa_cli("scan_results")
        networks = []
        lines = output.splitlines()
        if len(lines) > 1:
             for line in lines[1:]:
                 # bssid / frequency / signal level / flags / ssid
                 parts = line.split('\t')
                 if len(parts) >= 5:
                     networks.append({
                         "bssid": parts[0],
                         "frequency": parts[1],
                         "signal_level": parts[2],
                         "flags": parts[3],
                         "ssid": parts[4]
                     })
        return networks

    def connect_wifi(self, ssid, password=None):
        """连接普通 Wi-Fi (WPA2-PSK)"""
        network_id = self._run_wpa_cli("add_network")
        if not network_id.isdigit():
            return False

        self._run_wpa_cli("set_network", network_id, "ssid", f'"{ssid}"')
        if password:
            self._run_wpa_cli("set_network", network_id, "psk", f'"{password}"')
        else:
            self._run_wpa_cli("set_network", network_id, "key_mgmt", "NONE")

        self._run_wpa_cli("enable_network", network_id)
        self._run_wpa_cli("select_network", network_id)
        self._run_wpa_cli("save_config")
        return True

    def connect_eap(self, ssid, identity, password):
        """连接企业级 Wi-Fi (WPA-EAP PEAP/MSCHAPv2)"""
        network_id = self._run_wpa_cli("add_network")
        if not network_id.isdigit():
            return False

        self._run_wpa_cli("set_network", network_id, "ssid", f'"{ssid}"')
        self._run_wpa_cli("set_network", network_id, "key_mgmt", "WPA-EAP")
        self._run_wpa_cli("set_network", network_id, "eap", "PEAP")
        self._run_wpa_cli("set_network", network_id, "phase2", '"auth=MSCHAPV2"')
        self._run_wpa_cli("set_network", network_id, "identity", f'"{identity}"')
        self._run_wpa_cli("set_network", network_id, "password", f'"{password}"')

        self._run_wpa_cli("enable_network", network_id)
        self._run_wpa_cli("select_network", network_id)
        self._run_wpa_cli("save_config")
        return True

    def remove_network(self, network_id):
        """删除某个已保存的网络"""
        self._run_wpa_cli("remove_network", str(network_id))
        self._run_wpa_cli("save_config")

    def get_current_status(self):
        """获取当前网络状态"""
        output = self._run_wpa_cli("status")
        status = {}
        for line in output.splitlines():
            if '=' in line:
                key, val = line.split('=', 1)
                status[key] = val
        return status

    def open_hotspot(self, ssid, password, channel=6):
        """
        开启热点 (AP 模式)
        注意：要求网卡和 wpa_supplicant 均支持 AP 模式 (mode=2)
        """
        network_id = self._run_wpa_cli("add_network")
        if not network_id.isdigit():
            return False

        self._run_wpa_cli("set_network", network_id, "ssid", f'"{ssid}"')
        self._run_wpa_cli("set_network", network_id, "mode", "2")
        self._run_wpa_cli("set_network", network_id, "key_mgmt", "WPA-PSK")
        self._run_wpa_cli("set_network", network_id, "psk", f'"{password}"')
        self._run_wpa_cli("set_network", network_id, "frequency", str(2412 + (channel - 1) * 5))

        self._run_wpa_cli("select_network", network_id)
        self._run_wpa_cli("save_config")
        return network_id

    def close_hotspot(self, network_id=None):
        """关闭热点"""
        if network_id is not None:
             self._run_wpa_cli("remove_network", str(network_id))
        self._run_wpa_cli("reconfigure")
        self._run_wpa_cli("save_config")

if __name__ == "__main__":
    # Example Usage
    wifi = WifiManager("wlan0")
    print("Current status:", wifi.get_current_status())
    print("Saved networks:", wifi.list_saved_networks())