#define MMdebug

/**

   Hardware:
   - ESP8266 ESP-01 (GPIO0=SDA, GPIO2=SCL for I2C OLED; GPIO3 as input for config button).
   - OLED display SSD1306 128x64 via I2C (address 0x3C).

   Note: Uses custom fonts (FreeMonoBold18pt7b for time, FreeMonoBold12pt7b for temperature).
   All other text uses default font. Display is rotated 180 degrees (setRotation(2)).
*/

/* Source: https://github.com/witnessmenow/WiFi-Tetris-Clock/blob/master/ESP32%20or%20TinyPICO/EzTimeTetrisClockESP32/EzTimeTetrisClockESP32.ino */
/*******************************************************************
    Tetris clock that fetches its time Using the EzTimeLibrary

    For use with the ESP32 or TinyPICO
    MM: Display: 128*64   SSD_1306
 *                                                                 *
    Written by Brian Lough
    YouTube: https://www.youtube.com/brianlough
    Tindie: https://www.tindie.com/stores/brianlough/
    Twitter: https://twitter.com/witnessmenow
 *******************************************************************/
/*
  #if defined(ESP8266)
  #include <ESP8266WiFi.h>
  #include <ESP8266WebServer.h>
  // Pin definitions (ESP-01 / ESP8266):
  const uint8_t SDA_PIN = 0;           // I2C SDA connected to GPIO0
  const uint8_t SCL_PIN = 2;           // I2C SCL connected to GPIO2
  const uint8_t BUTTON_PIN = 3;        // Button on GPIO3 (RX pin)
  #elif defined(ESP32)
  #include <WiFi.h>
  #include <WebServer.h>
    // MM:
    // Pin definitions (ESP32):
    const uint8_t SDA_PIN = 21;           // I2C SDA connected to GPIO0
    const uint8_t SCL_PIN = 22;           // I2C SCL connected to GPIO2
    const uint8_t BUTTON_PIN = 16;        // Button on GPIO3 (RX pin)
  #endif
*/


#include <Wire.h>
#include <Adafruit_GFX.h>
#include <Adafruit_SSD1306.h>

#ifdef ESP32
#include <WiFi.h>
#include "esp_wifi.h"
#elif defined(ESP8266)
#include <ESP8266WiFi.h>
extern "C" {
#include "user_interface.h"
}
#endif

#define SCREEN_WIDTH 128
#define SCREEN_HEIGHT 64
#define OLED_RESET -1
#define OLED_ADDR 0x3C

Adafruit_SSD1306 display(SCREEN_WIDTH, SCREEN_HEIGHT, &Wire, OLED_RESET);
// MM: TODO: wiring muss noch vorgenommen werden, um generisch bei SCL /SDA zu sein.

#define BTN_PIN 16 // 0
#define LED_PIN 2 // to check...
#define MAX_CH 13
#define MAX_MACS 25
#define LINE_Y 12
#define LINE_VAL 47

unsigned long prevTime = 0, curTime = 0, switchTime = 0;
int curChannel = 1;

struct Device {
  String mac;
  int rssi;
};

Device devices[MAX_MACS];
int deviceCount = 0;

unsigned int val[128];
unsigned long pkts = 0, no_deauths = 0, deauths = 0;
unsigned long maxVal = 0;
double multiplicator = 0.0;
bool showGraph = true; // Start mit Graphen

bool ssid_locked = false;
char ssid[33] = "unknown";


#ifdef ESP32
void sniffer(void *buf, wifi_promiscuous_pkt_type_t type) {
  const wifi_promiscuous_pkt_t *pkt = (wifi_promiscuous_pkt_t *)buf;
  const wifi_pkt_rx_ctrl_t *ctrl = &pkt->rx_ctrl;
  const uint8_t *payload = pkt->payload;

  pkts++;
  if (payload[0] == 0xC0 || payload[0] == 0xA0) {
    deauths++;
  }

 if (!ssid_locked && payload[0] == 0x80) {
    uint8_t len = payload[37];
    if (len > 0 && len < 33) {
      memcpy(ssid, &payload[38], len);
      ssid[len] = 0;
      ssid_locked = true;
    }
  }

  String mac = "";
  for (int i = 10; i < 16; i++) {
    mac += String(pkt->payload[i], HEX);
    if (i < 15) mac += ":";
  }
  mac.toUpperCase();

  int rssi = ctrl->rssi;

  // Prüfen, ob MAC bereits erfasst wurde
  bool found = false;
  for (int i = 0; i < deviceCount; i++) {
    if (devices[i].mac == mac) {
      devices[i].rssi = rssi; // RSSI aktualisieren
      found = true;
      break;
    }
  }

  // Neue MAC-Adresse hinzufügen
  if (!found && deviceCount < MAX_MACS) {
    devices[deviceCount].mac = mac;
    devices[deviceCount].rssi = rssi;
    deviceCount++;
  }

  pkts++; // Paketanzahl erhöhen
}
#elif defined(ESP8266)
void sniffer(uint8_t *buf, uint16_t len) {
  const sniffer_buf *pkt = (sniffer_buf *)buf;
  const sniffer_buf2 *pkt2 = (sniffer_buf2 *)buf;
  int len_pkt = pkt->rx_ctrl.legacy_length;
  if (len_pkt == 0) return; // Leere Pakete ignorieren
  const int8_t rssi = pkt->rx_ctrl.rssi;

  String mac = "";
  for (int i = 10; i < 16; i++) {
    mac += String(pkt->buf[i], HEX);
    if (i < 15) mac += ":";
  }
  mac.toUpperCase();

  // Prüfen, ob MAC bereits erfasst wurde
  bool found = false;
  for (int i = 0; i < deviceCount; i++) {
    if (devices[i].mac == mac) {
      devices[i].rssi = rssi; // RSSI aktualisieren
      found = true;
      break;
    }
  }

  // Neue MAC-Adresse hinzufügen
  if (!found && deviceCount < MAX_MACS) {
    devices[deviceCount].mac = mac;
    devices[deviceCount].rssi = rssi;
    deviceCount++;
  }

  pkts++; // Paketanzahl erhöhen
}
#endif

void getMultiplicator() {
  maxVal = val[0];
  for (int i = 1; i < SCREEN_WIDTH; i++) {
    if (val[i] > maxVal) maxVal = val[i];
  }
  if (maxVal < 1) maxVal = 1;
  multiplicator = maxVal > LINE_VAL ? (double)LINE_VAL / maxVal : 1;
}

void displayGraph() {
  display.clearDisplay();
  display.setRotation(2);  // for "Weatherclock" kit
  display.setTextSize(1);
  display.setCursor(0, 0);
  display.print("Ch:");
  display.print(curChannel);
      display.print(" ");
    display.print(ssid);
    display.print(" ");
    display.print(no_deauths);
    /*
  display.print(" Pkts:");
  display.print(pkts);
*/
  display.drawLine(0, LINE_Y, SCREEN_WIDTH - 1, LINE_Y, SSD1306_WHITE);

  for (int i = 0; i < SCREEN_WIDTH; i++) {
    int bar = val[i] * multiplicator;
    display.drawLine(i, SCREEN_HEIGHT - 1, i, SCREEN_HEIGHT - 1 - bar, SSD1306_WHITE);
  }

  display.display();
}

void displayDevices() {
  display.clearDisplay();
  display.setTextSize(1);
  display.setCursor(0, 0);
  display.print("Ch:");
  display.print(curChannel);
  display.println();

  for (int i = 0; i < deviceCount && i < 5; i++) { // Maximal 5 Geräte anzeigen
    display.setCursor(0, 12 + i * 10);
    display.print(devices[i].mac);
    display.setCursor(90, 12 + i * 10);
    display.print(devices[i].rssi);
    display.print(" dBm");
  }

  display.display();
}

void setup() {
  Serial.begin(115200);
  pinMode(BTN_PIN, INPUT_PULLUP);

  if (!display.begin(SSD1306_SWITCHCAPVCC, OLED_ADDR)) {
    while (1)
      ;
  }

  display.clearDisplay();
  display.setTextSize(1);
  display.setTextColor(SSD1306_WHITE);
  display.display();
  delay(2000);

#ifdef ESP32
  wifi_init_config_t cfg = WIFI_INIT_CONFIG_DEFAULT();
  esp_wifi_init(&cfg);
  esp_wifi_set_storage(WIFI_STORAGE_RAM);
  esp_wifi_set_mode(WIFI_MODE_NULL);
  esp_wifi_start();
  esp_wifi_set_promiscuous(true);
  esp_wifi_set_promiscuous_rx_cb(&sniffer);
#elif defined(ESP8266)
  wifi_set_opmode(STATION_MODE);
  wifi_promiscuous_enable(1);
  wifi_set_promiscuous_rx_cb(sniffer);
#endif
}

void loop() {
  curTime = millis();

  if (digitalRead(BTN_PIN) == LOW) {
    curChannel++;
    if (curChannel > MAX_CH) curChannel = 1;
#ifdef ESP32
    esp_wifi_set_channel(curChannel, WIFI_SECOND_CHAN_NONE);
#elif defined(ESP8266)
    wifi_set_channel(curChannel);
#endif
    deviceCount = 0; // Liste leeren
    delay(500);
  }

  if (curTime - prevTime >= 1000) {
    prevTime = curTime;

    // Werte für Graph aktualisieren
    for (int i = 0; i < SCREEN_WIDTH - 1; i++) {
      val[i] = val[i + 1];
    }
    val[SCREEN_WIDTH - 1] = pkts;

    getMultiplicator();

    digitalWrite(LED_PIN, deauths > PACKET_RATE ? LOW : HIGH);

    if (pkts == 0) pkts = deauths;
    no_deauths = pkts - deauths;
    // pkts = 0;

    // Zwischen Graph und Geräteliste wechseln
    if (curTime - switchTime >= 5000) { // Alle 5 Sekunden wechseln
      switchTime = curTime;
      showGraph = !showGraph;
    }

    if (showGraph) {
      displayGraph();
    } else {
      displayDevices();
    }
  }
}