/*
Copyright (C) 2020 Simon Wunderlich, Marek Sobe
Copyright (C) 2020 Doodle Labs

SPDX-License-Identifier: Apache-2.0

Open Drone ID C Library

Maintainer:
Simon Wunderlich
sw@simonwunderlich.de
*/
/*该文件是对开源wifi修改适配国产标准*/
#if defined(ARDUINO_ARCH_ESP32)
#include <Arduino.h>
int clock_gettime(clockid_t, struct timespec *);
#else 
#include <string.h>
#include <stddef.h>
#include <stdio.h>
#endif

#include <errno.h>
#include <time.h>

#include "did_GB2025_wifi.h"


#if __BYTE_ORDER__ == __ORDER_LITTLE_ENDIAN__
#define cpu_to_le16(x)  (x)
#define cpu_to_le64(x)  (x)
#else
#define cpu_to_le16(x)      (bswap_16(x))
#define cpu_to_le64(x)      (bswap_64(x))
#endif

#define IEEE80211_FCTL_FTYPE          0x000c
#define IEEE80211_FCTL_STYPE          0x00f0

#define IEEE80211_FTYPE_MGMT            0x0000
#define IEEE80211_STYPE_ACTION          0x00D0
#define IEEE80211_STYPE_BEACON          0x0080

/* IEEE 802.11-2016 capability info */
#define IEEE80211_CAPINFO_ESS               0x0001
#define IEEE80211_CAPINFO_IBSS              0x0002
#define IEEE80211_CAPINFO_CF_POLLABLE       0x0004
#define IEEE80211_CAPINFO_CF_POLLREQ        0x0008
#define IEEE80211_CAPINFO_PRIVACY           0x0010
#define IEEE80211_CAPINFO_SHORT_PREAMBLE    0x0020
/* bits 6-7 reserved */
#define IEEE80211_CAPINFO_SPECTRUM_MGMT     0x0100
#define IEEE80211_CAPINFO_QOS               0x0200
#define IEEE80211_CAPINFO_SHORT_SLOTTIME    0x0400
#define IEEE80211_CAPINFO_APSD              0x0800
#define IEEE80211_CAPINFO_RADIOMEAS         0x1000
/* bit 13 reserved */
#define IEEE80211_CAPINFO_DEL_BLOCK_ACK     0x4000
#define IEEE80211_CAPINFO_IMM_BLOCK_ACK     0x8000

/* IEEE 802.11 Element IDs */
#define IEEE80211_ELEMID_SSID		0x00
#define IEEE80211_ELEMID_RATES		0x01
#define IEEE80211_ELEMID_VENDOR		0xDD

/* 《邻居感知网络规范》v3.1第2.8.2节
* NAN集群ID是一个MAC地址，其取值范围为50-6F-9A-01-00-00至50-6F-9A-01-FF-FF，
* 承载于部分NAN帧的A3字段中。NAN集群ID由发起NAN集群的设备随机选择。
* 然而，《ASTM远程ID规范》v1.1规定，NAN集群ID必须固定为50-6F-9A-01-00-FF这一值。
 */
static const uint8_t *get_nan_cluster_id(void)
{
    static const uint8_t cluster_id[6] = { 0x50, 0x6F, 0x9A, 0x01, 0x00, 0xFF };
    return cluster_id;
}

static int buf_fill_ieee80211_mgmt(uint8_t *buf, size_t *len, size_t buf_size,
                                   const uint16_t subtype,
                                   const uint8_t *dst_addr,
                                   const uint8_t *src_addr,
                                   const uint8_t *bssid)
{
    if (*len + sizeof(struct ieee80211_mgmt) > buf_size)
        return -ENOMEM;

    struct ieee80211_mgmt *mgmt = (struct ieee80211_mgmt *)(buf + *len);
    mgmt->frame_control = (uint16_t) cpu_to_le16(IEEE80211_FTYPE_MGMT | subtype);
    mgmt->duration = cpu_to_le16(0x0000);
    memcpy(mgmt->da, dst_addr, sizeof(mgmt->da));
    memcpy(mgmt->sa, src_addr, sizeof(mgmt->sa));
    memcpy(mgmt->bssid, bssid, sizeof(mgmt->bssid));
    mgmt->seq_ctrl = cpu_to_le16(0x0000);
    *len += sizeof(*mgmt);

    return 0;
}

static int buf_fill_ieee80211_beacon(uint8_t *buf, size_t *len, size_t buf_size, uint16_t interval_tu)
{
    if (*len + sizeof(struct ieee80211_beacon) > buf_size)
        return -ENOMEM;

    struct ieee80211_beacon *beacon = (struct ieee80211_beacon *)(buf + *len);
    struct timespec ts;
    uint64_t mono_us = 0;

#if defined(CLOCK_MONOTONIC)
    clock_gettime(CLOCK_MONOTONIC, &ts);
    mono_us = (uint64_t)((double) ts.tv_sec * 1e6 + (double) ts.tv_nsec * 1e-3);
#elif defined(CLOCK_REALTIME)
    clock_gettime(CLOCK_REALTIME, &ts);
    mono_us = (uint64_t)((double) ts.tv_sec * 1e6 + (double) ts.tv_nsec * 1e-3);
#elif defined(ARDUINO)
#warning "No REALTIME or MONOTONIC clock, using micros()."
    mono_us = micros();
#else
#warning "Unable to set wifi timestamp."
#endif
    beacon->timestamp = cpu_to_le64(mono_us);
    beacon->beacon_interval = cpu_to_le16(interval_tu);
    beacon->capability = cpu_to_le16(IEEE80211_CAPINFO_SHORT_SLOTTIME | IEEE80211_CAPINFO_SHORT_PREAMBLE);
    *len += sizeof(*beacon);

    return 0;
}

/* void* 不能用 可能是编译器版本的问题 改成uint8_t* 了*/
int did_GB2025_message_build_pack(UavIdentificationData *UAS_Data, uint8_t *pack, size_t buflen)
{
    if (UAS_Data == NULL || pack == NULL || buflen == 0) {
        return -1;  // 参数错误
    }
    
    UavIdentificationPacket *msg_pack_enc = (UavIdentificationPacket *)pack;
 
    
    // 1. 调用组包函数
    int8_t ret = uav_identification_pack(msg_pack_enc, UAS_Data);
    if (ret != 0) {
        return -2;  // 组包失败
    }
    
    // 2. 计算实际包大小（头部 + 内容）
    // 数据包总大小 = 固定头部(1+1+1+3) + 内容长度
    size_t total_len = sizeof(msg_pack_enc->data_type) + 
                       sizeof(msg_pack_enc->version) + 
                       sizeof(msg_pack_enc->content_len) + 
                       sizeof(msg_pack_enc->data_flag) + 
                       msg_pack_enc->content_len;
    
    // 3. 检查缓冲区大小
    if (total_len > buflen) {
        return -3;  // 缓冲区不足
    }
    
    return (int)total_len;  // 返回实际使用的字节数
}
// 里面的服务id可能后面要改 还有一些指向 open droneid的
int did_GB2025_wifi_build_nan_sync_beacon_frame(char *mac, uint8_t *buf, size_t buf_size)
{
    /* Broadcast address */
    uint8_t target_addr[6] = { 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF };
    uint8_t wifi_alliance_oui[3] = { 0x50, 0x6F, 0x9A };
    /* "org.opendroneid.remoteid" hash */
    uint8_t service_id[6] = { 0x88, 0x69, 0x19, 0x9D, 0x92, 0x09 };
    const uint8_t *cluster_id = get_nan_cluster_id();
    struct ieee80211_vendor_specific *vendor;
    struct nan_master_indication_attribute *master_indication_attr;
    struct nan_cluster_attribute *cluster_attr;
    struct nan_service_id_list_attribute *nsila;
    int ret;
    size_t len = 0;

    /* IEEE 802.11 Management Header */
    ret = buf_fill_ieee80211_mgmt(buf, &len, buf_size, IEEE80211_STYPE_BEACON, target_addr, (uint8_t *)mac, cluster_id);
    if (ret <0)
        return ret;

    /* Beacon */
    ret = buf_fill_ieee80211_beacon(buf, &len, buf_size, 0x0200);
    if (ret <0)
        return ret;

    /* Vendor Specific */
    if (len + sizeof(*vendor) > buf_size)
        return -ENOMEM;

    vendor = (struct ieee80211_vendor_specific *)(buf + len);
    memset(vendor, 0, sizeof(*vendor));
    vendor->element_id = IEEE80211_ELEMID_VENDOR;
    vendor->length = 0x22;
    memcpy(vendor->oui, wifi_alliance_oui, sizeof(vendor->oui));
    vendor->oui_type = 0x13; 
    len += sizeof(*vendor);

    /* NAN Master Indication attribute */
    if (len + sizeof(*master_indication_attr) > buf_size)
        return -ENOMEM;

    master_indication_attr = (struct nan_master_indication_attribute *)(buf + len);
    memset(master_indication_attr, 0, sizeof(*master_indication_attr));
    master_indication_attr->header.attribute_id = 0x00;
    master_indication_attr->header.length = cpu_to_le16(0x0002);
    /* Information that is used to indicate a NAN Device’s preference to serve
     * as the role of Master, with a larger value indicating a higher
     * preference. Values 1 and 255 are used for testing purposes only.
     */
    master_indication_attr->master_preference = 0xFE;
    /* Random factor value 0xEA is recommended by the European Standard */
    master_indication_attr->random_factor = 0xEA;
    len += sizeof(*master_indication_attr);

    /* NAN Cluster attribute */
    if (len + sizeof(*cluster_attr) > buf_size)
        return -ENOMEM;

    cluster_attr = (struct nan_cluster_attribute *)(buf + len);
    memset(cluster_attr, 0, sizeof(*cluster_attr));
    cluster_attr->header.attribute_id = 0x1;
    cluster_attr->header.length = cpu_to_le16(0x000D);
    memcpy(cluster_attr->device_mac, mac, sizeof(cluster_attr->device_mac));
    cluster_attr->random_factor = 0xEA;
    cluster_attr->master_preference = 0xFE;
    cluster_attr->hop_count_to_anchor_master = 0x00;
    memset(cluster_attr->anchor_master_beacon_transmission_time, 0, sizeof(cluster_attr->anchor_master_beacon_transmission_time));
    len += sizeof(*cluster_attr);

    /* NAN attributes */
    if (len + sizeof(*nsila) > buf_size)
        return -ENOMEM;

    nsila = (struct nan_service_id_list_attribute *)(buf + len);
    memset(nsila, 0, sizeof(*nsila));
    nsila->header.attribute_id = 0x02;
    nsila->header.length = cpu_to_le16(0x0006);
    memcpy(nsila->service_id, service_id, sizeof(service_id));
    len += sizeof(*nsila);

    return (int) len;
}

int did_GB2025_wifi_build_message_pack_nan_action_frame(UavIdentificationData *UAS_Data, char *mac,
                                                  uint8_t send_counter,
                                                  uint8_t *buf, size_t buf_size)
{
    /* Neighbor Awareness Networking Specification v3.0 in section 2.8.1
     * NAN网络ID要求目标MAC地址为51-6F-9A-01-00-00*/
    uint8_t target_addr[6] = { 0x51, 0x6F, 0x9A, 0x01, 0x00, 0x00 };
    /* "org.opendroneid.remoteid 哈希值 国际通用的*/
    uint8_t service_id[6] = { 0x88, 0x69, 0x19, 0x9D, 0x92, 0x09 };
    uint8_t wifi_alliance_oui[3] = { 0x50, 0x6F, 0x9A };
    const uint8_t *cluster_id = get_nan_cluster_id();
    struct nan_service_discovery *nsd;
    struct nan_service_descriptor_attribute *nsda;
    struct nan_service_descriptor_extension_attribute *nsdea;
    struct DID_GB2025_service_info *si;
    int ret;
    size_t len = 0;

    /* IEEE 802.11 Management Header */
    ret = buf_fill_ieee80211_mgmt(buf, &len, buf_size, IEEE80211_STYPE_ACTION, target_addr, (uint8_t *)mac, cluster_id);
    if (ret <0)
        return ret;

    /* NAN Service Discovery header */
    if (len + sizeof(*nsd) > buf_size)
        return -ENOMEM;

    nsd = (struct nan_service_discovery *)(buf + len);
    memset(nsd, 0, sizeof(*nsd));
    nsd->category = 0x04;               /* IEEE 802.11 Public Action frame */
    nsd->action_code = 0x09;            /* IEEE 802.11 Public Action frame Vendor Specific*/
    memcpy(nsd->oui, wifi_alliance_oui, sizeof(nsd->oui));
    nsd->oui_type = 0x13;               /* Identify Type and version of the NAN */
    len += sizeof(*nsd);

    /* NAN Attribute for Service Descriptor header */
    if (len + sizeof(*nsda) > buf_size)
        return -ENOMEM;

    nsda = (struct nan_service_descriptor_attribute *)(buf + len);
    nsda->header.attribute_id = 0x3;    /* Service Descriptor Attribute type */
    memcpy(nsda->service_id, service_id, sizeof(service_id));
    /* always 1 */
    nsda->instance_id = 0x01;           /* always 1 */
    nsda->requestor_instance_id = 0x00; /* from triggering frame */
    nsda->service_control = 0x10;       /* follow up */
    len += sizeof(*nsda);

    /* CNDID Service Info Attribute header */
    if (len + sizeof(*si) > buf_size)
        return -ENOMEM;

    si = (struct DID_GB2025_service_info *)(buf + len);
    memset(si, 0, sizeof(*si));
    si->message_counter = send_counter;
    len += sizeof(*si);

    ret = did_GB2025_message_build_pack(UAS_Data, buf + len, buf_size - len);
    if (ret < 0)
        return ret;
    len += ret;

    /* set the lengths according to the message pack lengths */
    nsda->service_info_length = sizeof(*si) + ret;
    nsda->header.length = cpu_to_le16(sizeof(*nsda) - sizeof(struct nan_attribute_header) + nsda->service_info_length);

    /* NAN Attribute for Service Descriptor extension header */
    if (len + sizeof(*nsdea) > buf_size)
        return -ENOMEM;

    nsdea = (struct nan_service_descriptor_extension_attribute *)(buf + len);
    nsdea->header.attribute_id = 0xE;
    nsdea->header.length = cpu_to_le16(0x0004);
    nsdea->instance_id = 0x01;
    nsdea->control = cpu_to_le16(0x0200);
    nsdea->service_update_indicator = send_counter;
    len += sizeof(*nsdea);

    return (int) len;
}

int did_GB2025_wifi_build_message_pack_beacon_frame(UavIdentificationData *UAS_Data, char *mac,
                                              const char *SSID, size_t SSID_len,
                                              uint16_t interval_tu, uint8_t send_counter,
                                              uint8_t *buf, size_t buf_size)
{
    /* Broadcast address */
    uint8_t target_addr[6] = { 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF };
    uint8_t asd_stan_oui[3] = { 0xFA, 0x0B, 0xBC }; // 国标也是0xFA0BBC
    /* Mgmt Beacon frame mandatory fields + IE 221 */
    struct ieee80211_ssid *ssid_s;
    struct ieee80211_supported_rates *rates;
    struct ieee80211_vendor_specific *vendor;

    /* Message Pack */
    struct DID_GB2025_service_info *si;

    int ret;
    size_t len = 0;

    /* IEEE 802.11 Management Header */
    ret = buf_fill_ieee80211_mgmt(buf, &len, buf_size, IEEE80211_STYPE_BEACON, target_addr, (uint8_t *)mac, (uint8_t *)mac);
    if (ret <0)
        return ret;

    /* Mandatory Beacon as of 802.11-2016 Part 11 */
    ret = buf_fill_ieee80211_beacon(buf, &len, buf_size, interval_tu);
    if (ret <0)
        return ret;

    /* SSID: 1-32 bytes */
    if (len + sizeof(*ssid_s) > buf_size)
        return -ENOMEM;

    ssid_s = (struct ieee80211_ssid *)(buf + len);
    if(!SSID || (SSID_len ==0) || (SSID_len > 32))
        return -EINVAL;
    ssid_s->element_id = IEEE80211_ELEMID_SSID;
    ssid_s->length = (uint8_t) SSID_len;
    memcpy(ssid_s->ssid, SSID, ssid_s->length);
    len += sizeof(*ssid_s) + SSID_len;

    /* Supported Rates: 1 record at minimum */
    if (len + sizeof(*rates) > buf_size)
        return -ENOMEM;

    rates = (struct ieee80211_supported_rates *)(buf + len);
    rates->element_id = IEEE80211_ELEMID_RATES;
    rates->length = 1; // One rate only
    rates->supported_rates = 0x8C;     // 6 Mbps
    len += sizeof(*rates);

    /* Vendor Specific Information Element (IE 221) */
    if (len + sizeof(*vendor) > buf_size)
        return -ENOMEM;

    vendor = (struct ieee80211_vendor_specific *)(buf + len);
    vendor->element_id = IEEE80211_ELEMID_VENDOR;
    vendor->length = 0x00;  // Length updated at end of function
    memcpy(vendor->oui, asd_stan_oui, sizeof(vendor->oui));
    vendor->oui_type = 0x0D;
    len += sizeof(*vendor);

    /* DID GB2025 Service Info Attribute header */
    if (len + sizeof(*si) > buf_size)
        return -ENOMEM;

    si = (struct DID_GB2025_service_info *)(buf + len);
    memset(si, 0, sizeof(*si));
    si->message_counter = send_counter;
    len += sizeof(*si);

    ret = did_GB2025_message_build_pack(UAS_Data, buf + len, buf_size - len);
    if (ret < 0)
        return ret;
    len += ret;

    /* set the lengths according to the message pack lengths */
    vendor->length = sizeof(vendor->oui) + sizeof(vendor->oui_type) + sizeof(*si) + ret; // ret 是个负的表示空闲字节数

    return (int) len;
}
// 解包
int did_GB2025_message_process_pack(UavIdentificationData  *UAS_Data, uint8_t *pack, size_t buflen)
{
    if (UAS_Data == NULL || pack == NULL || buflen == 0) {
        return -1;
    }
    
    UavIdentificationPacket *msg_pack_enc = (UavIdentificationPacket *) pack;
    
    // 固定头部大小：data_type(1) + version(1) + content_len(1) + data_flag(3) = 6字节
    const size_t FIXED_HEADER_SIZE = 6;
    
    // 1. 首先检查是否能安全读取固定头部
    if (buflen < FIXED_HEADER_SIZE) {
        return -2;  // 缓冲区太小，连头部都不够
    }
    
    // 2. 现在可以安全地读取 content_len
    uint8_t content_len = msg_pack_enc->content_len;
    
    // 3. 检查 content_len 是否合理
    if (content_len > MAX_CONTENT_BYTES) {
        return -3;  // 内容长度超限
    }
    
    // 4. 计算完整的包大小
    size_t total_size = FIXED_HEADER_SIZE + content_len;
    
    // 5. 验证整个包是否在缓冲区范围内
    if (total_size > buflen) {
        return -4;  // 包不完整
    }
    
    // 6. 检查数据类型（可选）
    if (msg_pack_enc->data_type != PACKET_DATA_TYPE) {
        return -5;  // 数据类型错误
    }
    
    // 7. 初始化解包目标
    uav_identification_data_init(UAS_Data);
    
    // 8. 执行解包
    if (uav_identification_unpack(UAS_Data, msg_pack_enc) != 0) {
        return -6;  // 解包失败
    }
    
    // 9. 返回实际处理的字节数
    return (int)total_size;
}

int did_GB2025_wifi_receive_message_pack_nan_action_frame(UavIdentificationData  *UAS_Data,
                                                    char *mac, uint8_t *buf, size_t buf_size)
{
    struct ieee80211_mgmt *mgmt; // 802.11管理帧头
    struct nan_service_discovery *nsd; // NAN服务发现帧头
    struct nan_service_descriptor_attribute *nsda; // NAN服务描述符属性
    struct nan_service_descriptor_extension_attribute *nsdea;  // NAN服务描述符扩展属性
    struct DID_GB2025_service_info *si; // 中国远程ID服务信息
    uint8_t target_addr[6] = { 0x51, 0x6F, 0x9A, 0x01, 0x00, 0x00 };  // 远程ID（Remote ID）分配的专用组播地址
    uint8_t wifi_alliance_oui[3] = { 0x50, 0x6F, 0x9A }; // Wi-Fi联盟OUI
    uint8_t service_id[6] = { 0x88, 0x69, 0x19, 0x9D, 0x92, 0x09 }; // 远程ID服务ID 国际通用的 不需要改
    int ret;
    size_t len = 0;

    /* IEEE 802.11 Management Header */
    if (len + sizeof(*mgmt) > buf_size) // 检查缓冲区是否足够容纳一个管理帧头
        return -EINVAL;
    mgmt = (struct ieee80211_mgmt *)(buf + len);
    if ((mgmt->frame_control & cpu_to_le16(IEEE80211_FCTL_FTYPE | IEEE80211_FCTL_STYPE)) !=
        cpu_to_le16(IEEE80211_FTYPE_MGMT | IEEE80211_STYPE_ACTION))
        return -EINVAL;
    if (memcmp(mgmt->da, target_addr, sizeof(mgmt->da)) != 0)
        return -EINVAL;
    memcpy(mac, mgmt->sa, sizeof(mgmt->sa));

    len += sizeof(*mgmt);

    /* NAN Service Discovery header */
    if (len + sizeof(*nsd) > buf_size)
        return -EINVAL;
    nsd = (struct nan_service_discovery *)(buf + len);
    if (nsd->category != 0x04)
        return -EINVAL;
    if (nsd->action_code != 0x09)
        return -EINVAL;
    if (memcmp(nsd->oui, wifi_alliance_oui, sizeof(wifi_alliance_oui)) != 0)
        return -EINVAL;
    if (nsd->oui_type != 0x13)
        return -EINVAL;
    len += sizeof(*nsd);

    /* NAN Attribute for Service Descriptor header */
    if (len + sizeof(*nsda) > buf_size)
        return -EINVAL;
    nsda = (struct nan_service_descriptor_attribute *)(buf + len);
    if (nsda->header.attribute_id != 0x3)
        return -EINVAL;
    if (memcmp(nsda->service_id, service_id, sizeof(service_id)) != 0)
        return -EINVAL;
    if (nsda->instance_id != 0x01)
        return -EINVAL;
    if (nsda->service_control != 0x10)
        return -EINVAL;
    len += sizeof(*nsda);

    // 在读取 si 之前应该检查空间
    if (len + sizeof(*si) > buf_size)
        return -EINVAL;

    si = (struct DID_GB2025_service_info *)(buf + len);
    ret = did_GB2025_message_process_pack(UAS_Data, buf + len + sizeof(*si), buf_size - len - sizeof(*nsdea));
    if (ret < 0)
        return -EINVAL;
    if (nsda->service_info_length != (sizeof(*si) + ret))
        return -EINVAL;
    if (nsda->header.length != (cpu_to_le16(sizeof(*nsda) - sizeof(struct nan_attribute_header) + nsda->service_info_length)))
        return -EINVAL;
    len += sizeof(*si) + ret;

    /* NAN Attribute for Service Descriptor extension header */
    if (len + sizeof(*nsdea) > buf_size)
        return -ENOMEM;
    nsdea = (struct nan_service_descriptor_extension_attribute *)(buf + len);
    if (nsdea->header.attribute_id != 0xE)
        return -EINVAL;
    if (nsdea->header.length != cpu_to_le16(0x0004))
        return -EINVAL;
    if (nsdea->instance_id != 0x01)
        return -EINVAL;
    if (nsdea->control != cpu_to_le16(0x0200))
        return -EINVAL;

    return 0;
}