#include <U8g2lib.h>
#include <EEPROM.h>

// OLED Display
U8G2_SSD1306_128X64_NONAME_F_HW_I2C u8g2(U8G2_R0, U8X8_PIN_NONE);

// Pins
#define BTN_UP     2
#define BTN_DOWN   3
#define BTN_LEFT   4
#define BTN_RIGHT  5
#define BUZZER     6
#define TRIG_PIN   9
#define ECHO_PIN   10

// EEPROM Addresses
#define ADDR_UNIT           0
#define ADDR_BRIGHTNESS     1
#define ADDR_ALARM_DISTANCE 2

// Menu arrays
const char* mainMenu[] = { "Distance Meter", "Set Alarm", "Object Counter", "Settings" };
const int mainMenuLength = 4;

const char* settingsMenu[] = { "Unit", "Brightness", "Alarm Dist", "Back" };
const int settingsLength = 4;

int mainMenuIndex = 0;
int settingsMenuIndex = 0;

bool inSubMenu = false;
bool inSettings = false;

// Settings
int unitMode = 0;         // 0 = cm, 1 = mm, 2 = in
int brightnessLevel = 5;  // 1–10
int alarmDistance = 10;   // cm

int objectCount = 0;
bool objectDetected = false;

// ==== SETUP ====

void setup() {
  u8g2.begin();

  pinMode(BTN_UP, INPUT_PULLUP);
  pinMode(BTN_DOWN, INPUT_PULLUP);
  pinMode(BTN_LEFT, INPUT_PULLUP);
  pinMode(BTN_RIGHT, INPUT_PULLUP);
  pinMode(BUZZER, OUTPUT);
  pinMode(TRIG_PIN, OUTPUT);
  pinMode(ECHO_PIN, INPUT);

  loadSettings();
  applyBrightness();

  showSplashScreen();
}

// ==== MAIN LOOP ====

void loop() {
  long dist = getDistance();

  if (!inSubMenu) {
    handleMenuNavigation();
    drawMainMenu();
  } else if (inSettings) {
    handleSettingsMenu();
  } else {
    switch (mainMenuIndex) {
      case 0: drawDistance(dist); break;
      case 1: drawAlarm(dist); break;
      case 2: drawCounter(dist); break;
      case 3: inSettings = true; settingsMenuIndex = 0; break;
    }

    if (digitalRead(BTN_LEFT) == LOW) {
      beep();
      inSubMenu = false;
      inSettings = false;
      digitalWrite(BUZZER, LOW);
      delay(200);
    }
  }

  delay(100);
}

// ==== DISTANCE SENSOR ====

long getDistance() {
  digitalWrite(TRIG_PIN, LOW);
  delayMicroseconds(2);
  digitalWrite(TRIG_PIN, HIGH);
  delayMicroseconds(10);
  digitalWrite(TRIG_PIN, LOW);

  long duration = pulseIn(ECHO_PIN, HIGH, 30000);
  float cm = duration * 0.0343 / 2;
  return (cm > 0) ? cm : 999;
}

// ==== SETTINGS ====

void loadSettings() {
  unitMode = EEPROM.read(ADDR_UNIT);
  if (unitMode > 2) unitMode = 0;

  brightnessLevel = EEPROM.read(ADDR_BRIGHTNESS);
  if (brightnessLevel < 1 || brightnessLevel > 10) brightnessLevel = 5;

  alarmDistance = EEPROM.read(ADDR_ALARM_DISTANCE);
  if (alarmDistance < 5 || alarmDistance > 100) alarmDistance = 10;
}

void saveSettings() {
  EEPROM.write(ADDR_UNIT, unitMode);
  EEPROM.write(ADDR_BRIGHTNESS, brightnessLevel);
  EEPROM.write(ADDR_ALARM_DISTANCE, alarmDistance);
}

void applyBrightness() {
  int contrast = map(brightnessLevel, 1, 10, 10, 255);
  u8g2.setContrast(contrast);
}

// ==== BEEP ====

void beep() {
  tone(BUZZER, 1000, 50);
  delay(60);
}

// ==== MENU NAVIGATION ====

void handleMenuNavigation() {
  if (digitalRead(BTN_UP) == LOW) {
    mainMenuIndex = (mainMenuIndex - 1 + mainMenuLength) % mainMenuLength;
    beep();
    delay(200);
  }
  if (digitalRead(BTN_DOWN) == LOW) {
    mainMenuIndex = (mainMenuIndex + 1) % mainMenuLength;
    beep();
    delay(200);
  }
  if (digitalRead(BTN_RIGHT) == LOW) {
    inSubMenu = true;
    beep();
    delay(200);
  }
}

void handleSettingsNavigation() {
  if (digitalRead(BTN_UP) == LOW) {
    settingsMenuIndex = (settingsMenuIndex - 1 + settingsLength) % settingsLength;
    beep();
    delay(200);
  }
  if (digitalRead(BTN_DOWN) == LOW) {
    settingsMenuIndex = (settingsMenuIndex + 1) % settingsLength;
    beep();
    delay(200);
  }

  if (digitalRead(BTN_RIGHT) == LOW) {
    beep();
    if (settingsMenuIndex == 0) {
      unitMode = (unitMode + 1) % 3;
      saveSettings();
    } else if (settingsMenuIndex == 1) {
      brightnessLevel++;
      if (brightnessLevel > 10) brightnessLevel = 1;
      applyBrightness();
      saveSettings();
    } else if (settingsMenuIndex == 2) {
      alarmDistance += 1;
      if (alarmDistance > 100) alarmDistance = 5;
      saveSettings();
    } else {
      inSubMenu = false;
      inSettings = false;
    }
    delay(200);
  }

  if (digitalRead(BTN_LEFT) == LOW && settingsMenuIndex == 2) {
    beep();
    alarmDistance -= 1;
    if (alarmDistance < 5) alarmDistance = 100;
    saveSettings();
    delay(200);
  }
}

// ==== DRAW MENUS ====

void drawCenteredTitle(const char* title) {
  u8g2.setFont(u8g2_font_ncenB08_tr);
  int x = (128 - u8g2.getStrWidth(title)) / 2;
  u8g2.drawStr(x, 14, title);
}

void drawMainMenu() {
  u8g2.clearBuffer();
  drawCenteredTitle("===== MENU =====");

  u8g2.setFont(u8g2_font_6x10_tr);
  for (int i = 0; i < mainMenuLength; i++) {
    u8g2.setCursor(6, 30 + i * 12); // vertical spacing increased
    if (i == mainMenuIndex) {
      u8g2.print("-> ");
    } else {
      u8g2.print("   ");
    }
    u8g2.print(mainMenu[i]);
  }

  u8g2.sendBuffer();
}

void handleSettingsMenu() {
  handleSettingsNavigation();

  u8g2.clearBuffer();
  drawCenteredTitle("SETTINGS");

  u8g2.setFont(u8g2_font_6x10_tr);
  for (int i = 0; i < settingsLength; i++) {
    u8g2.setCursor(10, 30 + i * 12); // vertical spacing increased
    if (i == settingsMenuIndex) {
      u8g2.print("-> ");
    } else {
      u8g2.print("   ");
    }

    if (i == 0) {
      u8g2.print("Unit: ");
      if (unitMode == 0) u8g2.print("cm");
      else if (unitMode == 1) u8g2.print("mm");
      else u8g2.print("in");
    } else if (i == 1) {
      u8g2.print("Brightness: ");
      u8g2.print(brightnessLevel);
    } else if (i == 2) {
      u8g2.print("Alarm Dist: ");
      u8g2.print(alarmDistance);
      u8g2.print(" cm");
    } else {
      u8g2.print("Back");
    }
  }

  u8g2.sendBuffer();
}

// ==== FEATURES ====

void drawDistance(long dist) {
  u8g2.clearBuffer();
  drawCenteredTitle("DISTANCE");

  float val;
  const char* unitStr;

  switch (unitMode) {
    case 0: val = dist; unitStr = "cm"; break;
    case 1: val = dist * 10.0; unitStr = "mm"; break;
    case 2: val = dist / 2.54; unitStr = "in"; break;
  }

  u8g2.setFont(u8g2_font_6x10_tr);
  u8g2.setCursor(10, 40);
  u8g2.print("Value: ");
  u8g2.print(val, 1);
  u8g2.print(" ");
  u8g2.print(unitStr);

  u8g2.sendBuffer();
}

void drawAlarm(long dist) {
  u8g2.clearBuffer();
  drawCenteredTitle("ALARM");

  u8g2.setFont(u8g2_font_6x10_tr);
  if (dist <= alarmDistance) {
    digitalWrite(BUZZER, HIGH);
    u8g2.setCursor(10, 40);
    u8g2.print("TOO CLOSE!");
  } else {
    digitalWrite(BUZZER, LOW);
    u8g2.setCursor(10, 40);
    u8g2.print("Distance OK");
  }

  u8g2.setCursor(10, 55);
  u8g2.print(dist);
  u8g2.print(" cm");

  u8g2.sendBuffer();
}

void drawCounter(long dist) {
  static unsigned long lastSeen = 0;

  if (dist > 0 && dist < 30) {
    if (!objectDetected) {
      objectDetected = true;
      objectCount++;
      lastSeen = millis();
    }
  } else if (millis() - lastSeen > 500) {
    objectDetected = false;
  }

  u8g2.clearBuffer();
  drawCenteredTitle("COUNTER");

  u8g2.setFont(u8g2_font_6x10_tr);
  u8g2.setCursor(10, 35);
  u8g2.print("Objects: ");
  u8g2.print(objectCount);

  u8g2.setCursor(10, 50);
  u8g2.print("Dist: ");
  u8g2.print(dist);
  u8g2.print(" cm");

  u8g2.sendBuffer();
}

// ==== SPLASH SCREEN ====

void showSplashScreen() {
  u8g2.clearBuffer();

  u8g2.setFont(u8g2_font_ncenB10_tr);
  int titleX = (128 - u8g2.getStrWidth("Ultrasonic Sensor Tool")) / 2;
  u8g2.drawStr(titleX, 22, "Ultrasonic Sensor Tool");

  u8g2.setFont(u8g2_font_6x10_tr);
  int subX = (128 - u8g2.getStrWidth("Via MaDDy Regal")) / 2;
  u8g2.drawStr(subX, 40, "Via MaDDy Regal");

  u8g2.sendBuffer();

  // Optional tone sequence
  tone(BUZZER, 1000, 150);
  delay(200);
  tone(BUZZER, 1200, 150);
  delay(200);
  tone(BUZZER, 1400, 150);
  delay(500);
}