/* Alice STA 端 — 连接 Bob 的 AP 热点，发送原始 802.11 数据帧并采集 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 "STA"

/* Bob AP 的 SSID 和密码 */
#define AP_SSID      "test_claude_Bob"
#define AP_PASSWORD  "12345678"

/* CSI 目标 MAC（Bob AP 的 BSSID，从 ap_info 获取） */
static uint8_t g_target_mac[6];
static bool g_target_mac_set = false;

/* Alice 自己的 MAC 地址（作为 802.11 帧的 TA，也用于 CSI dmac 过滤） */
static uint8_t g_own_mac[6];
static bool g_own_mac_set = false;

/* ── Sniffer 回调 — 全帧放行，对比驱动分类 vs FC 字段真值 ── */
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;

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

    /* 从 FC 字段解析真实 802.11 帧类型 */
    uint8_t fc_type = (pl[0] >> 2) & 0x03;
    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 数据接收回调 — 采集 Bob→Alice 的 ACK 帧 CSI
 *
 * 802.11 ACK 帧无 source 地址，data->mac=00:00:00:00:00:00。
 * 因此用 dmac (目标地址) 匹配 Alice MAC + sig_len<50 来识别 ACK 帧。 */
static void wifi_csi_rx_cb(void *ctx, wifi_csi_info_t *data)
{
    bool mac_match = g_target_mac_set && (memcmp(data->mac, g_target_mac, 6) == 0);
    bool dmac_match = g_own_mac_set && (memcmp(data->dmac, g_own_mac, 6) == 0);
    bool looks_like_ack = !data->rx_ctrl.is_group && data->rx_ctrl.sig_len < 50;

    if (!mac_match && !(dmac_match && looks_like_ack)) {
        return;
    }

    if (data->rx_ctrl.is_group || data->rx_ctrl.sig_len >= 50) {
        return;
    }

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

/* CSI 初始化 — Alice 作为 STA 采集来自 Bob 的 ACK 帧 CSI */
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));
    wifi_promiscuous_filter_t ctrl_filter = {
        .filter_mask = WIFI_PROMIS_CTRL_FILTER_MASK_ACK,
    };
    ESP_ERROR_CHECK(esp_wifi_set_promiscuous_ctrl_filter(&ctrl_filter));

    // API 调用顺序必须为 config → rx_cb → enable
    wifi_csi_config_t csi_cfg = {
        .enable = true,
        .acquire_csi_legacy = true,   // ACK 帧为 Legacy 格式 (L-LTF, len=106)
        .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 = true,
        .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 初始化完成 — dump_ack_en=true, legacy+ht20");
}

/* WiFi 连接完成标志 */
static bool g_wifi_connected = false;

/* ─────────────────────────────────────────────────────────────
 * 发送原始 802.11 数据帧
 *
 * 绕过 LwIP/UDP 协议栈，直接构造 802.11 Data 帧发给 Bob。
 * Bob 作为 AP，MAC 层对单播 Data 帧必须回复 ACK。
 * 这样可以验证 Alice 的物理层是否能接收到 ACK 帧。
 *
 * 802.11 Data 帧格式 (STA→AP, To DS=1, 非加密):
 *   [0-1]  Frame Control = 0x0108
 *   [2-3]  Duration = 0x0000
 *   [4-9]  Address 1 = RA  = Bob BSSID
 *   [10-15] Address 2 = TA  = Alice MAC
 *   [16-21] Address 3 = BSSID = Bob BSSID
 *   [22-23] Sequence Control (en_sys_seq=true 时由驱动自动填充)
 *   [24+]  Frame Body (payload)
 * ───────────────────────────────────────────────────────────── */
static esp_err_t send_80211_data_frame(void)
{
    if (!g_target_mac_set || !g_own_mac_set) {
        return ESP_ERR_INVALID_STATE;
    }

    const char *body = "Hello Bob from 802.11 raw frame";
    size_t body_len = strlen(body);
    size_t buf_len = 24 + body_len;  // 802.11 header(24) + payload

    uint8_t *buf = malloc(buf_len);
    if (!buf) {
        return ESP_ERR_NO_MEM;
    }

    // Frame Control (little-endian): Type=Data(10), Subtype=Data(0000), To DS=1
    buf[0] = 0x08;  // Protocol=00, Type=10, Subtype=0000
    buf[1] = 0x01;  // To DS=1, From DS=0, no frag/retry/pm/md/protected/order

    // Duration
    buf[2] = 0x00;
    buf[3] = 0x00;

    // Address 1 = RA (Receiver) → Bob BSSID
    memcpy(&buf[4], g_target_mac, 6);

    // Address 2 = TA (Transmitter) → Alice MAC
    memcpy(&buf[10], g_own_mac, 6);

    // Address 3 = BSSID → Bob BSSID
    memcpy(&buf[16], g_target_mac, 6);

    // Sequence Control: let driver fill (en_sys_seq=true)
    buf[22] = 0x00;
    buf[23] = 0x00;

    // Frame Body
    memcpy(&buf[24], body, body_len);

    esp_err_t ret = esp_wifi_80211_tx(WIFI_IF_STA, buf, buf_len, true);
    free(buf);
    return ret;
}

/* 原始 802.11 数据帧发送任务 — 每秒发一帧 */
static void raw_80211_tx_task(void *pvParameters)
{
    while (!g_wifi_connected || !g_own_mac_set) {
        vTaskDelay(pdMS_TO_TICKS(100));
    }
    vTaskDelay(pdMS_TO_TICKS(500));

    ESP_LOGI(TAG, "802.11 raw TX 任务启动, own_mac=" MACSTR, MAC2STR(g_own_mac));

    uint32_t pkt_count = 0;
    while (1) {
        esp_err_t ret = send_80211_data_frame();
        if (ret == ESP_OK) {
            ESP_LOGI(TAG, "80211TX[%lu] 原始 Data 帧已发送 → Bob", pkt_count++);
        } else {
            ESP_LOGW(TAG, "80211TX[%lu] 发送失败: %s", pkt_count, esp_err_to_name(ret));
        }
        vTaskDelay(pdMS_TO_TICKS(1000));
    }
}

/* 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_STA_START) {
        ESP_LOGI(TAG, "STA 启动成功，开始连接 AP %s", AP_SSID);
        esp_wifi_connect();
    } else if (event_base == WIFI_EVENT && event_id == WIFI_EVENT_STA_DISCONNECTED) {
        wifi_event_sta_disconnected_t *disconn = (wifi_event_sta_disconnected_t *)event_data;
        ESP_LOGI(TAG, "连接断开，reason=%d，尝试重连", disconn->reason);
        g_wifi_connected = false;
        esp_wifi_connect();
    } else if (event_base == IP_EVENT && event_id == IP_EVENT_STA_GOT_IP) {
        ip_event_got_ip_t *event = (ip_event_got_ip_t *)event_data;
        ESP_LOGI(TAG, "获取到 IP 地址: " IPSTR, IP2STR(&event->ip_info.ip));
        ESP_LOGI(TAG, ">>> WiFi 连接成功！Alice 已连接到 Bob <<<");
        g_wifi_connected = true;

        // 获取 Alice 自己的 MAC 地址 (用于 802.11 raw frame TA)
        if (esp_wifi_get_mac(WIFI_IF_STA, g_own_mac) == ESP_OK) {
            g_own_mac_set = true;
            ESP_LOGI(TAG, "自己的 MAC: " MACSTR, MAC2STR(g_own_mac));
        }

        // 获取 AP (Bob) 的 BSSID 用于 CSI 回调过滤
        wifi_ap_record_t ap_info;
        if (esp_wifi_sta_get_ap_info(&ap_info) == ESP_OK) {
            memcpy(g_target_mac, ap_info.bssid, 6);
            g_target_mac_set = true;
            ESP_LOGI(TAG, "目标 MAC (Bob) = " MACSTR, MAC2STR(g_target_mac));
        }
    }
}

/* WiFi STA 初始化 */
static void wifi_sta_init(void)
{
    // 初始化网络协议栈
    esp_netif_init();
    ESP_ERROR_CHECK(esp_event_loop_create_default());
    esp_netif_create_default_wifi_sta();

    // 初始化 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));
    ESP_ERROR_CHECK(esp_event_handler_register(IP_EVENT, IP_EVENT_STA_GOT_IP,
                                               &wifi_event_handler, NULL));

    // 配置 STA
    wifi_config_t sta_config = {
        .sta = {
            .ssid = AP_SSID,
            .password = AP_PASSWORD,
        },
    };
    ESP_ERROR_CHECK(esp_wifi_set_mode(WIFI_MODE_STA));
    ESP_ERROR_CHECK(esp_wifi_set_config(WIFI_IF_STA, &sta_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_STA, WIFI_PROTOCOL_11A | WIFI_PROTOCOL_11N));
    ESP_ERROR_CHECK(esp_wifi_set_bandwidth(WIFI_IF_STA, WIFI_BW20));

    // 初始化 CSI（必须在上述配置之后，esp_wifi_connect 之前）
    csi_init();

    ESP_LOGI(TAG, "Alice STA 初始化完成，5GHz / Channel跟随AP / 802.11an / HT20");
}

void app_main(void)
{
    ESP_LOGI(TAG, "=== Alice STA 端启动 ===");

    // 初始化 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 STA
    wifi_sta_init();

    // 创建原始 802.11 数据帧发送任务
    xTaskCreate(raw_80211_tx_task, "raw_80211_tx", 4096, NULL, 5, NULL);
}