#include <Wire.h>
#include <LiquidCrystal_I2C.h>

// LCD Configuration (confirm I2C address with scanner)
LiquidCrystal_I2C lcd(0x27, 16, 2);

// Button Pins (adjust according to your actual wiring)
#define BTN_LEFT   2
#define BTN_RIGHT  3
#define BTN_UP     4
#define BTN_DOWN   5
#define BTN_MODE   6

// Frequency management
int frequency = 1000;  // Initial frequency
int selectedDigit = 0;
int digits[4] = {1, 0, 0, 0};
unsigned long lastDebounce = 0;
#define DEBOUNCE_DELAY 200

// System modes
enum Mode {FREQ_MODE, WIFI_MODE, PREPARE_MODE, WAVE_MODE};
Mode currentMode = FREQ_MODE;

void setup() {
  // Initialize buttons with internal pullups
  pinMode(BTN_LEFT, INPUT_PULLUP);
  pinMode(BTN_RIGHT, INPUT_PULLUP);
  pinMode(BTN_UP, INPUT_PULLUP);
  pinMode(BTN_DOWN, INPUT_PULLUP);
  pinMode(BTN_MODE, INPUT_PULLUP);

  // Initialize LCD
  lcd.init();
  lcd.backlight();
  lcd.blink();
  
  updateDisplay();
}

void loop() {
  handleModeButton();
  handleButtons();
}

void handleModeButton() {
  if (debounceRead(BTN_MODE)) {
    currentMode = static_cast<Mode>((currentMode + 1) % 4);
    lcd.clear();
    lcd.noBlink();
    
    switch(currentMode) {
      case FREQ_MODE:
        lcd.blink();
        updateFrequencyDisplay();
        break;
      case WIFI_MODE:
        showWiFiPassword();
        break;
      case PREPARE_MODE:
        lcd.print("Prepare Mode");
        lcd.setCursor(0, 1);
        lcd.print("Press Start");
        break;
      case WAVE_MODE:
        showWaveform();
        break;
    }
  }
}

void handleButtons() {
  if (currentMode == FREQ_MODE) {
    if (debounceRead(BTN_LEFT)) {
      selectedDigit = constrain(selectedDigit - 1, 0, 3);
      updateFrequencyDisplay();
    }
    if (debounceRead(BTN_RIGHT)) {
      selectedDigit = constrain(selectedDigit + 1, 0, 3);
      updateFrequencyDisplay();
    }
    if (debounceRead(BTN_UP)) {
      digits[selectedDigit] = (digits[selectedDigit] + 1) % 10;
      updateFrequency();
      updateFrequencyDisplay();
    }
    if (debounceRead(BTN_DOWN)) {
      digits[selectedDigit] = (digits[selectedDigit] - 1 + 10) % 10;
      updateFrequency();
      updateFrequencyDisplay();
    }
  }
  else if (currentMode == WAVE_MODE) {
    if (debounceRead(BTN_UP) || debounceRead(BTN_DOWN)) {
      showWaveform();
    }
  }
}

bool debounceRead(int pin) {
  if (digitalRead(pin) == LOW && (millis() - lastDebounce) > DEBOUNCE_DELAY) {
    lastDebounce = millis();
    while(digitalRead(pin) == LOW); // Wait for release
    return true;
  }
  return false;
}

void updateFrequency() {
  frequency = digits[0]*1000 + digits[1]*100 + digits[2]*10 + digits[3];
}

void updateFrequencyDisplay() {
  lcd.setCursor(0, 0);
  lcd.print("Freq: ");
  for(int i=0; i<4; i++) {
    lcd.print(digits[i]);
  }
  lcd.print(" Hz  ");
  lcd.setCursor(6 + selectedDigit, 0);
}

void showWiFiPassword() {
  lcd.setCursor(0, 0);
  lcd.print("WiFi Password: ");
  lcd.setCursor(0, 1);
  lcd.print(generatePassword());
}

String generatePassword() {
  String pass;
  for(int i=0; i<8; i++) {
    pass += (char)random(48, 90);
  }
  return pass;
}

void showWaveform() {
  static int waveType = 0;
  const String waves[4] = {"Sine", "Square", "Triangle", "Sawtooth"};
  
  if (debounceRead(BTN_UP)) waveType = (waveType + 1) % 4;
  if (debounceRead(BTN_DOWN)) waveType = (waveType - 1 + 4) % 4;
  
  lcd.clear();
  lcd.print(waves[waveType]);
  lcd.setCursor(0, 1);
  lcd.print("Freq: ");
  lcd.print(frequency);
  lcd.print(" Hz");
}

void updateDisplay() {
  lcd.clear();
  switch(currentMode) {
    case FREQ_MODE:
      updateFrequencyDisplay();
      break;
    case WIFI_MODE:
      showWiFiPassword();
      break;
    case PREPARE_MODE:
      lcd.print("Prepare Mode");
      break;
    case WAVE_MODE:
      showWaveform();
      break;
  }
}