// ============================================================================
// network_id.cpp — 802.11 frame-parser / network identification
//
// See network_id.h for full API and wire-format documentation.
// ============================================================================

#include "network_id.h"

// Under UNIT_TEST the ESP-IDF Wi-Fi types / calls are stubbed; on device they
// come from the Arduino ESP32 framework.
#ifndef UNIT_TEST
#  include "esp_wifi.h"
#  include "esp_wifi_types.h"
#  include "Arduino.h"        // millis()
#  include "mesh_publish.h"   // mesh_publish_raw()
#else
// ---- Host-side stubs (populated by test harness) -------------------------
#  include <cstdlib>
#  include <cstring>
#  include <cstdio>
// Minimal type mirrors so this TU compiles on the native platform.
// (the enum tag is now provided by network_id.h under UNIT_TEST)
#ifndef WIFI_PROMISCUOUS_PKT_TYPE_T_DEFINED
#define WIFI_PROMISCUOUS_PKT_TYPE_T_DEFINED
typedef enum { WIFI_PKT_MGMT = 0, WIFI_PKT_CTRL = 1, WIFI_PKT_DATA = 2, WIFI_PKT_MISC = 3 } wifi_promiscuous_pkt_type_t;
#else
// Header already defined the type — provide alias constants the .cpp expects.
static const wifi_promiscuous_pkt_type_t WIFI_PKT_MGMT = NETWORK_ID_PKT_MGMT;
static const wifi_promiscuous_pkt_type_t WIFI_PKT_CTRL = NETWORK_ID_PKT_CTRL;
static const wifi_promiscuous_pkt_type_t WIFI_PKT_DATA = NETWORK_ID_PKT_DATA;
static const wifi_promiscuous_pkt_type_t WIFI_PKT_MISC = NETWORK_ID_PKT_MISC;
#endif
struct wifi_pkt_rx_ctrl_t {
    int8_t   rssi;
    uint8_t  channel;
    uint16_t sig_len;   // mirrors real esp_wifi_types.h's sig_len bitfield:
                         // total received frame length in payload[].
};
struct wifi_promiscuous_pkt_t {
    wifi_pkt_rx_ctrl_t rx_ctrl;
    uint8_t payload[0];
};
// Stub millis() — provided by the test harness or falls back to 0.
#  ifndef STUB_MILLIS_PROVIDED
static uint32_t stub_millis_val = 1000;
static inline uint32_t millis() { return stub_millis_val; }
#  endif
// mesh_publish_raw stub — declared extern so the test can intercept.
extern bool mesh_publish_raw(const uint8_t* data, size_t len);
// legacy_rssi_observe stub
__attribute__((weak)) void legacy_rssi_observe(uint8_t channel, int8_t rssi) {
    (void)channel; (void)rssi;
}
#endif // UNIT_TEST

#include <cstring>
#include <cstdlib>

// ---------------------------------------------------------------------------
// Legacy RSSI-bin shim (weak symbol — real implementation in rf_classifier.cpp
// or wherever the old per-channel accounting lived). When network_id receives
// a non-management frame it calls this to preserve prior behaviour.
// ---------------------------------------------------------------------------
#ifndef UNIT_TEST
extern void legacy_rssi_observe(uint8_t channel, int8_t rssi) __attribute__((weak));
void legacy_rssi_observe(uint8_t channel, int8_t rssi) {
    // No-op if the original sniffer callback is not present.
    (void)channel; (void)rssi;
}
#endif

// ---------------------------------------------------------------------------
// Internal state
// ---------------------------------------------------------------------------

static NetworkEntry s_table[NETWORK_ID_TABLE_SIZE];
static int          s_count = 0;   // number of valid entries

// Helper: find entry by BSSID; returns index or -1.
static int find_by_bssid(const uint8_t* bssid) {
    for (int i = 0; i < s_count; i++) {
        if (memcmp(s_table[i].bssid, bssid, 6) == 0) return i;
    }
    return -1;
}

// Helper: evict the entry with the smallest last_seen_ms and return its index.
static int evict_lru() {
    int oldest = 0;
    for (int i = 1; i < NETWORK_ID_TABLE_SIZE; i++) {
        if (s_table[i].last_seen_ms < s_table[oldest].last_seen_ms) oldest = i;
    }
    return oldest;
}

// ---------------------------------------------------------------------------
// 802.11 frame-control field helpers
//   Frame Control Word (2 bytes, little-endian):
//     bits [1:0]  = version (always 0)
//     bits [3:2]  = type   (0=mgmt, 1=ctrl, 2=data)
//     bits [7:4]  = subtype
// ---------------------------------------------------------------------------
#define FCF_TYPE(fcf)    (((fcf) >> 2) & 0x03u)
#define FCF_SUBTYPE(fcf) (((fcf) >> 4) & 0x0Fu)

#define FRAME_TYPE_MGMT  0x00u
#define FRAME_TYPE_CTRL  0x01u
#define FRAME_TYPE_DATA  0x02u

#define MGMT_SUBTYPE_BEACON        0x08u
#define MGMT_SUBTYPE_PROBE_RESP    0x05u

// 802.11 management frame layout (fixed header, bytes):
//   [0..1]   Frame Control
//   [2..3]   Duration
//   [4..9]   Address 1 (destination / RA)
//   [10..15] Address 2 (source / TA)
//   [16..21] Address 3 (BSSID for beacon/probe-resp)
//   [22..23] Sequence Control
//   [24..]   Frame body (beacon: fixed 12B params + tagged params)
#define MGMT_HDR_LEN     24
#define MGMT_ADDR3_OFF   16   // BSSID offset in management frame

// Beacon / probe-response fixed params (12 bytes after 24B header):
//   [0..7]  timestamp
//   [8..9]  beacon interval
//   [10..11] capability info
#define BEACON_FIXED_LEN 12

// Tagged parameter IE IDs we care about
#define IE_SSID       0x00u
#define IE_DS_PARAM   0x03u
#define IE_RSN        0x30u

// RSN cipher suite selectors (last byte of 4-byte OUI+suite)
#define RSN_CIPHER_TKIP    0x02u
#define RSN_CIPHER_CCMP    0x04u

// Parse tagged parameters from a raw IE blob.
// ie_buf points to the first IE; ie_len is the remaining bytes.
// Fills ssid, channel, cipher_id in the supplied NetworkEntry.
static void parse_ies(const uint8_t* ie_buf, size_t ie_len, NetworkEntry* entry) {
    size_t pos = 0;
    bool ssid_found = false;  // only parse the first SSID IE; subsequent ones
                               // (including garbage from over-read) are ignored
    while (pos + 2 <= ie_len) {
        uint8_t id  = ie_buf[pos];
        uint8_t len = ie_buf[pos + 1];
        pos += 2;
        if (pos + len > ie_len) break;  // truncated IE, stop

        if (id == IE_SSID && !ssid_found) {
            // SSID: up to 32 bytes, not null-terminated on wire.
            // Only the first SSID IE is stored; later re-advertisements of
            // the same IE (e.g. in multi-BSSID frames or garbage over-reads)
            // are ignored so a zero-length follow-on IE cannot wipe the name.
            uint8_t slen = (len < 32) ? len : 32;
            memcpy(entry->ssid, ie_buf + pos, slen);
            entry->ssid[slen] = '\0';
            ssid_found = true;
        } else if (id == IE_DS_PARAM && len == 1) {
            entry->channel = ie_buf[pos];
        } else if (id == IE_RSN && len >= 8) {
            // RSN IE layout:
            //   [0..1]  version
            //   [2..5]  group cipher suite (OUI 3B + suite 1B)
            //   [6..7]  pairwise count
            //   [8..]   pairwise cipher suites (4B each)
            uint8_t group_suite = ie_buf[pos + 3];
            if (group_suite == RSN_CIPHER_TKIP) entry->cipher_id |= CIPHER_TKIP;
            if (group_suite == RSN_CIPHER_CCMP) entry->cipher_id |= CIPHER_CCMP;
            // Inspect first pairwise if present
            uint16_t pw_count;
            memcpy(&pw_count, ie_buf + pos + 6, 2);
            if (pw_count >= 1 && len >= 12) {
                uint8_t pw_suite = ie_buf[pos + 11];
                if (pw_suite == RSN_CIPHER_TKIP) entry->cipher_id |= CIPHER_TKIP;
                if (pw_suite == RSN_CIPHER_CCMP) entry->cipher_id |= CIPHER_CCMP;
            }
            // If RSN present but no recognized cipher, still mark as non-open
            if (entry->cipher_id == CIPHER_UNKNOWN) entry->cipher_id = CIPHER_CCMP;
        }
        pos += len;
    }
}

// Upsert a parsed entry into the hash table.
static void upsert(NetworkEntry* incoming) {
    int idx = find_by_bssid(incoming->bssid);
    if (idx < 0) {
        // New entry
        if (s_count < NETWORK_ID_TABLE_SIZE) {
            idx = s_count++;
        } else {
            idx = evict_lru();
        }
        memcpy(&s_table[idx], incoming, sizeof(NetworkEntry));
    } else {
        // Update existing — preserve SSID if new parse came up empty
        s_table[idx].channel      = incoming->channel;
        s_table[idx].last_rssi    = incoming->last_rssi;
        s_table[idx].last_seen_ms = incoming->last_seen_ms;
        s_table[idx].cipher_id    = incoming->cipher_id;
        if (incoming->ssid[0] != '\0') {
            memcpy(s_table[idx].ssid, incoming->ssid, NETWORK_ID_SSID_MAX);
        }
    }
}

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

void network_id_init() {
    memset(s_table, 0, sizeof(s_table));
    s_count = 0;
}

void network_id_sniffer_cb(void* buf, wifi_promiscuous_pkt_type_t type) {
    if (!buf) return;
    const wifi_promiscuous_pkt_t* pkt =
        reinterpret_cast<const wifi_promiscuous_pkt_t*>(buf);

    int8_t  rssi    = pkt->rx_ctrl.rssi;
    uint8_t channel = pkt->rx_ctrl.channel;

    if (type != WIFI_PKT_MGMT) {
        // Non-management frame: hand off to legacy per-channel accounting.
        legacy_rssi_observe(channel, rssi);
        return;
    }

    // Management frame — need at least the fixed 24-byte header.
    // The promiscuous payload includes the full MPDU (no FCS on ESP32).
    // We don't have the length from the struct alone; rely on rx_ctrl fields.
    // Use a conservative minimum: header + beacon fixed params.
    const uint8_t* frame = pkt->payload;
    // Minimum required: 24B header + 12B beacon fixed = 36B
    // We read up to 350B of IEs safely.

    uint16_t fcw;
    memcpy(&fcw, frame, 2);
    uint8_t ftype   = FCF_TYPE(fcw);
    uint8_t fsubtype = FCF_SUBTYPE(fcw);

    if (ftype != FRAME_TYPE_MGMT) {
        legacy_rssi_observe(channel, rssi);
        return;
    }

    if (fsubtype != MGMT_SUBTYPE_BEACON && fsubtype != MGMT_SUBTYPE_PROBE_RESP) {
        // Other mgmt (assoc, auth, etc.) — still fall through to rssi observe
        legacy_rssi_observe(channel, rssi);
        return;
    }

    NetworkEntry entry;
    memset(&entry, 0, sizeof(entry));
    entry.last_rssi    = rssi;
    entry.last_seen_ms = (uint32_t)millis();
    entry.channel      = channel;  // default from rx_ctrl; IE may override
    entry.cipher_id    = CIPHER_UNKNOWN;

    // BSSID = Address 3 (bytes 16-21)
    memcpy(entry.bssid, frame + MGMT_ADDR3_OFF, 6);

    // IEs start at byte 24 (header) + 12 (fixed beacon params) = 36
    const size_t ie_start_off = MGMT_HDR_LEN + BEACON_FIXED_LEN;
    const uint8_t* ie_start = frame + ie_start_off;
    // BUG FIX: this previously ALWAYS capped at a flat 350 bytes regardless
    // of the real received frame length -- pkt->rx_ctrl.sig_len (the ESP-IDF
    // WiFi driver's actual reported frame length) was available the whole
    // time but never read. Since most real beacons carry well under 350
    // bytes of IEs, this meant parse_ies() was routinely reading PAST THE
    // END of the real captured frame into whatever memory happened to sit
    // next in the DMA/stack buffer -- occasionally landing on bytes that
    // happen to look like a valid RSN IE (tag 0x30) and spuriously
    // classifying an open/WEP network as WPA2/WPA3-CCMP. Bound the parse to
    // the real frame length instead; 350 remains only as a defensive
    // ceiling in case sig_len itself is ever bogus (0 or implausibly huge).
    size_t ie_len_cap = 350;
    uint16_t sig_len = pkt->rx_ctrl.sig_len;
    if (sig_len > ie_start_off && sig_len < 2000) {
        size_t avail = (size_t)sig_len - ie_start_off;
        if (avail < ie_len_cap) ie_len_cap = avail;
    }
    parse_ies(ie_start, ie_len_cap, &entry);

    // BUG FIX (found while writing wifi_security_scan's first native tests):
    // no RSN IE (the only thing parse_ies inspects) does NOT mean "open" --
    // WEP networks never advertise an RSN IE either (RSN is WPA/WPA2/WPA3-
    // only), so a WEP network was silently defaulting to CIPHER_OPEN here,
    // even though CIPHER_WEP is a real, distinct flag that wifi_security_
    // scan.cpp has been checking for since it was written (network_id.h's
    // own header comment already claimed this parses WEP -- it never did).
    // Fix: fall back to the beacon's Capability Info "Privacy" bit (802.11
    // Section 9.4.1.4, bit 0 of the 2-byte field at the end of the fixed
    // beacon params) to distinguish the two -- Privacy set + no RSN IE =
    // WEP; Privacy clear + no RSN IE = genuinely open.
    if (entry.cipher_id == CIPHER_UNKNOWN) {
        // Fixed beacon params: timestamp[8] + beacon_interval[2] + capability[2].
        // Capability is the LAST 2 bytes -> offset = MGMT_HDR_LEN + 8 + 2 = +10.
        uint16_t capability_info;
        memcpy(&capability_info, frame + MGMT_HDR_LEN + 10, 2);
        bool privacy_bit = (capability_info & 0x0001u) != 0;
        entry.cipher_id = privacy_bit ? CIPHER_WEP : CIPHER_OPEN;
    }

    upsert(&entry);
}

int network_id_count() {
    return s_count;
}

// Simple index-based accessor — no re-sort; callers can sort by last_seen_ms
// externally if needed. Index is stable until next upsert.
bool network_id_get(int idx, NetworkEntry* out) {
    if (!out || idx < 0 || idx >= s_count) return false;
    memcpy(out, &s_table[idx], sizeof(NetworkEntry));
    return true;
}

bool network_id_publish(uint32_t sensor_id) {
    if (s_count == 0) return false;

    // Packet layout:
    //   [0]     = MSG_TYPE_NETWORK_LIST (0x15)
    //   [1..4]  = sensor_id LE
    //   [5]     = entry count in this packet
    //   [6..]   = entries (16B each): bssid(6) + ssid_first8(8) + channel(1) + rssi(1)
    // Max packet size: 1 + 4 + 1 + 4*16 = 70 bytes — well under LoRa MTU.

    static const size_t HDR_LEN   = 6;
    static const size_t ENTRY_LEN = NETWORK_ID_ENTRY_BYTES;
    uint8_t pkt[HDR_LEN + NETWORK_ID_ENTRIES_PER_PKT * ENTRY_LEN];

    bool any_sent = false;
    int  idx      = 0;

    while (idx < s_count) {
        int batch = s_count - idx;
        if (batch > NETWORK_ID_ENTRIES_PER_PKT) batch = NETWORK_ID_ENTRIES_PER_PKT;

        // Build header
        pkt[0] = MSG_TYPE_NETWORK_LIST;
        pkt[1] = (uint8_t)(sensor_id & 0xFF);
        pkt[2] = (uint8_t)((sensor_id >> 8)  & 0xFF);
        pkt[3] = (uint8_t)((sensor_id >> 16) & 0xFF);
        pkt[4] = (uint8_t)((sensor_id >> 24) & 0xFF);
        pkt[5] = (uint8_t)batch;

        // Build entries
        for (int b = 0; b < batch; b++) {
            uint8_t* ep = pkt + HDR_LEN + b * ENTRY_LEN;
            const NetworkEntry& e = s_table[idx + b];
            memcpy(ep, e.bssid, 6);              // [0..5] bssid
            // SSID first 8 chars, zero-padded
            size_t slen = strlen(e.ssid);
            if (slen > 8) slen = 8;
            memcpy(ep + 6, e.ssid, slen);
            if (slen < 8) memset(ep + 6 + slen, 0, 8 - slen);
            ep[14] = e.channel;
            ep[15] = (uint8_t)(int8_t)e.last_rssi;
        }

        size_t pkt_len = HDR_LEN + (size_t)batch * ENTRY_LEN;
        if (mesh_publish_raw(pkt, pkt_len)) any_sent = true;
        idx += batch;
    }
    return any_sent;
}
