/*
  WiFi NAN and Beacon driver

  with thanks to https://github.com/sxjack/uav_electronic_ids for WiFi calls
 */

#include "WiFi_TX.h"
#include <esp_wifi.h>
#include <WiFi.h>
#include <esp_system.h>
#include "parameters.h"

bool WiFi_TX::init(void)
{
    if (initialised) {
        return true;
    }
    initialised = true;

    //use a local MAC address to avoid tracking transponders based on their MAC address
    uint8_t mac_addr[6];
    generate_random_mac(mac_addr);

    mac_addr[0] |= 0x02;  // set MAC local bit
    mac_addr[0] &= 0xFE;  // unset MAC multicast bit
    Serial.printf("Setting MAC: %02X:%02X:%02X:%02X:%02X:%02X\n",
                mac_addr[0], mac_addr[1], mac_addr[2],
                mac_addr[3], mac_addr[4], mac_addr[5]);
    //set MAC address
    esp_base_mac_addr_set(mac_addr);
    // 2.4Gwifi信道默认是信道6   5.8GHZ应使用信道149
    if (g.webserver_enable == 0) {
        WiFi.softAP(g.wifi_ssid, g.wifi_password, g.wifi_channel, false, 0); //make it visible and allow no connection
    } else {
        WiFi.softAP(g.wifi_ssid, g.wifi_password, g.wifi_channel, false, 1); //make it visible and allow only 1 connection
    }

    if (esp_wifi_set_bandwidth(WIFI_IF_AP, WIFI_BW_HT20) != ESP_OK) { // 国标要求 20MHZ带宽
        return false;
    }

    memcpy(WiFi_mac_addr,mac_addr,6); //use generated random MAC address for OpenDroneID messages

    esp_wifi_set_max_tx_power(dBm_to_tx_power(g.wifi_power)); // 按国标要求来说设置不同的功率 

    return true;
}

// 发送中国国标的数据包
bool WiFi_TX::transmit_cn_nan(CNDID_UAS_Data &UAS_data)
{
    init();

    uint8_t buffer[1024] {};

    int length;

     if ((length = cndid_wifi_build_nan_sync_beacon_frame((char *)WiFi_mac_addr,
                  buffer,sizeof(buffer))) > 0) {
        if (esp_wifi_80211_tx(WIFI_IF_AP,buffer,length,true) != ESP_OK) {
            return false;
        }
    }

    if ((length = cndid_wifi_build_message_pack_nan_action_frame(&UAS_data, (char *)WiFi_mac_addr,
                  ++send_counter_cn_nan,
                  buffer,sizeof(buffer))) > 0) {
        if (esp_wifi_80211_tx(WIFI_IF_AP,buffer,length,true) != ESP_OK) {
            return false;
        }
    }

    return true;
}

// 发送国际标准的数据包
bool WiFi_TX::transmit_nan(ODID_UAS_Data &UAS_data)
{
    init();

    uint8_t buffer[1024] {};

    int length;

    if ((length = odid_wifi_build_nan_sync_beacon_frame((char *)WiFi_mac_addr,
                  buffer,sizeof(buffer))) > 0) {
        if (esp_wifi_80211_tx(WIFI_IF_AP,buffer,length,true) != ESP_OK) {
            return false;
        }
    }

    if ((length = odid_wifi_build_message_pack_nan_action_frame(&UAS_data,(char *)WiFi_mac_addr,
                  ++send_counter_nan,
                  buffer,sizeof(buffer))) > 0) {
        if (esp_wifi_80211_tx(WIFI_IF_AP,buffer,length,true) != ESP_OK) {
            return false;
        }
    }

    return true;
}

bool WiFi_TX::transmit_cn_beacon(CNDID_UAS_Data &UAS_data)
{
    init();

    uint8_t buffer[1024] {};

    int length;
    if ((length = cndid_wifi_build_message_pack_beacon_frame(&UAS_data,(char *)WiFi_mac_addr,
                   "UAS_ID_OPEN", strlen("UAS_ID_OPEN"), //use dummy SSID, as we only extract payload data
                  1000/g.wifi_beacon_rate, ++send_counter_cn_beacon, buffer, sizeof(buffer))) > 0) {

        //set the RID IE element
        uint8_t header_offset = 58; // 固定包头长度
        vendor_ie_data_t IE_data;
        IE_data.element_id = 0xDD; // 国标要求
        // 符合国标
        IE_data.vendor_oui[0] = 0xFA;
        IE_data.vendor_oui[1] = 0x0B;
        IE_data.vendor_oui[2] = 0xBC;
        IE_data.vendor_oui_type = 0x0D;
        IE_data.length = length - header_offset + 4; // 添加4作为esp_wifi_set_vendor_ie的定义
        // 也就是 uoi vender类型是4个字节
        memcpy(IE_data.payload,&buffer[header_offset],length - header_offset);

        // 所以首先移除旧元素，之后添加新元素
        if (esp_wifi_set_vendor_ie(false, WIFI_VND_IE_TYPE_BEACON, WIFI_VND_IE_ID_0, &IE_data) != ESP_OK){
            return false;
        }

        if (esp_wifi_set_vendor_ie(true, WIFI_VND_IE_TYPE_BEACON, WIFI_VND_IE_ID_0, &IE_data) != ESP_OK){
            return false;
        }

        //也将有效载荷设置为探测请求，以提高手机上的更新速率
        //因此，先移除旧元素，之后再添加新元素
        if (esp_wifi_set_vendor_ie(false, WIFI_VND_IE_TYPE_PROBE_RESP, WIFI_VND_IE_ID_0, &IE_data) != ESP_OK){
            return false;
        }

        if (esp_wifi_set_vendor_ie(true, WIFI_VND_IE_TYPE_PROBE_RESP, WIFI_VND_IE_ID_0, &IE_data) != ESP_OK){
            return false;
        }

        return true;
	}
	else {
		return false;
	}
}

//update the payload of the beacon frames in this function
bool WiFi_TX::transmit_beacon(ODID_UAS_Data &UAS_data)
{
    init();

    uint8_t buffer[1024] {};

    int length;
    if ((length = odid_wifi_build_message_pack_beacon_frame(&UAS_data,(char *)WiFi_mac_addr,
                   "UAS_ID_OPEN", strlen("UAS_ID_OPEN"), //use dummy SSID, as we only extract payload data
                  1000/g.wifi_beacon_rate, ++send_counter_beacon, buffer, sizeof(buffer))) > 0) {

        //设置厂商 RID的元素
        uint8_t header_offset = 58; // 定义IE元素头部偏移（58字节是Beacon帧固定头部长度，跳过头部取核心数据）
        vendor_ie_data_t IE_data; // 定义厂商IE元素结构体（ESP32 WiFi驱动要求的格式）
        IE_data.element_id = WIFI_VENDOR_IE_ELEMENT_ID; // 固定值，标识是厂商自定义IE
        IE_data.vendor_oui[0] = 0xFA; // 设置Open Drone ID标准的厂商OUI（FA:0B:BC是无人机Remote ID的专属OUI）
        IE_data.vendor_oui[1] = 0x0B;
        IE_data.vendor_oui[2] = 0xBC;
        IE_data.vendor_oui_type = 0x0D; // Open Drone ID定义的WiFi IE类型值
        // // 步骤4：设置IE元素长度（核心Remote ID数据长度 + 4字节ESP32驱动要求的额外长度）
        IE_data.length = length - header_offset + 4; //add 4 as of definition esp_wifi_set_vendor_ie
        // 拷贝Remote ID核心数据到IE元素的payload（跳过Beacon帧固定头部）
        memcpy(IE_data.payload,&buffer[header_offset],length - header_offset);
        // 删除旧的Beacon帧厂商IE元素（false=禁用）
        // so first remove old element, add new afterwards
        if (esp_wifi_set_vendor_ie(false, WIFI_VND_IE_TYPE_BEACON, WIFI_VND_IE_ID_0, &IE_data) != ESP_OK){
            return false;
        }
        // 添加新的Beacon帧厂商IE元素（true=启用）
        if (esp_wifi_set_vendor_ie(true, WIFI_VND_IE_TYPE_BEACON, WIFI_VND_IE_ID_0, &IE_data) != ESP_OK){
            return false;
        }
        // 为Probe Response帧（手机探测WiFi时的响应帧）也添加相同的Remote ID IE元素
        //set the payload also to probe requests, to increase update rate on mobile phones
        // so first remove old element, add new afterwards
        if (esp_wifi_set_vendor_ie(false, WIFI_VND_IE_TYPE_PROBE_RESP, WIFI_VND_IE_ID_0, &IE_data) != ESP_OK){
            return false;
        }

        if (esp_wifi_set_vendor_ie(true, WIFI_VND_IE_TYPE_PROBE_RESP, WIFI_VND_IE_ID_0, &IE_data) != ESP_OK){
            return false;
        }

        return true;
	}
	else {
		return false;
	}
}
bool WiFi_TX::transmit_GB2025_nan(UavIdentificationData &UAS_data)
{
    init();

    uint8_t buffer[1024] {};

    int length;

    if ((length = did_GB2025_wifi_build_nan_sync_beacon_frame((char *)WiFi_mac_addr,
                  buffer,sizeof(buffer))) > 0) {
        if (esp_wifi_80211_tx(WIFI_IF_AP,buffer,length,true) != ESP_OK) {
            return false;
        }
    }

    if ((length = did_GB2025_wifi_build_message_pack_nan_action_frame(&UAS_data,(char *)WiFi_mac_addr,
                  ++send_counter_GB2025_nan,
                  buffer,sizeof(buffer))) > 0) {
        if (esp_wifi_80211_tx(WIFI_IF_AP,buffer,length,true) != ESP_OK) {
            return false;
        }
    }

    return true;
}
bool WiFi_TX::transmit_GB2025_beacon(UavIdentificationData &UAS_data)
{
    init();

    uint8_t buffer[1024] {};

    int length;
    if ((length = did_GB2025_wifi_build_message_pack_beacon_frame(&UAS_data,(char *)WiFi_mac_addr,
                   "UAS_ID_OPEN", strlen("UAS_ID_OPEN"), //use dummy SSID, as we only extract payload data
                  1000/g.wifi_beacon_rate, ++send_counter_GB2025_beacon, buffer, sizeof(buffer))) > 0) {

        //set the RID IE element
        uint8_t header_offset = 58; // 固定包头长度
        vendor_ie_data_t IE_data;
        IE_data.element_id = 0xDD; // 国标要求
        // 符合国标
        IE_data.vendor_oui[0] = 0xFA;
        IE_data.vendor_oui[1] = 0x0B;
        IE_data.vendor_oui[2] = 0xBC;
        IE_data.vendor_oui_type = 0x0D;
        IE_data.length = length - header_offset + 4; // 添加4作为esp_wifi_set_vendor_ie的定义
        // 也就是 uoi vender类型是4个字节
        memcpy(IE_data.payload,&buffer[header_offset],length - header_offset);

        // 所以首先移除旧元素，之后添加新元素
        if (esp_wifi_set_vendor_ie(false, WIFI_VND_IE_TYPE_BEACON, WIFI_VND_IE_ID_0, &IE_data) != ESP_OK){
            return false;
        }

        if (esp_wifi_set_vendor_ie(true, WIFI_VND_IE_TYPE_BEACON, WIFI_VND_IE_ID_0, &IE_data) != ESP_OK){
            return false;
        }

        //也将有效载荷设置为探测请求，以提高手机上的更新速率
        //因此，先移除旧元素，之后再添加新元素
        if (esp_wifi_set_vendor_ie(false, WIFI_VND_IE_TYPE_PROBE_RESP, WIFI_VND_IE_ID_0, &IE_data) != ESP_OK){
            return false;
        }

        if (esp_wifi_set_vendor_ie(true, WIFI_VND_IE_TYPE_PROBE_RESP, WIFI_VND_IE_ID_0, &IE_data) != ESP_OK){
            return false;
        }

        return true;
	}
	else {
		return false;
	}
}

/*
  map dBm to a TX power
 */
uint8_t WiFi_TX::dBm_to_tx_power(float dBm) const
{
    if (dBm < 2) {
        dBm = 2;
    }
    if (dBm > 20) {
        dBm = 20;
    }
    return uint8_t((dBm+1.125) * 4);
}