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

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

Adafruit_SSD1306 display(SCREEN_WIDTH, SCREEN_HEIGHT, &Wire, OLED_RESET);

const int UP_PIN = 32;
const int DOWN_PIN = 33;
const int OK_PIN = 27;
const int BACK_PIN = 35;

const char* menuItems[] = {"Scan", "Scrape", "Brute-force", "Decrypt", "Settings", "Internet Connection"};
int currentOption = 0;
int selectedNetwork = 0;
int numNetworks = 0;

void setup() {
  pinMode(UP_PIN, INPUT_PULLUP);
  pinMode(DOWN_PIN, INPUT_PULLUP);
  pinMode(OK_PIN, INPUT_PULLUP);
  pinMode(BACK_PIN, INPUT_PULLUP);

  Serial.begin(9600);

  if (!display.begin(SSD1306_SWITCHCAPVCC, SSD1306_I2C_ADDRESS)) {
    Serial.println(F("SSD1306 allocation failed"));
    for(;;);
  }

  display.clearDisplay();
  display.setTextSize(1);
  display.setTextColor(SSD1306_WHITE);
  display.setCursor(10, 10);
  display.println("  _   _       _____  ");
  display.println(" | \\ |\"|     |\"_  /u ");
  display.println("<|  \\| |>    U / //  ");
  display.println("U| |\\  |u    \\/ /_   ");
  display.println(" |_| \\_|     /____|  ");
  display.println(" ||   \\\\,-.  _//<<,- ");
  display.println(" (_\")  (_/  (__) (_/ ");
  display.display();

  delay(2000);
}

void loop() {
  displayMainMenu();

  if (digitalRead(UP_PIN) == LOW) {
    currentOption = (currentOption - 1 + 6) % 6;
    delay(200);
  }
  if (digitalRead(DOWN_PIN) == LOW) {
    currentOption = (currentOption + 1) % 6;
    delay(200);
  }
  if (digitalRead(OK_PIN) == LOW) {
    navigateToPage(currentOption);
    delay(200);
  }
}

void displayMainMenu() {
  display.clearDisplay();
  display.setCursor(0, 0);
  display.println("Main Menu:");
  for (int i = 0; i < 6; i++) {
    display.print((i == currentOption) ? "> " : "  ");
    display.println(menuItems[i]);
  }
  display.display();
}

void navigateToPage(int option) {
  switch (option) {
    case 0:
      displayScanPage();
      break;
    case 1:
      displaySubMenu("Scrape Menu");
      break;
    case 5:
      displayInternetConnectionPage();
      break;
    default:
      displaySubMenu(String(menuItems[option]) + " Menu");
      break;
  }
}

void displaySubMenu(String pageTitle) {
  display.clearDisplay();
  display.setCursor(0, 0);
  display.println(pageTitle);
  display.display();
}

void displayScanPage() {
  display.clearDisplay();
  display.setCursor(0, 0);
  display.println("Scanning for WiFi...");
  display.display();

  numNetworks = WiFi.scanNetworks();

  if (numNetworks == 0) {
    display.println("No networks found");
  } else {
    displayNetworkList();
  }

  delay(2000);
}

void displayNetworkList() {
  display.clearDisplay();
  display.setCursor(0, 0);
  display.println("WiFi Networks:");

  for (int i = 0; i < numNetworks; ++i) {
    if (i == selectedNetwork) {
      display.print("> ");
    } else {
      display.print("  ");
    }
    display.print(WiFi.SSID(i));
    display.print(" (");
    display.print(WiFi.RSSI(i));
    display.println("dBm)");
    if (i >= 5) { // Limit to first 6 networks to fit on the screen
      break;
    }
  }
  display.display();
}

void displayInternetConnectionPage() {
  display.clearDisplay();
  display.setCursor(0, 0);
  display.println("Internet Connection");
  display.println();
  if (WiFi.status() != WL_CONNECTED) {
    display.println("1. Search for Wi-Fi");
    display.println("2. Cancel");
    display.display();

    while (true) {
      if (digitalRead(OK_PIN) == LOW) {
        int option = currentOption;
        if (option == 0) {
          currentOption = 0;
          displayScanPage();
          break;
        } else if (option == 1) {
          break;
        }
      }
      if (digitalRead(BACK_PIN) == LOW) {
        break;
      }
    }
  } else {
    display.println("Already connected!");
    display.display();
    delay(2000);
  }
}