#include <U8g2lib.h>
#include <Bounce2.h>

byte button_pins[] = {9, 5, 6}; // button pins, 9 = up, 5 = down, 6 = select
#define NUMBUTTONS sizeof(button_pins)
Bounce *buttons = new Bounce[NUMBUTTONS];

U8X8_SSD1306_128X64_NONAME_HW_I2C display(U8X8_PIN_NONE);

#define MAIN_MENU_SIZE 3
char *mainMenu[MAIN_MENU_SIZE] = {"Set Pulses", "Set Frequency", "Start"};

int cursor = 0;
int numPulses = 10;     // Default number of pulses
int frequency = 5;   // Default frequency in Hz
const int outputPin = 7; // Pin to output pulses

bool inSubmenu = false;
int submenuValue = 0;

void setup() {
  Serial.begin(9600);
  pinMode(outputPin, OUTPUT);
  digitalWrite(outputPin, LOW);

  // Make input & enable pull-up resistors on switch pins
  for (int i = 0; i < NUMBUTTONS; i++) {
    buttons[i].attach(button_pins[i], INPUT_PULLUP); // setup the bounce instance for the current button
    buttons[i].interval(25); // interval in ms
  }

  display.begin();
  display.setPowerSave(0);
  display.setFont(u8x8_font_pxplusibmcgathin_f);

  showSplashScreen();
  showMainMenu();
}

void loop() {
  for (int i = 0; i < NUMBUTTONS; i++) {
    buttons[i].update(); // Update the Bounce instance
    if (buttons[i].fell()) { // If it fell
      if (inSubmenu) {
        handleSubmenuInput(i);
      } else {
        handleMainMenuInput(i);
      }
    }
  }
}

/**
 * Show a splash screen with the text "Pulse Generation".
 */
void showSplashScreen() {
  display.clearDisplay();

  // Display "Pulse" in a larger font
  display.setFont(u8x8_font_chroma48medium8_r); // Set a larger font
  display.drawString(2, 2, "Pulse"); // Draw "Pulse" at the top

  // Display "Generation" in a smaller font
  display.setFont(u8x8_font_pxplusibmcgathin_f); // Set a smaller font
  display.drawString(3, 4, "Generation"); // Draw "Generation" on the next line

  delay(2000); // Display the splash screen for 2 seconds
  display.clearDisplay();
}

/**
 * Show a finish screen with the text "Finished".
 */
void showFinishScreen() {
  display.clearDisplay();

  // Display "Finished" in a larger font
  display.setFont(u8x8_font_chroma48medium8_r); // Set a larger font
  display.drawString(4, 3, "Finished"); // Draw "Finished" centered vertically

  delay(2000); // Display the finish screen for 2 seconds
  display.clearDisplay();
}

/**
 * Clear display and show the main menu.
 */
void showMainMenu() {
  inSubmenu = false;
  cursor = 0;
  display.clearDisplay();
  for (int i = 0; i < MAIN_MENU_SIZE; i++) {
    display.drawString(2, i, mainMenu[i]);
  }
  display.setCursor(0, 0);
  display.print('>');
}

/**
 * Handle input in the main menu.
 */
void handleMainMenuInput(int buttonIndex) {
  // Erase previous cursor:
  display.setCursor(0, cursor);
  display.print(' ');

  if (buttonIndex == 0) { // Up button
    cursor++;
    if (cursor > (MAIN_MENU_SIZE - 1)) cursor = 0;
  } else if (buttonIndex == 1) { // Down button
    cursor--;
    if (cursor < 0) cursor = (MAIN_MENU_SIZE - 1);
  } else if (buttonIndex == 2) { // Select button
    executeChoice(cursor);
  }

  // Show cursor at new line:
  display.setCursor(0, cursor);
  display.print('>');
}

/**
 * Handle input in a submenu (for setting pulses or frequency).
 */
void handleSubmenuInput(int buttonIndex) {
  if (buttonIndex == 1) { // Up button
    submenuValue += 1; // Increment by 1
  } else if (buttonIndex == 0) { // Down button
    submenuValue -= 1; // Decrement by 1
    if (submenuValue < 0) submenuValue = 0; // Prevent negative values
  } else if (buttonIndex == 2) { // Select button
    if (cursor == 0) {
      numPulses = submenuValue;
      Serial.print("Number of pulses set to: ");
    } else {
      frequency = submenuValue;
      Serial.print("Frequency set to: ");
    }
    Serial.println(submenuValue);
    showMainMenu();
    return;
  }

  // Update submenu display:
  display.clearDisplay();
  display.setCursor(0, 0);
  if (cursor == 0) {
    display.print("Set Pulses:");
  } else {
    display.print("Set Frequency:");
  }
  display.setCursor(0, 1);
  display.print(submenuValue);
}

/**
 * Execute the task which matches the chosen menu item.
 */
void executeChoice(int choice) {
  inSubmenu = true;
  if (choice == 0) { // Set number of pulses
    submenuValue = numPulses;  // Initialize submenuValue with current pulse count
    display.clearDisplay();
    display.setCursor(0, 0);
    display.print("Set Pulses:");
    display.setCursor(0, 1);
    display.print(submenuValue);
  } else if (choice == 1) { // Set frequency
    submenuValue = frequency; // Initialize submenuValue with current frequency
    display.clearDisplay();
    display.setCursor(0, 0);
    display.print("Set Frequency:");
    display.setCursor(0, 1);
    display.print(submenuValue);
  } else if (choice == 2) { // Start pulse generation
    generatePulses(numPulses, frequency);
    showMainMenu();
  }
}

/**
 * Generate pulses on the output pin and display a progress bar.
 */
void generatePulses(int pulseCount, int freq) {
    if (freq <= 0) {
        Serial.println("Frequency must be greater than 0.");
        return;
    }

    long periodMicroseconds = 1000000L / freq; // Total period in microseconds
    long pulseWidth = periodMicroseconds / 2;  // High and Low time (half of the period)

    Serial.print("Generating ");
    Serial.print(pulseCount);
    Serial.print(" pulses at ");
    Serial.print(freq);
    Serial.println(" Hz");
    Serial.print("Pulse Width ");
    Serial.println(pulseWidth);

    unsigned long previousTime;
    int barWidth = 0;

    //display.clearDisplay();
    //display.drawString(0, 0, "Generating pulses:");

    for (int i = 0; i < pulseCount; i++) {
        digitalWrite(outputPin, HIGH);
        previousTime = micros();
        while (micros() - previousTime < pulseWidth) {
            // Wait for pulseWidth duration
        }

        digitalWrite(outputPin, LOW);
        previousTime = micros();
        while (micros() - previousTime < pulseWidth) {
            // Wait for pulseWidth duration
        }

        // Calculate progress and update progress bar
        barWidth = (i + 1) * 128 / pulseCount; // Progress bar width (128 is the screen width)
        display.clearLine(7); // Clear previous progress bar
        for (int x = 0; x < barWidth; x += 8) {
            display.drawTile(x / 8, 7, 1, "\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF"); // Draw the progress bar
        }
        display.refreshDisplay(); // Refresh the display to show changes

        // Debugging output to confirm the loop is running correctly
        Serial.print("Pulse ");
        Serial.println(i + 1);
    }

    Serial.println("Pulse generation complete");
    delay(250);
    // Show the finish screen
    showFinishScreen();
}