/* Bob AP 端 — 创建 WiFi 热点供 Alice 连接，采集 CSI 数据 */

#include <string.h>
#include "freertos/FreeRTOS.h"
#include "freertos/task.h"
#include "freertos/event_groups.h"
#include "esp_system.h"
#include "esp_wifi.h"
#include "esp_event.h"
#include "esp_log.h"
#include "nvs_flash.h"
#include "esp_mac.h"
#define TAG "AP"

/* AP 配置参数 */
#define AP_SSID          "test_claude_Bob"
#define AP_PASSWORD      "12345678"
#define AP_CHANNEL       40
#define AP_MAX_CONN      4

/* CSI 目标 MAC（Alice STA 的 MAC，从连接事件获取） */
static uint8_t g_target_mac[6];
static bool g_target_mac_set = false;

/* ── Sniffer 回调 — 全帧放行，对比驱动分类 vs FC 字段真值 ──
 *
 * 802.11 FC 字段 (pl[0]): bits[3:2]=Type  (0=MGMT, 1=CTRL, 2=DATA)
 *                           bits[7:4]=Subtype
 * driver type:  从 wifi_promiscuous_pkt_type_t 传入
 * fc_type:      从 FC 字段解析的真实 802.11 帧类型
 * 同时打印以判断是驱动 bug 还是硬件问题 */
static void wifi_sniffer_cb(void *buf, wifi_promiscuous_pkt_type_t type)
{
    static const char *tname[] = { "MGMT", "CTRL", "DATA", "MISC", "?" };
    wifi_promiscuous_pkt_t *pkt = (wifi_promiscuous_pkt_t *)buf;
    uint8_t *pl = pkt->payload;

    if (!g_target_mac_set) return;

    /* MAC 匹配：Data/Mgmt 检查 TA(offset10) 或 RA(offset4)，Ctrl 只查 RA */
    bool match = false;
    if (type == WIFI_PKT_CTRL) {
        match = (memcmp(&pl[4], g_target_mac, 6) == 0);
    } else {
        match = (memcmp(&pl[10], g_target_mac, 6) == 0)
             || (memcmp(&pl[4], g_target_mac, 6) == 0);
    }
    if (!match) return;

    /* 从 FC 字段解析真实 802.11 帧类型 */
    uint8_t fc_type = (pl[0] >> 2) & 0x03;  // bits[3:2]
    const char *fc_name = (fc_type <= 2) ? tname[fc_type] : "?";

    ESP_LOGI(TAG, "SNIFF drv=%s fc=%s | ra=" MACSTR " ta=" MACSTR " sig=%u rssi=%d fc=0x%02x%02x",
             (type <= WIFI_PKT_MISC) ? tname[type] : "?",
             fc_name,
             MAC2STR(&pl[4]), MAC2STR(&pl[10]),
             pkt->rx_ctrl.sig_len, pkt->rx_ctrl.rssi, pl[1], pl[0]);
}

/* CSI 数据接收回调 — 过滤：目标 MAC + 单播数据帧 (sig_len >= 50) */
static void wifi_csi_rx_cb(void *ctx, wifi_csi_info_t *data)
{
    if (!g_target_mac_set || memcmp(data->mac, g_target_mac, 6) != 0) {
        return;
    }

    ESP_LOGI(TAG, "CSI_DATA: mac=" MACSTR " len=%u rssi=%d ch=%d sig_len=%u",
             MAC2STR(data->mac), data->len, data->rx_ctrl.rssi,
             data->rx_ctrl.channel, data->rx_ctrl.sig_len);
}

/* CSI 初始化 — Bob 作为 AP */
static void csi_init(void)
{
    ESP_ERROR_CHECK(esp_wifi_set_promiscuous(true));

    // sniffer 透视 — 全帧放行
    ESP_ERROR_CHECK(esp_wifi_set_promiscuous_rx_cb(wifi_sniffer_cb));
    wifi_promiscuous_filter_t promis_filter = {
        .filter_mask = WIFI_PROMIS_FILTER_MASK_ALL,
    };
    ESP_ERROR_CHECK(esp_wifi_set_promiscuous_filter(&promis_filter));

    // 配置 CSI — legacy + ht20，覆盖 Alice raw 802.11 帧
    wifi_csi_config_t csi_cfg = {
        .enable = true,
        .acquire_csi_legacy = true,
        .acquire_csi_force_lltf = false,
        .acquire_csi_ht20 = true,
        .acquire_csi_ht40 = false,
        .acquire_csi_vht = false,
        .acquire_csi_su = false,
        .acquire_csi_mu = false,
        .acquire_csi_dcm = false,
        .acquire_csi_beamformed = false,
        .acquire_csi_he_stbc_mode = 0,
        .val_scale_cfg = 0,
        .dump_ack_en = false,
        .reserved = 0,
    };
    ESP_ERROR_CHECK(esp_wifi_set_csi_config(&csi_cfg));
    ESP_ERROR_CHECK(esp_wifi_set_csi_rx_cb(wifi_csi_rx_cb, NULL));
    ESP_ERROR_CHECK(esp_wifi_set_csi(true));

    ESP_LOGI(TAG, "CSI+Sniffer 初始化完成 — legacy + ht20");
}

/* WiFi 事件回调 */
static void wifi_event_handler(void *arg, esp_event_base_t event_base,
                               int32_t event_id, void *event_data)
{
    if (event_base == WIFI_EVENT && event_id == WIFI_EVENT_AP_START) {
        ESP_LOGI(TAG, "AP 启动成功！%s 开始广播", AP_SSID);
    } else if (event_base == WIFI_EVENT && event_id == WIFI_EVENT_AP_STACONNECTED) {
        wifi_event_ap_staconnected_t *event = (wifi_event_ap_staconnected_t *)event_data;
        ESP_LOGI(TAG, "检测到 STA 连接, MAC: " MACSTR, MAC2STR(event->mac));
        ESP_LOGI(TAG, ">>> 有设备连接到了 Bob AP <<<");
        memcpy(g_target_mac, event->mac, 6);
        g_target_mac_set = true;
    } else if (event_base == WIFI_EVENT && event_id == WIFI_EVENT_AP_STADISCONNECTED) {
        wifi_event_ap_stadisconnected_t *event = (wifi_event_ap_stadisconnected_t *)event_data;
        ESP_LOGI(TAG, "STA 断开, MAC: " MACSTR, MAC2STR(event->mac));
        g_target_mac_set = false;
    }
}

/* WiFi AP 初始化 */
static void wifi_ap_init(void)
{
    // 初始化网络协议栈
    esp_netif_init();
    ESP_ERROR_CHECK(esp_event_loop_create_default());
    esp_netif_create_default_wifi_ap();

    // 初始化 WiFi
    wifi_init_config_t cfg = WIFI_INIT_CONFIG_DEFAULT();
    ESP_ERROR_CHECK(esp_wifi_init(&cfg));

    // 注册事件回调
    ESP_ERROR_CHECK(esp_event_handler_register(WIFI_EVENT, ESP_EVENT_ANY_ID,
                                               &wifi_event_handler, NULL));

    // 配置 AP（channel 先设为 0 自动，避免 2G_only 模式下校验 5G 信道失败）
    wifi_config_t ap_config = {
        .ap = {
            .ssid = AP_SSID,
            .password = AP_PASSWORD,
            .ssid_len = strlen(AP_SSID),
            .channel = 0,    // 先自动选择，start 后再用 esp_wifi_set_channel 固定
            .max_connection = AP_MAX_CONN,
            .authmode = WIFI_AUTH_WPA2_PSK,
        },
    };
    ESP_ERROR_CHECK(esp_wifi_set_mode(WIFI_MODE_AP));
    ESP_ERROR_CHECK(esp_wifi_set_config(WIFI_IF_AP, &ap_config));
    ESP_ERROR_CHECK(esp_wifi_start());

    // 以下配置必须在 esp_wifi_start() 之后调用
    ESP_ERROR_CHECK(esp_wifi_set_band_mode(WIFI_BAND_MODE_5G_ONLY));
    ESP_ERROR_CHECK(esp_wifi_set_protocol(WIFI_IF_AP, WIFI_PROTOCOL_11A | WIFI_PROTOCOL_11N));
    ESP_ERROR_CHECK(esp_wifi_set_bandwidth(WIFI_IF_AP, WIFI_BW20));

    // 手动固定到 5GHz channel 40（HT20 下无需 second channel）
    ESP_ERROR_CHECK(esp_wifi_set_channel(AP_CHANNEL, WIFI_SECOND_CHAN_NONE));

    // 初始化 CSI（监听来自 Alice 的数据帧）
    csi_init();

    ESP_LOGI(TAG, "Bob AP 初始化完成，5GHz Ch40 / 802.11an / HT20，等待 STA 连接...");
}

void app_main(void)
{
    ESP_LOGI(TAG, "=== Bob AP 端启动 ===");

    // 初始化 NVS
    esp_err_t ret = nvs_flash_init();
    if (ret == ESP_ERR_NVS_NO_FREE_PAGES || ret == ESP_ERR_NVS_NEW_VERSION_FOUND) {
        ESP_ERROR_CHECK(nvs_flash_erase());
        ret = nvs_flash_init();
    }
    ESP_ERROR_CHECK(ret);

    // 启动 WiFi AP
    wifi_ap_init();

}