#include <Arduino.h>
#include <Wire.h>
#include <Adafruit_SSD1306.h>
#include <WiFi.h>
#include <BLEDevice.h>
#include <BLEUtils.h>
#include <BLEScan.h>
#include <esp_wifi.h>

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

// Buttons
#define BTN_RED D1
#define BTN_BLUE1 D2
#define BTN_BLUE2 D0

// Global counters
int wifiPackets = 0;
int bleDevices = 0;
int zigbeePackets = 0;
String lastAction = "Ready";

// Animation
unsigned long lastAnim = 0;
int animStep = 0;
unsigned long lastBleScan = 0;

// WiFi sniffer
void wifiSniffer(void* buf, wifi_promiscuous_pkt_type_t type) {
  wifiPackets++;
  lastAction = "WiFi pkt";
}

void setupWiFiSniffer() {
  WiFi.mode(WIFI_MODE_NULL);
  esp_wifi_set_promiscuous(true);
  esp_wifi_set_promiscuous_rx_cb(wifiSniffer);
}

// BLE scan
class MyAdvertisedDeviceCallbacks : public BLEAdvertisedDeviceCallbacks {
  void onResult(BLEAdvertisedDevice advertisedDevice) {
    bleDevices++;
    lastAction = "BLE: " + advertisedDevice.getAddress().toString().c_str();
  }
};

void setupBLESniffer() {
  BLEDevice::init("WORM");
  BLEScan* pScan = BLEDevice::getScan();
  pScan->setAdvertisedDeviceCallbacks(new MyAdvertisedDeviceCallbacks());
  pScan->setActiveScan(true);
  pScan->start(5, false);
}

void runBLEScan() {
  BLEScan* pScan = BLEDevice::getScan();
  pScan->start(5, false);
}

// Zigbee simulation
void zigbeeTask(void* pvParameters) {
  while (true) {
    zigbeePackets++;
    lastAction = "Zigbee #" + String(zigbeePackets);
    delay(random(500, 1500));
  }
}

void setupZigbeeSim() {
  xTaskCreate(zigbeeTask, "ZigbeeSim", 2048, NULL, 5, NULL);
}

// Buttons
unsigned long lastBtnTime = 0;
void checkButtons() {
  if (millis() - lastBtnTime < 200) return;
  
  if (digitalRead(BTN_RED) == LOW) {
    lastBtnTime = millis();
    lastAction = "RESET";
    wifiPackets = 0;
    bleDevices = 0;
    zigbeePackets = 0;
  }
  else if (digitalRead(BTN_BLUE1) == LOW) {
    lastBtnTime = millis();
    lastAction = "BLUE1";
  }
  else if (digitalRead(BTN_BLUE2) == LOW) {
    lastBtnTime = millis();
    lastAction = "BLUE2";
  }
}

void drawRoundedFrame(int x, int y, int w, int h, int r, int color) {
  display.fillRoundRect(x, y, w, h, r, color);
}

void drawUI() {
  display.clearDisplay();
  
  display.setTextSize(1);
  display.setTextColor(SSD1306_BLACK, SSD1306_WHITE);
  display.fillRect(0, 0, 128, 12, SSD1306_WHITE);
  display.setCursor(8, 2);
  display.print("W.O.R.M v3.0");
  
  display.fillRect(120 - animStep % 8, 4, 3, 5, SSD1306_BLACK);
  
  drawRoundedFrame(2, 16, 58, 20, 6, SSD1306_WHITE);
  drawRoundedFrame(68, 16, 58, 20, 6, SSD1306_WHITE);
  drawRoundedFrame(35, 40, 58, 22, 6, SSD1306_WHITE);
  
  display.setTextColor(SSD1306_WHITE);
  
  display.setCursor(6, 20);
  display.print("WiFi:");
  display.setCursor(40, 20);
  display.print(wifiPackets);
  display.setCursor(6, 28);
  display.print("pkts");
  
  display.setCursor(74, 20);
  display.print("BLE:");
  display.setCursor(112, 20);
  display.print(bleDevices);
  display.setCursor(74, 28);
  display.print("devs");
  
  display.setCursor(39, 44);
  display.print("Zigbee:");
  display.setCursor(80, 44);
  display.print(zigbeePackets);
  display.setCursor(39, 52);
  display.print("pkts");
  
  display.display();
}

void setup() {
  Serial.begin(115200);
  
  Wire.begin(D4, D5);
  if(!display.begin(SSD1306_SWITCHCAPVCC, 0x3C)) {
    Serial.println("OLED fail");
    for(;;);
  }
  display.clearDisplay();
  
  pinMode(BTN_RED, INPUT_PULLUP);
  pinMode(BTN_BLUE1, INPUT_PULLUP);
  pinMode(BTN_BLUE2, INPUT_PULLUP);
  
  WiFi.mode(WIFI_MODE_NULL);
  esp_wifi_init();
  esp_wifi_start();
  setupWiFiSniffer();
  
  setupBLESniffer();
  setupZigbeeSim();
  
  lastAction = "ACTIVE";
}

void loop() {
  checkButtons();
  
  if (millis() - lastBleScan > 10000) {
    lastBleScan = millis();
    runBLEScan();
  }
  
  if (millis() - lastAnim > 100) {
    lastAnim = millis();
    animStep++;
    drawUI();
  }
  
  delay(50);
}