#include <Wire.h>
#include <Adafruit_GFX.h>
#include <Adafruit_SSD1306.h>
#include <BLEDevice.h>
#include <BLEScan.h>
#include <BLEAdvertisedDevice.h>

#define SCREEN_WIDTH 128
#define SCREEN_HEIGHT 32 
#define OLED_RESET     -1
Adafruit_SSD1306 display(SCREEN_WIDTH, SCREEN_HEIGHT, &Wire, OLED_RESET);

int SCAN_TIME_SECONDS = 2; 
BLEScan* pBLEScan;
bool threatDetected = false;
int targetRSSI = -100;
String detectedDeviceName = "";

// Hardware Identifiers
const String META_UUID = "fd5f"; 
const uint16_t ID_META_FB      = 0x01AB; 
const uint16_t ID_META_TECH    = 0x058E; 
const uint16_t ID_LUXOTTICA    = 0x0D53; 
const uint16_t ID_SNAPCHAT     = 0x03C2; 

class MyAdvertisedDeviceCallbacks: public BLEAdvertisedDeviceCallbacks {
    void onResult(BLEAdvertisedDevice advertisedDevice) {
        String devName = advertisedDevice.getName().c_str();
        int rssi = advertisedDevice.getRSSI();

        // 1. Service UUID Filter (Meta Gen 2 Broadcast)
        if (advertisedDevice.haveServiceUUID()) {
            String serviceUUID = advertisedDevice.getServiceUUID().toString().c_str();
            if (serviceUUID.indexOf(META_UUID) != -1) {
                threatDetected = true;
                targetRSSI = rssi;
                detectedDeviceName = "Ray-Ban Meta G2";
                return;
            }
        }

        // 2. Manufacturer Data Filter (Hidden/Background Tracking)
        if (advertisedDevice.haveManufacturerData()) {
            String rawData = advertisedDevice.getManufacturerData();
            if (rawData.length() >= 2) {
                uint16_t companyId = (rawData[1] << 8) | rawData[0];
                
                if (companyId == ID_LUXOTTICA || companyId == ID_META_FB || companyId == ID_META_TECH) {
                    threatDetected = true;
                    targetRSSI = rssi;
                    detectedDeviceName = "Meta/Ray-Ban";
                    return;
                }
                
                if (companyId == ID_SNAPCHAT) {
                    threatDetected = true;
                    targetRSSI = rssi;
                    detectedDeviceName = "Snap Spectacles";
                    return;
                }
            }
        }

        // 3. Name Broadcast Backup Filter
        if (devName.indexOf("Meta") != -1 || devName.indexOf("Ray-Ban") != -1 || devName.indexOf("Spectacles") != -1) {
            threatDetected = true;
            targetRSSI = rssi;
            detectedDeviceName = devName.length() > 0 ? devName : "Smart Glass";
            return;
        }
    }
};

void showRadarUI(bool danger, String msg, int rssi) {
    display.clearDisplay();
    display.setTextSize(1);
    display.setTextColor(SSD1306_WHITE);
    
    if (danger) {
        display.setCursor(0, 0);
        display.println(F("!! THREAT DETECTED !!"));
        display.setCursor(0, 12);
        display.print(F("TRAP: "));
        display.println(msg.substring(0, 14)); 
        
        display.setCursor(0, 24);
        display.print(F("PROX: "));
        int bars = map(rssi, -90, -40, 1, 12);
        if (bars < 1) bars = 1; if (bars > 12) bars = 12;
        for(int i=0; i<bars; i++) display.print(F(">"));
    } else {
        display.setCursor(0, 4);
        display.println(F("MULTI-SHIELD ACTIVE"));
        display.setCursor(0, 20);
        display.print(F("Scanning Airspace..."));
    }
    display.display();
}

void setup() {
    Serial.begin(115200);
    Wire.begin(D4, D5);

    if(!display.begin(SSD1306_SWITCHCAPVCC, 0x3C)) { 
        for(;;);
    }
    
    display.clearDisplay();
    display.setTextSize(1);
    display.setTextColor(SSD1306_WHITE);
    display.setCursor(5, 12); 
    display.println(F("PRIVACY RADAR V3.0"));
    display.display();
    delay(2000);

    BLEDevice::init("Privacy_Shield_Pod");
    pBLEScan = BLEDevice::getScan();
    pBLEScan->setAdvertisedDeviceCallbacks(new MyAdvertisedDeviceCallbacks());
    pBLEScan->setActiveScan(true); 
    pBLEScan->setInterval(100);
    pBLEScan->setWindow(99);
}

void loop() {
    threatDetected = false;
    showRadarUI(false, "", 0);
    
    pBLEScan->start(SCAN_TIME_SECONDS, false);
    pBLEScan->clearResults();   
    
    if (threatDetected) {
        showRadarUI(true, detectedDeviceName, targetRSSI);
        delay(4000); 
    }
    delay(100); 
}