#include <WiFi.h>
#include <WiFiUdp.h>
#include "esp_efuse.h"
#include "esp_mac.h"
#include "config.h"

#if PROBE_SNIFF
#include "esp_wifi.h"
#include "freertos/queue.h"
#endif

WiFiUDP udp;
String  nodeId;

unsigned long lastScan      = 0;
unsigned long lastHeartbeat = 0;

static uint8_t homeChannel = 1;  // channel of sandbox AP — updated after connect and scan

// ─── Probe sniffing globals ───────────────────────────────────────────────────

#if PROBE_SNIFF

struct ProbeEvent {
    uint8_t src_mac[6];
    char    ssid[33];
    int8_t  rssi;
};

struct DeauthEvent {
    uint8_t  src[6];
    uint8_t  dst[6];
    uint8_t  bssid[6];
    uint8_t  subtype;   // 0xC0 = deauth, 0xA0 = disassoc
    uint16_t reason;
    int8_t   rssi;
};

struct AssocEvent {
    uint8_t src[6];    // client MAC
    uint8_t bssid[6];  // AP BSSID
    char    ssid[33];  // SSID the client is associating to
    uint8_t subtype;   // 0x00 = assoc req, 0x20 = reassoc req
    int8_t  rssi;
};

static QueueHandle_t probeQueue;
static QueueHandle_t deauthQueue;
static QueueHandle_t assocQueue;

static volatile uint32_t probeDrops  = 0;
static volatile uint32_t deauthDrops = 0;
static volatile uint32_t assocDrops  = 0;

// Set in promiscuous callback when a close-range attack-signature deauth is seen.
// Checked in loop() to trigger an immediate flush rather than waiting for next hop cycle.
static volatile bool alertPending = false;

// Dedup cache — suppress repeated MAC+SSID pairs within PROBE_DEDUP_SECS
struct DedupEntry { uint8_t mac[6]; char ssid[33]; uint32_t last_ms; };
static DedupEntry dedupCache[DEDUP_CACHE_SIZE];
static int        dedupNext = 0;

static bool isDuplicate(const uint8_t* mac, const char* ssid) {
    uint32_t now = millis();
    for (int i = 0; i < DEDUP_CACHE_SIZE; i++) {
        if (dedupCache[i].last_ms == 0) continue;
        if (memcmp(dedupCache[i].mac, mac, 6) == 0 &&
            strcmp(dedupCache[i].ssid, ssid)  == 0) {
            if ((now - dedupCache[i].last_ms) < (uint32_t)(PROBE_DEDUP_SECS * 1000))
                return true;
            dedupCache[i].last_ms = now;
            return false;
        }
    }
    // Not found — write to next slot (ring)
    memcpy(dedupCache[dedupNext].mac, mac, 6);
    strncpy(dedupCache[dedupNext].ssid, ssid, 32);
    dedupCache[dedupNext].ssid[32] = '\0';
    dedupCache[dedupNext].last_ms  = now;
    dedupNext = (dedupNext + 1) % DEDUP_CACHE_SIZE;
    return false;
}

// ─── Channel hop state ───────────────────────────────────────────────────────

static const uint8_t HOP_LIST[] = {1,2,3,4,5,6,7,8,9,10,11,12,13};
static const int     HOP_COUNT  = sizeof(HOP_LIST) / sizeof(HOP_LIST[0]);
static int           hopIdx     = 0;
static unsigned long lastHop    = 0;

// Promiscuous callback — runs in WiFi task context, not safe to call UDP here.
// Parse probe request, deauth, and disassoc frames and push to queues.
static void promiscuous_rx_cb(void* buf, wifi_promiscuous_pkt_type_t type) {
    if (type != WIFI_PKT_MGMT) return;

    const wifi_promiscuous_pkt_t* pkt = (const wifi_promiscuous_pkt_t*)buf;
    const uint8_t* d   = pkt->payload;
    uint16_t       len = pkt->rx_ctrl.sig_len;

    if (len < 24) return;
    uint8_t subtype = d[0];

    // ── Probe request (0x40) ─────────────────────────────────────────────────
    if (subtype == 0x40) {
        if (len < 28) return;
        const uint8_t* src = d + 10;

        char ssid[33] = "";
        if (d[24] == 0x00) {
            uint8_t slen = d[25];
            if (slen > 0 && slen <= 32 && (26 + slen) <= len) {
                bool printable = true;
                for (uint8_t i = 0; i < slen; i++) {
                    if (d[26 + i] < 32 || d[26 + i] > 126) { printable = false; break; }
                }
                if (printable) { memcpy(ssid, d + 26, slen); ssid[slen] = '\0'; }
            }
        }

        if (isDuplicate(src, ssid)) return;

        ProbeEvent ev;
        memcpy(ev.src_mac, src, 6);
        memcpy(ev.ssid,    ssid, 33);
        ev.rssi = (int8_t)pkt->rx_ctrl.rssi;
        if (xQueueSend(probeQueue, &ev, 0) != pdTRUE) probeDrops++;
        return;
    }

    // ── Deauth (0xC0) and Disassoc (0xA0) ───────────────────────────────────
    if (subtype == 0xC0 || subtype == 0xA0) {
        // 802.11 frame header: addr1 (dst) @ 4, addr2 (src) @ 10, addr3 (bssid) @ 16
        // Frame body starts at byte 24: 2-byte reason code
        if (len < 26) return;

        DeauthEvent ev;
        memcpy(ev.dst,   d + 4,  6);
        memcpy(ev.src,   d + 10, 6);
        memcpy(ev.bssid, d + 16, 6);
        ev.subtype = subtype;
        ev.reason  = (uint16_t)d[24] | ((uint16_t)d[25] << 8);
        ev.rssi    = (int8_t)pkt->rx_ctrl.rssi;
        if (xQueueSend(deauthQueue, &ev, 0) != pdTRUE) deauthDrops++;

        // Strong close-range attack signature — request an immediate flush
        if (ev.reason == 2 && ev.rssi >= RSSI_ALERT_THRESHOLD) {
            alertPending = true;
        }
        return;
    }

    // ── Association request (0x00) and Reassociation request (0x20) ─────────
    // Association req body:  capability(2) + listen_interval(2) + IEs
    // Reassociation req body: capability(2) + listen_interval(2) + current_AP(6) + IEs
    // SSID IE is always first IE: tag(1) + length(1) + ssid(n)
    if (subtype == 0x00 || subtype == 0x20) {
        uint16_t ie_offset = (subtype == 0x00) ? 28 : 34;  // body start + fixed fields
        if (len < ie_offset + 2) return;

        AssocEvent ev;
        memcpy(ev.src,   d + 10, 6);
        memcpy(ev.bssid, d + 16, 6);
        ev.subtype = subtype;
        ev.rssi    = (int8_t)pkt->rx_ctrl.rssi;
        ev.ssid[0] = '\0';

        if (d[ie_offset] == 0x00) {  // SSID IE tag
            uint8_t slen = d[ie_offset + 1];
            if (slen > 0 && slen <= 32 && (ie_offset + 2 + slen) <= len) {
                bool printable = true;
                for (uint8_t i = 0; i < slen; i++) {
                    if (d[ie_offset + 2 + i] < 32 || d[ie_offset + 2 + i] > 126) {
                        printable = false; break;
                    }
                }
                if (printable) {
                    memcpy(ev.ssid, d + ie_offset + 2, slen);
                    ev.ssid[slen] = '\0';
                }
            }
        }

        if (xQueueSend(assocQueue, &ev, 0) != pdTRUE) assocDrops++;
        return;
    }
}

#endif  // PROBE_SNIFF

// ─── Setup ───────────────────────────────────────────────────────────────────

void setup() {
    Serial.begin(115200);
    delay(500);

    // Derive node ID from MAC (last 4 bytes, no colons)
    uint8_t mac[6];
    esp_efuse_mac_get_default(mac);
    char macBuf[9];
    snprintf(macBuf, sizeof(macBuf), "%02X%02X%02X%02X", mac[2], mac[3], mac[4], mac[5]);
    nodeId = String(macBuf);

    Serial.printf("\n[NODE] ID: %s\n", nodeId.c_str());
    connectWiFi();

#if PROBE_SNIFF
    memset(dedupCache, 0, sizeof(dedupCache));
    probeQueue  = xQueueCreate(PROBE_QUEUE_SIZE,  sizeof(ProbeEvent));
    deauthQueue = xQueueCreate(DEAUTH_QUEUE_SIZE, sizeof(DeauthEvent));
    assocQueue  = xQueueCreate(ASSOC_QUEUE_SIZE,  sizeof(AssocEvent));
    esp_wifi_set_promiscuous_rx_cb(promiscuous_rx_cb);
    esp_wifi_set_promiscuous(true);
    Serial.println("[PROBE] Promiscuous mode enabled");
#endif
}

// ─── Loop ────────────────────────────────────────────────────────────────────

void loop() {
    if (WiFi.status() != WL_CONNECTED) {
        Serial.println("[WIFI] Lost connection, reconnecting...");
#if PROBE_SNIFF
        esp_wifi_set_promiscuous(false);
#endif
        connectWiFi();
#if PROBE_SNIFF
        esp_wifi_set_channel(homeChannel, WIFI_SECOND_CHAN_NONE);
        esp_wifi_set_promiscuous(true);
        hopIdx  = 0;
        lastHop = millis();
#endif
        return;
    }

    unsigned long now = millis();

    if (now - lastHeartbeat >= HEARTBEAT_INTERVAL_MS) {
        lastHeartbeat = now;
        sendHeartbeat();
    }

    if (now - lastScan >= SCAN_INTERVAL_MS) {
        lastScan = now;
        scanAndSend();
        return;
    }

#if PROBE_SNIFF
    // Alert path — close-range attack deauth detected; flush immediately
    // rather than waiting up to 4s for the next natural homeChannel pass.
    if (alertPending && WiFi.status() == WL_CONNECTED) {
        alertPending = false;
        esp_wifi_set_channel(homeChannel, WIFI_SECOND_CHAN_NONE);
        flushDeauthQueue();
        flushProbeQueue();
        flushAssocQueue();
        lastHop = millis();  // avoid an immediate hop right after flushing
    } else {
        advanceHop();
    }
#endif

    delay(10);
}

// ─── WiFi ────────────────────────────────────────────────────────────────────

void connectWiFi() {
    Serial.printf("[WIFI] Connecting to %s\n", WIFI_SSID);
    WiFi.mode(WIFI_STA);
    WiFi.begin(WIFI_SSID, WIFI_PASSWORD);

    int attempts = 0;
    while (WiFi.status() != WL_CONNECTED && attempts < 20) {
        delay(500);
        Serial.print(".");
        attempts++;
    }

    if (WiFi.status() == WL_CONNECTED) {
        Serial.printf("\n[WIFI] Connected — IP: %s\n", WiFi.localIP().toString().c_str());
#if PROBE_SNIFF
        uint8_t primary; wifi_second_chan_t second;
        esp_wifi_get_channel(&primary, &second);
        if (primary > 0) homeChannel = primary;
        Serial.printf("[HOP]  Home channel: %d\n", homeChannel);
#endif
    } else {
        Serial.println("\n[WIFI] Failed to connect, will retry in loop");
    }
}

// ─── Scan ────────────────────────────────────────────────────────────────────

void scanAndSend() {
#if PROBE_SNIFF
    // Stop hopping and return to home channel — scan requires a stable channel to start.
    esp_wifi_set_promiscuous(false);
    esp_wifi_set_channel(homeChannel, WIFI_SECOND_CHAN_NONE);
    delay(50);
#endif

    Serial.println("[SCAN] Starting...");
    int n = WiFi.scanNetworks(false, true); // async=false, show_hidden=true

    if (n == WIFI_SCAN_FAILED) {
        Serial.println("[SCAN] Failed");
#if PROBE_SNIFF
        esp_wifi_set_promiscuous(true);
#endif
        return;
    }

    Serial.printf("[SCAN] Found %d networks\n", n);

    for (int i = 0; i < n; i++) {
        sendBeaconEvent(i);
    }

    WiFi.scanDelete();

#if PROBE_SNIFF
    // Re-read home channel — scan resets the radio, AP channel may have shifted.
    uint8_t primary; wifi_second_chan_t second;
    esp_wifi_get_channel(&primary, &second);
    if (primary > 0) homeChannel = primary;

    // Re-enable promiscuous and flush anything queued before the scan.
    esp_wifi_set_promiscuous(true);
    flushProbeQueue();
    flushDeauthQueue();
    flushAssocQueue();

    // Reset hop state so next cycle starts cleanly from channel 1.
    hopIdx  = 0;
    lastHop = millis();
#endif
}

// ─── Event helpers ───────────────────────────────────────────────────────────

const char* encStr(wifi_auth_mode_t enc) {
    switch (enc) {
        case WIFI_AUTH_OPEN:          return "OPEN";
        case WIFI_AUTH_WEP:           return "WEP";
        case WIFI_AUTH_WPA_PSK:       return "WPA";
        case WIFI_AUTH_WPA2_PSK:      return "WPA2";
        case WIFI_AUTH_WPA_WPA2_PSK:  return "WPA/WPA2";
        case WIFI_AUTH_WPA3_PSK:      return "WPA3";
        case WIFI_AUTH_WPA2_WPA3_PSK: return "WPA2/WPA3";
        default:                       return "UNKNOWN";
    }
}

const char* importance(int rssi) {
    if (rssi >= -50) return "high";
    if (rssi <= -80) return "low";
    return "normal";
}

// ─── Send ────────────────────────────────────────────────────────────────────

void sendBeaconEvent(int idx) {
    String ssid = WiFi.SSID(idx);
    ssid.replace("\\", "\\\\");   // escape backslashes first, then quotes
    ssid.replace("\"", "\\\"");

    String bssid  = WiFi.BSSIDstr(idx);
    int    rssi   = WiFi.RSSI(idx);
    int    ch     = WiFi.channel(idx);
    const char* enc = encStr(WiFi.encryptionType(idx));
    const char* imp = importance(rssi);
    unsigned long ts = millis();

    char buf[512];
    snprintf(buf, sizeof(buf),
        "{\"node_id\":\"%s\",\"ts\":%lu,\"type\":\"beacon\","
        "\"conf\":\"high\",\"imp\":\"%s\","
        "\"ssid\":\"%s\",\"bssid\":\"%s\","
        "\"rssi\":%d,\"ch\":%d,\"enc\":\"%s\"}",
        nodeId.c_str(), ts, imp,
        ssid.c_str(), bssid.c_str(),
        rssi, ch, enc
    );

    udp.beginPacket(COORDINATOR_IP, COORDINATOR_PORT);
    udp.print(buf);
    udp.endPacket();

    Serial.printf("  → %-32s %s  ch%-3d %4ddBm  %s\n",
        ssid.c_str(), bssid.c_str(), ch, rssi, enc);
}

// ─── Heartbeat ───────────────────────────────────────────────────────────────

void sendHeartbeat() {
    unsigned long now = millis();
    char buf[320];
    snprintf(buf, sizeof(buf),
        "{\"node_id\":\"%s\",\"ts\":%lu,\"type\":\"heartbeat\","
        "\"uptime_ms\":%lu,\"free_heap\":%lu,\"wifi_rssi\":%d,"
        "\"probe_drops\":%lu,\"deauth_drops\":%lu,\"assoc_drops\":%lu}",
        nodeId.c_str(), now,
        now,
        (unsigned long)ESP.getFreeHeap(),
        WiFi.RSSI(),
        (unsigned long)probeDrops,
        (unsigned long)deauthDrops,
        (unsigned long)assocDrops
    );
    udp.beginPacket(COORDINATOR_IP, COORDINATOR_PORT);
    udp.print(buf);
    udp.endPacket();
    Serial.printf("[HB] uptime=%lus  heap=%luK  ap=%ddBm  drops(p/d/a)=%lu/%lu/%lu\n",
        now / 1000,
        (unsigned long)ESP.getFreeHeap() / 1024,
        WiFi.RSSI(),
        (unsigned long)probeDrops,
        (unsigned long)deauthDrops,
        (unsigned long)assocDrops
    );
}

// ─── Probe send / flush ───────────────────────────────────────────────────────

#if PROBE_SNIFF

void sendProbeEvent(const ProbeEvent& ev) {
    char mac_str[18];
    snprintf(mac_str, sizeof(mac_str), "%02X:%02X:%02X:%02X:%02X:%02X",
             ev.src_mac[0], ev.src_mac[1], ev.src_mac[2],
             ev.src_mac[3], ev.src_mac[4], ev.src_mac[5]);

    // Escape any quotes in SSID
    char ssidEsc[66];
    int j = 0;
    for (int i = 0; ev.ssid[i] && j < 64; i++) {
        if (ev.ssid[i] == '"' || ev.ssid[i] == '\\') ssidEsc[j++] = '\\';
        ssidEsc[j++] = ev.ssid[i];
    }
    ssidEsc[j] = '\0';

    char buf[256];
    snprintf(buf, sizeof(buf),
        "{\"node_id\":\"%s\",\"ts\":%lu,\"type\":\"probe\","
        "\"conf\":\"high\",\"imp\":\"%s\","
        "\"src_mac\":\"%s\",\"ssid\":\"%s\",\"rssi\":%d}",
        nodeId.c_str(), millis(), importance(ev.rssi),
        mac_str, ssidEsc, (int)ev.rssi
    );

    udp.beginPacket(COORDINATOR_IP, COORDINATOR_PORT);
    udp.print(buf);
    udp.endPacket();

    Serial.printf("  [PROBE] %s → \"%s\"  %ddBm\n",
        mac_str, ev.ssid[0] ? ev.ssid : "<wildcard>", (int)ev.rssi);
}

void flushProbeQueue() {
    if (WiFi.status() != WL_CONNECTED) return;
    ProbeEvent ev;
    while (xQueueReceive(probeQueue, &ev, 0) == pdTRUE) {
        sendProbeEvent(ev);
    }
}

void sendDeauthEvent(const DeauthEvent& ev) {
    auto macStr = [](const uint8_t* m, char* out) {
        snprintf(out, 18, "%02X:%02X:%02X:%02X:%02X:%02X",
                 m[0], m[1], m[2], m[3], m[4], m[5]);
    };
    char src[18], dst[18], bssid[18];
    macStr(ev.src,   src);
    macStr(ev.dst,   dst);
    macStr(ev.bssid, bssid);

    const char* stype = (ev.subtype == 0xC0) ? "deauth" : "disassoc";

    char buf[384];
    snprintf(buf, sizeof(buf),
        "{\"node_id\":\"%s\",\"ts\":%lu,\"type\":\"%s\","
        "\"src\":\"%s\",\"dst\":\"%s\",\"bssid\":\"%s\","
        "\"reason\":%u,\"rssi\":%d}",
        nodeId.c_str(), millis(), stype,
        src, dst, bssid,
        (unsigned)ev.reason, (int)ev.rssi
    );

    udp.beginPacket(COORDINATOR_IP, COORDINATOR_PORT);
    udp.print(buf);
    udp.endPacket();

    Serial.printf("  [%s] %s → %s  bssid=%s  reason=%u  %ddBm\n",
        stype, src, dst, bssid, (unsigned)ev.reason, (int)ev.rssi);
}

void flushDeauthQueue() {
    if (WiFi.status() != WL_CONNECTED) return;
    DeauthEvent ev;
    while (xQueueReceive(deauthQueue, &ev, 0) == pdTRUE) {
        sendDeauthEvent(ev);
    }
}

void sendAssocEvent(const AssocEvent& ev) {
    auto macStr = [](const uint8_t* m, char* out) {
        snprintf(out, 18, "%02X:%02X:%02X:%02X:%02X:%02X",
                 m[0], m[1], m[2], m[3], m[4], m[5]);
    };
    char src[18], bssid[18];
    macStr(ev.src,   src);
    macStr(ev.bssid, bssid);

    const char* stype = (ev.subtype == 0x00) ? "assoc" : "reassoc";

    // Escape quotes in SSID
    char ssidEsc[66];
    int j = 0;
    for (int i = 0; ev.ssid[i] && j < 64; i++) {
        if (ev.ssid[i] == '"' || ev.ssid[i] == '\\') ssidEsc[j++] = '\\';
        ssidEsc[j++] = ev.ssid[i];
    }
    ssidEsc[j] = '\0';

    char buf[320];
    snprintf(buf, sizeof(buf),
        "{\"node_id\":\"%s\",\"ts\":%lu,\"type\":\"%s\","
        "\"src\":\"%s\",\"bssid\":\"%s\",\"ssid\":\"%s\",\"rssi\":%d}",
        nodeId.c_str(), millis(), stype,
        src, bssid, ssidEsc, (int)ev.rssi
    );

    udp.beginPacket(COORDINATOR_IP, COORDINATOR_PORT);
    udp.print(buf);
    udp.endPacket();

    Serial.printf("  [%s] %s → %s  \"%s\"  %ddBm\n",
        stype, src, bssid, ev.ssid[0] ? ev.ssid : "<hidden>", (int)ev.rssi);
}

void flushAssocQueue() {
    if (WiFi.status() != WL_CONNECTED) return;
    AssocEvent ev;
    while (xQueueReceive(assocQueue, &ev, 0) == pdTRUE) {
        sendAssocEvent(ev);
    }
}

// Advance to the next channel in the hop list.
// Non-blocking — returns immediately if the dwell time hasn't elapsed.
// Flushes queues each time we land back on homeChannel (WiFi is usable there).
void advanceHop() {
    unsigned long now = millis();
    if (now - lastHop < HOP_DWELL_MS) return;
    lastHop = now;

    hopIdx = (hopIdx + 1) % HOP_COUNT;
    uint8_t ch = HOP_LIST[hopIdx];
    esp_wifi_set_channel(ch, WIFI_SECOND_CHAN_NONE);

    if (ch == homeChannel && WiFi.status() == WL_CONNECTED) {
        flushProbeQueue();
        flushDeauthQueue();
        flushAssocQueue();
    }
}

#endif  // PROBE_SNIFF