// ============================================================================
// wifi_odid_scanner.cpp — see wifi_odid_scanner.h for design notes.
//
// Section 1 (pure parsing) is a direct port of PDL Witness's pdl_witness.c
// classify_oui / find_droneid_vendor_ie logic, restructured to fill
// WifiOdidParseResult and to reuse ble_odid_scanner.h's message-body decoders
// instead of duplicating them.
//
// Section 2 (ESP32 integration) is new — extracts the raw frame from an
// ESP-IDF wifi_promiscuous_pkt_t and feeds it through the same
// blue_force/FFPS detection pipeline the BLE and RF-classifier paths use.
// ============================================================================

#include "wifi_odid_scanner.h"

#include <cstring>

// ============================================================================
// SECTION 1 — Pure parsing (no ESP-IDF dependency)
// ============================================================================

// 802.11 Frame Control low byte for a Beacon frame (type=Management,
// subtype=Beacon: 00 1000 00 -> 0x80).
#define WIFI_FRAME_CTRL_BEACON  0x80

// 802.11 IE tag for Vendor Specific.
#define WIFI_IE_TAG_VENDOR      0xDD

// Fixed beacon header + fixed parameters length before the IE list begins
// (frame ctrl(2) + duration(2) + da(6) + sa(6) + bssid(6) + seq(2) +
//  timestamp(8) + beacon interval(2) + capability(2) = 36 bytes).
#define WIFI_BEACON_IE_START    36

// Known drone manufacturer OUIs (3 bytes each), matching PDL Witness's table.
struct VendorOui { uint8_t b0, b1, b2; uint8_t vendor; };
static const VendorOui kVendorOuis[] = {
    { 0x60, 0x60, 0x1F, ODID_VENDOR_DJI    },
    { 0x70, 0x38, 0x69, ODID_VENDOR_AUTEL  },
    { 0x84, 0x38, 0x8E, ODID_VENDOR_AUTEL  },
    { 0x04, 0xE5, 0x48, ODID_VENDOR_SKYDIO },
    { 0x90, 0x3A, 0xE6, ODID_VENDOR_PARROT },
    { 0xA0, 0x14, 0x3D, ODID_VENDOR_PARROT },
};

// Locate a drone-vendor IE within a beacon frame body. Returns a pointer to
// the OUI's first byte (IE payload, after the tag/len header) and its
// length, or nullptr if no known drone OUI is present.
static const uint8_t* find_vendor_ie(const uint8_t* p, size_t len,
                                      uint16_t* out_ie_len, uint8_t* out_vendor) {
    if (!p || len < WIFI_BEACON_IE_START || p[0] != WIFI_FRAME_CTRL_BEACON) return nullptr;

    for (size_t i = WIFI_BEACON_IE_START; i + 5 < len; i++) {
        if (p[i] != WIFI_IE_TAG_VENDOR) continue;
        uint8_t ie_len = p[i + 1];
        if (i + 2 + ie_len > len) break;
        const uint8_t* o = &p[i + 2];

        for (const auto& v : kVendorOuis) {
            if (o[0] == v.b0 && o[1] == v.b1 && o[2] == v.b2) {
                *out_ie_len = ie_len;
                *out_vendor = v.vendor;
                return o;
            }
        }
    }
    return nullptr;
}

uint8_t wifi_odid_classify_oui(const uint8_t* frame, size_t frame_len) {
    uint16_t ie_len = 0;
    uint8_t vendor = ODID_VENDOR_UNKNOWN;
    if (find_vendor_ie(frame, frame_len, &ie_len, &vendor)) return vendor;
    return ODID_VENDOR_UNKNOWN;
}

// Brute-force scan a vendor-IE payload for a message whose header's low
// nibble matches want_type, starting after the 3-byte OUI + type byte.
// Mirrors PDL Witness's defensive scan (different vendors position the ODID
// payload at slightly different offsets within the vendor-specific IE,
// rather than a single fixed offset).
#define WIFI_ODID_IE_SCAN_START 3
#define ODID_MESSAGE_BYTES      25

static const uint8_t* find_message_of_type(const uint8_t* ie, uint16_t ie_len,
                                            uint8_t want_type) {
    if (!ie) return nullptr;
    for (int off = WIFI_ODID_IE_SCAN_START;
         off + ODID_MESSAGE_BYTES <= (int)ie_len; off++) {
        uint8_t hdr = ie[off];
        if ((hdr & 0x0F) == want_type) return &ie[off];
    }
    return nullptr;
}

WifiOdidParseResult wifi_odid_parse_beacon(const uint8_t* frame, size_t frame_len) {
    WifiOdidParseResult r{};
    r.hit = false;
    r.vendor = ODID_VENDOR_UNKNOWN;
    r.has_serial = false;
    r.serial[0] = '\0';
    r.location.valid = false;

    uint16_t ie_len = 0;
    uint8_t vendor = ODID_VENDOR_UNKNOWN;
    const uint8_t* ie = find_vendor_ie(frame, frame_len, &ie_len, &vendor);
    if (!ie) return r;

    r.hit = true;
    r.vendor = vendor;

    const uint8_t* basic_id_msg = find_message_of_type(ie, ie_len, 0x0);
    if (basic_id_msg) {
        r.has_serial = odid_extract_basic_id_serial(basic_id_msg, ODID_MESSAGE_BYTES, r.serial);
    }

    const uint8_t* location_msg = find_message_of_type(ie, ie_len, 0x1);
    if (location_msg) {
        odid_decode_location(location_msg, ODID_MESSAGE_BYTES, &r.location);
    }

    return r;
}

// ============================================================================
// SECTION 2 — ESP32/Arduino integration (real hardware only)
// ============================================================================
#ifdef ARDUINO

#include <Arduino.h>
#include <esp_wifi.h>
#include "config.h"
#include "blue_force.h"
#include "ffps.h"
#include "enrollment.h"
#include "sensor_location.h"
#include "health_report.h"

static WifiOdidScannerStats s_stats{};

// Shared by the real packet handler AND wifi_odid_scanner_test_inject() —
// so an on-device test that injects a synthetic frame exercises the exact
// same detection-building/blue-force/FFPS logic a real beacon would.
static void process_hit(const WifiOdidParseResult& r, int16_t rssi) {
    s_stats.rid_hits++;
    if (r.location.valid) s_stats.location_hits++;

    DetectionEvent evt;
    memset(&evt, 0, sizeof(evt));
    evt.msg_type      = MSG_TYPE_CANDIDATE;
    evt.sensor_id     = enrollment_get_sensor_id();
    evt.timestamp_utc = (uint32_t)(millis() / 1000);
    evt.class_label   = CLASS_RID_WIFI_DECODED;
    evt.confidence    = 255;  // decoded broadcast content, not a probabilistic guess
    evt.rssi          = rssi;
    evt.noise_floor   = 0;
    evt.battery_pct   = health_read_battery_pct();
    evt.flags         = 0;

    if (r.location.valid) {
        evt.lat_fixed = (int32_t)(r.location.lat_deg * 1e7);
        evt.lon_fixed = (int32_t)(r.location.lon_deg * 1e7);
    } else {
        location_get(&evt.lat_fixed, &evt.lon_fixed);
    }

    if (r.has_serial) {
        blue_force_compute_hash(r.serial, evt.signature_hash);
        strncpy(evt.drone_serial, r.serial, sizeof(evt.drone_serial) - 1);
        evt.drone_serial[sizeof(evt.drone_serial) - 1] = '\0';
    }

    if (r.has_serial && blue_force_check_hash(evt.signature_hash)) {
        evt.flags |= FLAG_BLUE_FORCE;
        s_stats.blue_force_matches++;
    } else {
        ffps_submit_candidate(evt);
    }
}

void wifi_odid_scanner_process_packet(void* buf) {
    const wifi_promiscuous_pkt_t* pkt = (const wifi_promiscuous_pkt_t*)buf;
    if (!pkt) return;

    s_stats.frames_seen++;

    const uint8_t* frame = pkt->payload;
    size_t frame_len = pkt->rx_ctrl.sig_len;

    WifiOdidParseResult r = wifi_odid_parse_beacon(frame, frame_len);
    if (!r.hit) return;

    process_hit(r, (int16_t)pkt->rx_ctrl.rssi);
}

void wifi_odid_scanner_get_stats(WifiOdidScannerStats* out) {
    if (out) *out = s_stats;
}

// Test-only injection: parses a synthetic 802.11 beacon frame through the
// EXACT same wifi_odid_parse_beacon() + process_hit() path a real captured
// beacon would take.
bool wifi_odid_scanner_test_inject(const uint8_t* frame, size_t len, int16_t rssi) {
    s_stats.frames_seen++;
    WifiOdidParseResult r = wifi_odid_parse_beacon(frame, len);
    if (!r.hit) return false;
    process_hit(r, rssi);
    return true;
}

#endif // ARDUINO
