// ============================================================================
// handset_fingerprint.cpp — see handset_fingerprint.h for design notes.
// ============================================================================

#include "handset_fingerprint.h"
#include "oui_table.h"

#include <cstring>
#include <cctype>

#ifdef ARDUINO
#include <Arduino.h>
static inline uint32_t hfp_now_ms() { return millis(); }
#else
#include <chrono>
static inline uint32_t hfp_now_ms() {
    using namespace std::chrono;
    static auto start = steady_clock::now();
    return (uint32_t)duration_cast<milliseconds>(steady_clock::now() - start).count();
}
#endif

// ---------------------------------------------------------------------------
// General consumer-vendor OUI table -- the REAL IEEE MA-L registry
// (oui_table.h, 39,809 entries, downloaded from standards-oui.ieee.org
// 2026-07-27), not a small hand-curated example list. Sorted ascending by
// OUI for binary search, since a linear scan through 40k entries per
// received frame would be too slow for the promiscuous RX callback's IRQ
// context. This replaces the earlier ~16-entry curated table entirely --
// every OUI that table had is already covered by the real registry, so
// this is a strict coverage improvement, not a removed capability.
// ---------------------------------------------------------------------------
static const char* lookup_vendor_oui_real(const uint8_t oui[3]) {
    uint32_t target = ((uint32_t)oui[0] << 16) | ((uint32_t)oui[1] << 8) | oui[2];
    int lo = 0, hi = OUI_TABLE_SIZE - 1;
    while (lo <= hi) {
        int mid = lo + (hi - lo) / 2;
        uint32_t v = kOuiTable[mid].oui;
        if (v == target) return kOuiTable[mid].vendor;
        if (v < target) lo = mid + 1;
        else hi = mid - 1;
    }
    return nullptr;
}

// ---------------------------------------------------------------------------
// Drone-controller-app signature table — ported directly from PDL Witness's
// probe_fingerprint.c (same OUI sourcing: DJI, Autel, Skydio, Parrot,
// Walksnail, HDZero). Matches the controller-app PHONE, not the drone
// itself (DroneID/wifi_odid_scanner handle that side).
//
// Iridium GO! deliberately has NO signature here — its actual WiFi-hotspot
// OUI/SSID pattern has not been verified against a real capture, and this
// module does not ship unverified vendor signatures. HANDSET_APP_IRIDIUM_GO
// stays defined for when that data exists.
// ---------------------------------------------------------------------------
struct ControllerSig {
    uint8_t oui[2][3];
    uint8_t oui_count;
    const char* ssid_pattern;   // nullptr = no SSID requirement
    HandsetAppClass app_class;
    const char* label;
};

static const ControllerSig kControllerSigs[] = {
    {{{0x60,0x60,0x1F},{0x34,0xD2,0x62}}, 2, "DJI",    HANDSET_APP_DJI_PILOT2,        "DJI Pilot 2"},
    {{{0x9C,0xDD,0x1F},{0x60,0x60,0x1F}}, 2, "Mavic",  HANDSET_APP_DJI_FLY,           "DJI Fly"},
    {{{0x30,0xDE,0x4B},{0x18,0xB6,0x5F}}, 2, "Autel",  HANDSET_APP_AUTEL_GO,          "AutelGo"},
    {{{0xA4,0xB1,0x97},{0x2C,0x6D,0xC1}}, 2, nullptr,  HANDSET_APP_SKYDIO,            "Skydio"},
    {{{0x00,0x12,0x1C},{0x90,0x03,0xB7}}, 2, "Anafi",  HANDSET_APP_PARROT_FREEFLIGHT, "Parrot FreeFlight"},
    {{{0x3C,0x71,0xBF},{0xF0,0x08,0xD1}}, 2, nullptr,  HANDSET_APP_WALKSNAIL,         "Walksnail"},
    {{{0xDC,0x54,0x75},{0,0,0}},          1, "HDZero", HANDSET_APP_HDZERO,            "HDZero VRX"},
};
#define N_CONTROLLER_SIGS (sizeof(kControllerSigs) / sizeof(kControllerSigs[0]))

// ---------------------------------------------------------------------------
// Known CCTV/surveillance-camera manufacturer table -- whitehat "unexpected
// camera nearby" lead. Deliberately restricted to vendors whose product
// line is DEDICATED camera/DVR/NVR hardware (not general consumer-
// electronics brands that happen to also sell a camera, e.g. TP-Link,
// D-Link, Logitech are excluded -- those would false-positive on routers/
// mice/webcams that aren't fixed surveillance installs). Substrings are
// matched case-insensitively against the exact vendor string text as it
// appears in the real IEEE OUI registry (oui_table.h), so these must match
// IEEE's registered organization name, not a marketing brand name -- e.g.
// Hikvision registers as "Hangzhou Hikvision Digital Technology", so the
// substring is "Hikvision", not the full legal name.
// ---------------------------------------------------------------------------
static const char* kCameraVendorSubstrings[] = {
    "Hikvision",
    "Dahua",
    "Reolink",
    "Amcrest",
    "Foscam",
    "Wyze",
    "Anjielosmart",
    "Axis Communications",
    "Vivotek",
    "Swann",
    "Lorex",
    "Arlo",
    "Ubiquiti Networks",   // UniFi Protect camera line specifically
    "Uniview",
    "Zhejiang Uniview",
};
#define N_CAMERA_VENDOR_SUBSTRINGS (sizeof(kCameraVendorSubstrings) / sizeof(kCameraVendorSubstrings[0]))

static bool ci_substr_cstr(const char* hay, const char* needle) {
    if (!hay || !needle || !*needle) return false;
    size_t hlen = strlen(hay), nlen = strlen(needle);
    if (nlen > hlen) return false;
    for (size_t i = 0; i + nlen <= hlen; i++) {
        size_t k = 0;
        for (; k < nlen; k++) {
            if (tolower((unsigned char)hay[i + k]) != tolower((unsigned char)needle[k])) break;
        }
        if (k == nlen) return true;
    }
    return false;
}

static bool is_camera_vendor(const char* vendor) {
    if (!vendor || !*vendor) return false;
    for (size_t i = 0; i < N_CAMERA_VENDOR_SUBSTRINGS; i++) {
        if (ci_substr_cstr(vendor, kCameraVendorSubstrings[i])) return true;
    }
    return false;
}

// BLE manufacturer company IDs (Bluetooth SIG assigned, little-endian on air).
struct BleCompanyId { uint16_t id; const char* vendor; };
static const BleCompanyId kBleCompanyIds[] = {
    {0x004C, "Apple"},
    {0x0075, "Samsung"},
    {0x00E0, "Google"},
    {0x0006, "Microsoft"},
    {0x038F, "Xiaomi"},
};
#define N_BLE_COMPANY_IDS (sizeof(kBleCompanyIds) / sizeof(kBleCompanyIds[0]))

// ---------------------------------------------------------------------------
// Nearby-device presence ring (generalizes probe_fingerprint.c's pilot ring)
// ---------------------------------------------------------------------------
#define PRESENCE_RING_SLOTS 48
#define PRESENCE_TTL_MS      (60u * 1000u)

struct PresenceSlot {
    uint8_t  mac[6];
    uint32_t last_seen_ms;
    bool     active;
    bool     vendor_matched;
    bool     controller_matched;
    bool     camera_matched;
    char     vendor[16];
    char     app_label[24];
    int8_t   last_rssi;
    bool     via_ble;
    bool     is_public_addr;
};
static PresenceSlot s_ring[PRESENCE_RING_SLOTS];
static HandsetStats s_stats{};

static void note_device(const uint8_t mac[6], const char* vendor, const char* app_label,
                         int8_t rssi, bool via_ble, bool is_public_addr) {
    bool vendor_matched = (vendor != nullptr);
    bool controller_matched = (app_label != nullptr);
    bool camera_matched = is_camera_vendor(vendor);
    uint32_t now = hfp_now_ms();
    int slot = -1, oldest = -1;
    uint32_t oldest_ms = 0xFFFFFFFFu;
    for (int i = 0; i < PRESENCE_RING_SLOTS; i++) {
        if (s_ring[i].active && memcmp(s_ring[i].mac, mac, 6) == 0) { slot = i; break; }
        if (!s_ring[i].active && slot < 0) slot = i;
        if (s_ring[i].active && s_ring[i].last_seen_ms < oldest_ms) {
            oldest = i; oldest_ms = s_ring[i].last_seen_ms;
        }
    }
    if (slot < 0) slot = oldest;
    if (slot < 0) return;
    bool is_new = !s_ring[slot].active || memcmp(s_ring[slot].mac, mac, 6) != 0;
    if (is_new) {
        memset(&s_ring[slot], 0, sizeof(PresenceSlot));
        memcpy(s_ring[slot].mac, mac, 6);
    }
    s_ring[slot].last_seen_ms = now;
    s_ring[slot].active = true;
    s_ring[slot].last_rssi = rssi;
    s_ring[slot].via_ble = via_ble;
    s_ring[slot].is_public_addr = is_public_addr;
    if (vendor_matched && (is_new || !s_ring[slot].vendor_matched)) s_stats.vendor_matches_total++;
    if (controller_matched && (is_new || !s_ring[slot].controller_matched)) s_stats.controller_app_matches_total++;
    if (camera_matched && (is_new || !s_ring[slot].camera_matched)) s_stats.camera_vendor_matches_total++;
    s_ring[slot].vendor_matched = s_ring[slot].vendor_matched || vendor_matched;
    s_ring[slot].controller_matched = s_ring[slot].controller_matched || controller_matched;
    s_ring[slot].camera_matched = s_ring[slot].camera_matched || camera_matched;
    if (vendor) { strncpy(s_ring[slot].vendor, vendor, sizeof(s_ring[slot].vendor) - 1); }
    if (app_label) { strncpy(s_ring[slot].app_label, app_label, sizeof(s_ring[slot].app_label) - 1); }
}

static uint16_t count_active_recent() {
    uint32_t now = hfp_now_ms();
    uint16_t n = 0;
    for (int i = 0; i < PRESENCE_RING_SLOTS; i++) {
        if (s_ring[i].active && (now - s_ring[i].last_seen_ms) < PRESENCE_TTL_MS) n++;
    }
    return n;
}

static const char* lookup_vendor_oui(const uint8_t oui[3]) {
    return lookup_vendor_oui_real(oui);
}

static bool ci_substr(const uint8_t* hay, uint8_t hay_len, const char* needle) {
    if (!needle || !*needle) return true;
    size_t nlen = strlen(needle);
    if (nlen > hay_len) return false;
    for (uint8_t i = 0; i + nlen <= hay_len; i++) {
        size_t k = 0;
        for (; k < nlen; k++) {
            char a = (char)hay[i + k];
            char b = needle[k];
            if (tolower((unsigned char)a) != tolower((unsigned char)b)) break;
        }
        if (k == nlen) return true;
    }
    return false;
}

void handset_fingerprint_init() {
    memset(s_ring, 0, sizeof(s_ring));
    s_stats = HandsetStats{};
}

// ---------------------------------------------------------------------------
// WiFi probe-request parsing (802.11 management frame, subtype 0x04).
// Mgmt header is 24 bytes; IE chain starts there. Source MAC = Address 2,
// header bytes 10..15.
// ---------------------------------------------------------------------------
#define MGMT_HDR_LEN 24

void handset_fingerprint_process_wifi_probe(const uint8_t* frame, size_t len, int8_t rssi) {
    if (!frame || len < MGMT_HDR_LEN + 2) return;
    s_stats.wifi_probes_seen++;

    const uint8_t* p = frame + MGMT_HDR_LEN;
    const uint8_t* end = frame + len;

    const uint8_t* ssid = nullptr;
    uint8_t ssid_len = 0;
    uint8_t vie_ouis[16][3];
    uint8_t vie_n = 0;

    while (p + 2 <= end) {
        uint8_t tag = p[0];
        uint8_t l = p[1];
        if (p + 2 + l > end) break;
        const uint8_t* body = p + 2;
        if (tag == 0) {
            ssid = body; ssid_len = l;
        } else if (tag == 221 && l >= 3) {
            if (vie_n < 16) { memcpy(vie_ouis[vie_n], body, 3); vie_n++; }
        }
        p += 2 + l;
    }

    const uint8_t* src_mac = frame + 10;
    const char* vendor = nullptr;
    const char* app_label = nullptr;

    // Vendor-IE OUIs identify the general device vendor (best-effort).
    for (uint8_t v = 0; v < vie_n && !vendor; v++) {
        vendor = lookup_vendor_oui(vie_ouis[v]);
    }
    // Source MAC's own OUI (first 3 bytes) is a second, independent vendor hint.
    if (!vendor) vendor = lookup_vendor_oui(src_mac);

    // Drone-controller-app match (specific case, same technique as ODID work).
    for (size_t s = 0; s < N_CONTROLLER_SIGS && !app_label; s++) {
        const ControllerSig& sig = kControllerSigs[s];
        bool oui_hit = false;
        for (uint8_t i = 0; i < sig.oui_count && !oui_hit; i++) {
            for (uint8_t v = 0; v < vie_n; v++) {
                if (memcmp(sig.oui[i], vie_ouis[v], 3) == 0) { oui_hit = true; break; }
            }
        }
        if (!oui_hit) continue;
        if (sig.ssid_pattern && !ci_substr(ssid, ssid_len, sig.ssid_pattern)) continue;
        app_label = sig.label;
    }

    note_device(src_mac, vendor, app_label, rssi, /*via_ble=*/false, /*is_public_addr=*/true);
    s_stats.distinct_devices_last_60s = count_active_recent();
}

// ---------------------------------------------------------------------------
// BLE advertisement parsing — manufacturer-specific data (AD type 0xFF):
// first 2 bytes (little-endian) are the Bluetooth SIG company ID.
// ---------------------------------------------------------------------------
void handset_fingerprint_process_ble_adv(const uint8_t mac[6],
                                          const uint8_t* adv_data, size_t adv_len,
                                          int8_t rssi, bool is_public_addr) {
    if (!mac) return;
    s_stats.ble_adverts_seen++;

    const char* vendor = nullptr;
    size_t i = 0;
    while (i + 1 < adv_len) {
        uint8_t field_len = adv_data[i];
        if (field_len == 0 || i + 1 + field_len > adv_len) break;
        uint8_t ad_type = adv_data[i + 1];
        if (ad_type == 0xFF && field_len >= 3) {
            uint16_t company_id = (uint16_t)adv_data[i + 2] | ((uint16_t)adv_data[i + 3] << 8);
            for (size_t c = 0; c < N_BLE_COMPANY_IDS; c++) {
                if (kBleCompanyIds[c].id == company_id) { vendor = kBleCompanyIds[c].vendor; break; }
            }
        }
        i += 1 + field_len;
    }
    // Only meaningful for a genuine public (static) address -- a random
    // address's "OUI-looking" first 3 bytes are not a real vendor prefix.
    if (!vendor && is_public_addr) vendor = lookup_vendor_oui(mac);

    note_device(mac, vendor, /*app_label=*/nullptr, rssi, /*via_ble=*/true, is_public_addr);
    s_stats.distinct_devices_last_60s = count_active_recent();
}

void handset_fingerprint_get_stats(HandsetStats* out) {
    if (out) *out = s_stats;
}

size_t handset_fingerprint_list_devices(HandsetDeviceInfo* out, size_t max) {
    if (!out || max == 0) return 0;
    uint32_t now = hfp_now_ms();
    size_t n = 0;
    for (int i = 0; i < PRESENCE_RING_SLOTS && n < max; i++) {
        if (!s_ring[i].active) continue;
        if (now - s_ring[i].last_seen_ms >= PRESENCE_TTL_MS) continue;
        memcpy(out[n].mac, s_ring[i].mac, 6);
        strncpy(out[n].vendor, s_ring[i].vendor, sizeof(out[n].vendor) - 1);
        out[n].vendor[sizeof(out[n].vendor) - 1] = '\0';
        strncpy(out[n].app_label, s_ring[i].app_label, sizeof(out[n].app_label) - 1);
        out[n].app_label[sizeof(out[n].app_label) - 1] = '\0';
        out[n].last_rssi = s_ring[i].last_rssi;
        out[n].last_seen_ms = s_ring[i].last_seen_ms;
        out[n].via_ble = s_ring[i].via_ble;
        out[n].is_public_addr = s_ring[i].is_public_addr;
        out[n].is_known_camera_vendor = s_ring[i].camera_matched;
        n++;
    }
    return n;
}

bool handset_fingerprint_test_inject_wifi(const uint8_t* frame, size_t len, int8_t rssi) {
    uint32_t before = s_stats.vendor_matches_total + s_stats.controller_app_matches_total;
    handset_fingerprint_process_wifi_probe(frame, len, rssi);
    uint32_t after = s_stats.vendor_matches_total + s_stats.controller_app_matches_total;
    return after > before;
}

bool handset_fingerprint_test_inject_ble(const uint8_t mac[6],
                                          const uint8_t* adv_data, size_t adv_len,
                                          int8_t rssi, bool is_public_addr) {
    uint32_t before = s_stats.vendor_matches_total;
    handset_fingerprint_process_ble_adv(mac, adv_data, adv_len, rssi, is_public_addr);
    uint32_t after = s_stats.vendor_matches_total;
    return after > before;
}
