#include <WiFi.h>
#include <esp_wifi.h>

// WiFi配置参数
const char* TEST_SSID_2G = "TEST_2G_SSID";
const char* TEST_SSID_5G = "TEST_5G_SSID"; 
const char* PASSWORD = "test_password";

// 性能测试结构体
typedef struct {
    int rssi;
    int channel;
    String ssid;
    String bssid;
    wifi_auth_mode_t authmode;
    bool is_5G;
} WiFiAPInfo;

// 测试结果结构体
typedef struct {
    String ssid;
    int connect_time;
    int rssi;
    float throughput;
    bool success;
} TestResult;

void setup() {
    Serial.begin(115200);
    delay(1000);
    
    Serial.println("====================================");
    Serial.println("ESP32-C5 无线协议性能测试");
    Serial.println("====================================");
    
    // 初始化WiFi
    WiFi.mode(WIFI_STA);
    WiFi.disconnect();
    delay(100);
}

void loop() {
    // 执行WiFi扫描测试
    performWiFiScanTest();
    
    // 执行连接性能测试
    performConnectionTest();
    
    Serial.println("\n测试完成，10秒后重新开始...");
    delay(10000);
}

// 执行WiFi扫描测试
void performWiFiScanTest() {
    Serial.println("\n--- WiFi扫描测试开始 ---");
    
    // 执行扫描
    int apCount = WiFi.scanNetworks();
    if (apCount == 0) {
        Serial.println("未发现任何WiFi网络");
        return;
    }
    
    // 分析扫描结果
    analyzeScanResults(apCount);
}

// 分析扫描结果
void analyzeScanResults(int apCount) {
    int count_2G = 0;
    int count_5G = 0;
    
    Serial.println("发现网络数量: " + String(apCount));
    
    for (int i = 0; i < apCount; i++) {
        WiFiAPInfo apInfo;
        apInfo.ssid = WiFi.SSID(i);
        apInfo.rssi = WiFi.RSSI(i);
        apInfo.channel = WiFi.channel(i);
        apInfo.authmode = WiFi.encryptionType(i);
        apInfo.bssid = WiFi.BSSIDstr(i);
        apInfo.is_5G = (apInfo.channel > 14);
        
        if (apInfo.is_5G) {
            count_5G++;
        } else {
            count_2G++;
        }
        
        // 打印AP信息
        printAPInfo(apInfo);
    }
    
    Serial.println("\n频段统计:");
    Serial.println("2.4GHz网络: " + String(count_2G));
        Serial.println("5GHz网络: " + String(count_5G));
}

// 打印AP详细信息
void printAPInfo(WiFiAPInfo apInfo) {
    Serial.println("--------------------------------");
    Serial.println("SSID: " + apInfo.ssid);
    Serial.println("BSSID: " + apInfo.bssid);
    Serial.println("信号强度: " + String(apInfo.rssi) + " dBm");
    Serial.println("信道: " + String(apInfo.channel));
    Serial.println("频段: " + String(apInfo.is_5G ? "5GHz" : "2.4GHz"));
    Serial.println("加密方式: " + getAuthModeString(apInfo.authmode));
}

// 执行连接性能测试
void performConnectionTest() {
    Serial.println("\n--- 连接性能测试开始 ---");
    
    // 测试2.4GHz连接
    TestResult result_2G = testConnection(TEST_SSID_2G, PASSWORD, false);
    printTestResult(result_2G, "2.4GHz");
    
    // 测试5GHz连接  
    TestResult result_5G = testConnection(TEST_SSID_5G, PASSWORD, true);
    printTestResult(result_5G, "5GHz");
}

// 测试特定连接
TestResult testConnection(const char* ssid, const char* password, bool is_5G) {
    TestResult result;
    result.ssid = String(ssid);
    
    Serial.println("尝试连接: " + String(ssid));
    
    unsigned long startTime = millis();
    WiFi.begin(ssid, password);
    
    int attempts = 0;
    while (WiFi.status() != WL_CONNECTED && attempts < 20) {
        delay(500);
        attempts++;
        Serial.print(".");
    }
    
    result.connect_time = (millis() - startTime) / 1000;
    result.rssi = WiFi.RSSI();
    result.success = (WiFi.status() == WL_CONNECTED);
    
    if (result.success) {
        Serial.println("\n连接成功!");
        result.throughput = estimateThroughput(result.rssi, is_5G);
    
    // 收集连接信息
    collectConnectionInfo(is_5G);
    
    // 断开连接
    WiFi.disconnect();
    delay(1000);
    
    return result;
}

// 打印测试结果
void printTestResult(TestResult result, String band) {
    Serial.println("\n" + band + "测试结果:");
    Serial.println("状态: " + String(result.success ? "成功" : "失败"));
    if (result.success) {
        Serial.println("连接时间: " + String(result.connect_time) + " 秒");
    Serial.println("信号强度: " + String(result.rssi) + " dBm");
    Serial.println("预估吞吐量: " + String(result.throughput) + " Mbps");
    }
}

// 收集连接信息
void collectConnectionInfo(bool is_5G) {
    wifi_ap_record_t ap_info;
    if (esp_wifi_sta_get_ap_info(&ap_info) == ESP_OK) {
        Serial.println("实际连接信息:");
        Serial.println("  信道: " + String(ap_info.primary));
    Serial.println("  最大数据速率: " + String(ap_info.max_data_rate) + " Mbps");
  }
}

// 预估吞吐量
float estimateThroughput(int rssi, bool is_5G) {
    // 基于信号强度和频段的简单估算
    float base_throughput = is_5G ? 200.0 : 80.0; // 基础吞吐量
    float quality_factor = 0.0;
    
    if (rssi >= -50) {
        quality_factor = 1.0;
    } else if (rssi >= -60) {
        quality_factor = 0.8;
    } else if (rssi >= -70) {
        quality_factor = 0.6;
    } else if (rssi >= -80) {
        quality_factor = 0.4;
    } else {
        quality_factor = 0.2;
    }
    
    return base_throughput * quality_factor;
}

// 获取认证模式字符串
String getAuthModeString(wifi_auth_mode_t authmode) {
    switch (authmode) {
        case WIFI_AUTH_OPEN: return "开放";
        case WIFI_AUTH_WEP: return "WEP";
        case WIFI_AUTH_WPA_PSK: return "WPA-PSK";
        case WIFI_AUTH_WPA2_PSK: return "WPA2-PSK";
        case WIFI_AUTH_WPA_WPA2_PSK: return "WPA/WPA2";
        case WIFI_AUTH_WPA3_PSK: return "WPA3-PSK";
        default: return "未知";
    }
}