#include "air.hpp"

#include <cstring>

#include "armor_time.hpp"
#include "esp_event.h"
#include "esp_log.h"
#include "esp_netif.h"
#include "esp_wifi.h"
#include "freertos/FreeRTOS.h"
#include "freertos/task.h"
#include "storage.hpp"

namespace {

constexpr const char *TAG = "air";

portMUX_TYPE s_dbus_lock = portMUX_INITIALIZER_UNLOCKED;

// ESP-NOW payloads are 802.11 frames; the promiscuous callback is used to keep
// track of the peer's signal strength.
typedef struct {
    unsigned frame_ctrl : 16;
    unsigned duration_id : 16;
    uint8_t addr1[6]; /* receiver address */
    uint8_t addr2[6]; /* sender address */
    uint8_t addr3[6]; /* filtering address */
    unsigned sequence_ctrl : 16;
    uint8_t addr4[6]; /* optional */
} wifi_ieee80211_mac_hdr_t;

typedef struct {
    wifi_ieee80211_mac_hdr_t hdr;
    uint8_t payload[0]; /* network data ended with 4 bytes csum (CRC32) */
} wifi_ieee80211_packet_t;

} // namespace

// public
Air::packet_t Air::my_packet, Air::peer_packet;
volatile uint32_t Air::last_air_ms = 0;
volatile bool Air::is_pairing = false;
volatile bool Air::pairing_found = false;
// private
const uint8_t Air::empty_mac[6] = {0x00, 0x00, 0x00, 0x00, 0x00, 0x00};
const uint8_t Air::pairing_data[4] = {0x11, 0x45, 0x14, 0x00};
esp_now_peer_info_t Air::broadcast_peer;
esp_now_peer_info_t Air::peer;
volatile int8_t Air::pairing_best_rssi = -128;

void Air::begin() {
    // ESP-NOW needs the WiFi stack up on the station interface.
    ESP_ERROR_CHECK(esp_netif_init());
    ESP_ERROR_CHECK(esp_event_loop_create_default());

    wifi_init_config_t wifi_cfg = WIFI_INIT_CONFIG_DEFAULT();
    ESP_ERROR_CHECK(esp_wifi_init(&wifi_cfg));
    // The peer MAC lives in our own NVS namespace; don't let the WiFi driver
    // write extra calibration data it does not need here.
    ESP_ERROR_CHECK(esp_wifi_set_storage(WIFI_STORAGE_RAM));
    ESP_ERROR_CHECK(esp_wifi_set_mode(WIFI_MODE_STA));
    ESP_ERROR_CHECK(esp_wifi_start());

    // 使用 WiFi 14 信道
    wifi_country_t country = {};
    memcpy(country.cc, "JP", sizeof(country.cc));
    country.schan = 1;
    country.nchan = 14;
    country.policy = WIFI_COUNTRY_POLICY_MANUAL;
    ESP_ERROR_CHECK(esp_wifi_set_country(&country));
    ESP_ERROR_CHECK(esp_wifi_set_channel(14, WIFI_SECOND_CHAN_NONE));

    // 20dBm 功率（84 = 21dBm；C3 上限为 20dBm，失败时退回 80）
    esp_err_t power_err = esp_wifi_set_max_tx_power(84);
    if (power_err != ESP_OK) {
        ESP_LOGW(TAG, "set_max_tx_power(84) failed: %s, retrying with 80",
                 esp_err_to_name(power_err));
        power_err = esp_wifi_set_max_tx_power(80);
        if (power_err != ESP_OK) {
            ESP_LOGW(TAG, "set_max_tx_power(80) failed: %s", esp_err_to_name(power_err));
        }
    }

    // 启用混杂模式，以获取发送方 rssi
    wifi_promiscuous_filter_t filter = {};
    filter.filter_mask = WIFI_PROMIS_FILTER_MASK_MGMT;
    ESP_ERROR_CHECK(esp_wifi_set_promiscuous_filter(&filter));
    ESP_ERROR_CHECK(esp_wifi_set_promiscuous_rx_cb(promiscuous_rx_cb));
    ESP_ERROR_CHECK(esp_wifi_set_promiscuous(true));

    // 初始化 ESP-NOW
    ESP_ERROR_CHECK(esp_now_init());
    ESP_ERROR_CHECK(esp_now_register_recv_cb(onRecv));

    // 添加广播 peer
    memset(&broadcast_peer, 0, sizeof(broadcast_peer));
    memset(broadcast_peer.peer_addr, 0xFF, sizeof(broadcast_peer.peer_addr));
    broadcast_peer.channel = 0;
    broadcast_peer.ifidx = WIFI_IF_STA;
    broadcast_peer.encrypt = false;
    ESP_ERROR_CHECK(esp_now_add_peer(&broadcast_peer));

    memset(&peer, 0, sizeof(peer));
    peer.channel = 0;
    peer.ifidx = WIFI_IF_STA;
    peer.encrypt = false;

    if (Storage::hasKey("peer_mac")) {
        Storage::getBytes("peer_mac", peer.peer_addr, sizeof(peer.peer_addr));
        if (memcmp(peer.peer_addr, empty_mac, sizeof(empty_mac)) == 0) {
            ESP_LOGI(TAG, "stored peer MAC is empty, staying unpaired");
        } else {
            esp_err_t err = esp_now_add_peer(&peer);
            if (err != ESP_OK) {
                ESP_LOGW(TAG, "esp_now_add_peer(stored) failed: %s", esp_err_to_name(err));
            }
        }
    } else {
        memset(peer.peer_addr, 0x00, sizeof(peer.peer_addr));
    }
}

void Air::entryPairing() {
    // 清除之前的配对
    esp_now_del_peer(peer.peer_addr);

    is_pairing = true;
    pairing_found = false;

    // 用于只配对最强的信号
    my_packet.rssi = -128;
    pairing_best_rssi = -128;
}

void Air::savePairing() {
    if (pairing_found) {
        // 添加 peer
        esp_err_t err = esp_now_add_peer(&peer);
        if (err != ESP_OK) {
            ESP_LOGW(TAG, "esp_now_add_peer(pairing) failed: %s", esp_err_to_name(err));
        }
    } else {
        // 配对失败，清空设备 mac
        memset(peer.peer_addr, 0x00, sizeof(peer.peer_addr));
    }

    // 写入 flash
    Storage::putBytes("peer_mac", peer.peer_addr, sizeof(peer.peer_addr));

    is_pairing = false;
}

void Air::onLoop() {
    if (is_pairing) {
        // 发送配对请求
        esp_now_send(broadcast_peer.peer_addr, pairing_data, sizeof(pairing_data));
    } else if (memcmp(peer.peer_addr, empty_mac, sizeof(empty_mac)) != 0) {
        // 已有配对设备
        esp_now_send(peer.peer_addr, reinterpret_cast<const uint8_t *>(&my_packet), sizeof(packet_t));
    }
}

void Air::dbusSet(const uint8_t *src, size_t len) {
    portENTER_CRITICAL(&s_dbus_lock);
    memcpy(my_packet.dbus, src, len);
    portEXIT_CRITICAL(&s_dbus_lock);
}

void Air::dbusClear() {
    portENTER_CRITICAL(&s_dbus_lock);
    memset(my_packet.dbus, 0, sizeof(my_packet.dbus));
    portEXIT_CRITICAL(&s_dbus_lock);
}

void Air::onRecv(const esp_now_recv_info_t *info, const uint8_t *data, int data_len) {
    if (info == nullptr || info->src_addr == nullptr || data == nullptr) {
        return;
    }

    // The original relied on the promiscuous callback to have recorded the RSSI
    // of this frame beforehand; rx_ctrl carries the same value synchronously
    // with the packet we are handling.
    if (info->rx_ctrl != nullptr) {
        my_packet.rssi = info->rx_ctrl->rssi;
    }

    if (is_pairing) {
        // 配对 magic 正确
        if (data_len == static_cast<int>(sizeof(pairing_data)) &&
            memcmp(pairing_data, data, sizeof(pairing_data)) == 0) {
            // 只配对最强的信号
            if (my_packet.rssi > pairing_best_rssi) {
                pairing_best_rssi = my_packet.rssi;
                memcpy(peer.peer_addr, info->src_addr, sizeof(peer.peer_addr));
                pairing_found = true;
            }
        }
    } else if (memcmp(peer.peer_addr, info->src_addr, sizeof(peer.peer_addr)) == 0) {
        // 是从配对设备发过来的，且数据长度正确
        if (data_len == static_cast<int>(sizeof(packet_t))) {
            memcpy(&peer_packet, data, sizeof(packet_t));
            last_air_ms = millis();
        }
    }
}

void Air::promiscuous_rx_cb(void *buf, wifi_promiscuous_pkt_type_t type) {
    if (type != WIFI_PKT_MGMT) {
        return;
    }

    const auto *ppkt = static_cast<const wifi_promiscuous_pkt_t *>(buf);
    const auto *ipkt = reinterpret_cast<const wifi_ieee80211_packet_t *>(ppkt->payload);
    const wifi_ieee80211_mac_hdr_t *hdr = &ipkt->hdr;

    if (is_pairing) {
        // 配对模式
        my_packet.rssi = ppkt->rx_ctrl.rssi;
    } else if (memcmp(hdr->addr2, peer.peer_addr, sizeof(peer.peer_addr)) == 0) {
        my_packet.rssi = ppkt->rx_ctrl.rssi;
    }
}