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

#define SCREEN_WIDTH 128
#define SCREEN_HEIGHT 64
#define OLED_SDA 0  // D3
#define OLED_SCL 2  // D4
Adafruit_SSD1306 display(SCREEN_WIDTH, SCREEN_HEIGHT, &Wire, -1);

#define OUTPUT_PIN 15  // D8
#define INPUT_PIN  3   // RX
#define BTN_MODE   4   // D2
#define BTN_UP     5   // D1
#define BTN_DOWN   14  // D5
#define BTN_STEP   16  // D0
#define BTN_SWEEP  12  // D6
#define BTN_OUTPUT 13  // D7

enum Mode { FREQ_GEN, FREQ_READ, TONE_GEN };
Mode currentMode = FREQ_GEN;

int freq = 10; // Starting frequency in Hz
int stepSize = 10; // Initial step size
bool outputEnabled = true;
bool sweepEnabled = false;
unsigned long lastSweepTime = 0;
int sweepDir = 1;

enum StepSizeIndex { STEP_1HZ = 0, STEP_10HZ = 1, STEP_100HZ = 2 };
int currentStepIndex = STEP_10HZ;

const unsigned int notes[] = {261, 294, 329, 349, 392, 440, 493}; // C4, D4, E4, F4, G4, A4, B4
const char* noteNames[] = {"C4", "D4", "E4", "F4", "G4", "A4", "B4"};
int currentNoteIndex = 0;

void setup() {
  // Pin setup
  pinMode(OUTPUT_PIN, OUTPUT);
  pinMode(INPUT_PIN, INPUT);
  pinMode(BTN_MODE, INPUT_PULLUP);
  pinMode(BTN_UP, INPUT_PULLUP);
  pinMode(BTN_DOWN, INPUT_PULLUP);
  pinMode(BTN_STEP, INPUT_PULLUP);
  pinMode(BTN_SWEEP, INPUT_PULLUP);
  pinMode(BTN_OUTPUT, INPUT_PULLUP);

  Wire.begin(OLED_SDA, OLED_SCL);
  display.begin(SSD1306_SWITCHCAPVCC, 0x3C);
  display.clearDisplay();
  display.setTextColor(WHITE);
  display.setTextSize(1);

  tone(OUTPUT_PIN, freq); // Start with initial frequency
}

void loop() {
  handleButtons();
  handleSweep();
  updateDisplay();
  if (outputEnabled) {
    if (currentMode == FREQ_GEN || currentMode == TONE_GEN) {
      if (currentMode == TONE_GEN) {
        tone(OUTPUT_PIN, notes[currentNoteIndex]);
      } else {
        tone(OUTPUT_PIN, freq);
      }
    }
  } else {
    noTone(OUTPUT_PIN); // Disable tone when output is off
  }
}

void handleButtons() {
  static unsigned long lastDebounceTime = 0;
  static const unsigned long debounceDelay = 150;
  static bool lastModeState = HIGH, lastUpState = HIGH, lastDownState = HIGH, lastStepState = HIGH, lastSweepState = HIGH, lastOutputState = HIGH;

  if (millis() - lastDebounceTime < debounceDelay) return;

  bool buttonPressed = false;

  bool currentModeState = digitalRead(BTN_MODE);
  if (currentModeState == LOW && lastModeState == HIGH) {
    currentMode = (Mode)((currentMode + 1) % 3);
    if (currentMode == TONE_GEN) {
      currentNoteIndex = 0;
      freq = notes[currentNoteIndex];
    }
    buttonPressed = true;
  }
  lastModeState = currentModeState;

  bool currentUpState = digitalRead(BTN_UP);
  if (currentUpState == LOW && lastUpState == HIGH) {
    if (currentMode == TONE_GEN) {
      currentNoteIndex = (currentNoteIndex + 1) % 7;
      freq = notes[currentNoteIndex];
    } else {
      freq = min(freq + stepSize, 10000); // Increase frequency, max 1000 Hz
    }
    buttonPressed = true;
  }
  lastUpState = currentUpState;

  bool currentDownState = digitalRead(BTN_DOWN);
  if (currentDownState == LOW && lastDownState == HIGH) {
    if (currentMode == TONE_GEN) {
      currentNoteIndex = (currentNoteIndex + 6) % 7;
      freq = notes[currentNoteIndex];
    } else {
      freq = max(freq - stepSize, 1); // Decrease frequency, min 1 Hz
    }
    buttonPressed = true;
  }
  lastDownState = currentDownState;

  bool currentStepState = digitalRead(BTN_STEP);
  if (currentStepState == LOW && lastStepState == HIGH) {
    currentStepIndex = (currentStepIndex + 1) % 3;
    switch (currentStepIndex) {
      case STEP_1HZ: stepSize = 1; break;
      case STEP_10HZ: stepSize = 10; break;
      case STEP_100HZ: stepSize = 100; break;
    }
    buttonPressed = true;
  }
  lastStepState = currentStepState;

  bool currentSweepState = digitalRead(BTN_SWEEP);
  if (currentSweepState == LOW && lastSweepState == HIGH) {
    sweepEnabled = !sweepEnabled;
    buttonPressed = true;
  }
  lastSweepState = currentSweepState;

  bool currentOutputState = digitalRead(BTN_OUTPUT);
  if (currentOutputState == LOW && lastOutputState == HIGH) {
    outputEnabled = !outputEnabled;
    if (!outputEnabled) noTone(OUTPUT_PIN);
    buttonPressed = true;
  }
  lastOutputState = currentOutputState;

  if (buttonPressed) lastDebounceTime = millis();
}

void handleSweep() {
  if (!sweepEnabled || currentMode != FREQ_GEN) return;
  if (millis() - lastSweepTime > 100) {
    freq += sweepDir * (stepSize / 10.0);
    if (freq > 10000) { freq = 10000; sweepDir = -1; } // Max 1000 Hz
    else if (freq < 1) { freq = 1; sweepDir = 1; }
    lastSweepTime = millis();
  }
}

void updateDisplay() {
  display.clearDisplay();

  // Horizontal line separators for professional look
  display.drawLine(0, 16, SCREEN_WIDTH - 1, 16, WHITE);
  display.drawLine(0, 44, SCREEN_WIDTH - 1, 44, WHITE);

  // Mode (Top section)
  display.setTextSize(1);
  display.setCursor(0, 2);
  display.print("Mode: ");
  display.setTextSize(2);
  if (currentMode == FREQ_GEN) display.print("GEN");
  else if (currentMode == FREQ_READ) display.print("READ");
  else display.print("TONE");

  // Main value (Middle section)
  display.setTextSize(1);
  display.setCursor(0, 20);
  if (currentMode == TONE_GEN) {
    display.print("Note: ");
    display.setTextSize(2);
    // Center the note name horizontally (approximate, adjust x based on text length if needed)
    display.setCursor(36, 24);
    display.print(noteNames[currentNoteIndex]);
  } else if (currentMode == FREQ_READ) {
    display.print("In: ");
    display.setTextSize(2);
    float measuredFreq = measureInputFreq();
    if (measuredFreq > 1) {
      char inFreqStr[10];
      snprintf(inFreqStr, sizeof(inFreqStr), "%.1f Hz", measuredFreq);
      // Center the frequency string (approximate, adjust x based on text length if needed)
      display.setCursor(32, 24);
      display.print(inFreqStr);
    } else {
      // Center "No Signal" for neat alignment
      display.setCursor(20, 24);
      display.print("No Signal");
    }
  } else {
    display.print("Freq: ");
    display.setTextSize(2);
    char freqStr[10];
    snprintf(freqStr, sizeof(freqStr), "%d Hz", freq);
    // Center the frequency string (approximate, adjust x based on text length if needed)
    display.setCursor(36, 24);
    display.print(freqStr);
  }

  // Status (Bottom section)
  display.setTextSize(1);
  display.setCursor(0, 48);
  display.print("Step: ");
  display.print(stepSize);
  display.print(" Hz  Out: ");
  display.print(outputEnabled ? "ON" : "OFF");
  display.print("  Swp: ");
  display.print(sweepEnabled ? "ON" : "OFF");

  display.display();
}

float measureInputFreq() {
  unsigned long high = pulseIn(INPUT_PIN, HIGH, 50000);
  unsigned long low = pulseIn(INPUT_PIN, LOW, 50000);

  if (high == 0 || low == 0 || high > 45000 || low > 45000 || high < 5 || low < 5) {
    return 0;
  }
  unsigned long period = high + low;
  float freq = 1000000.0 / period;
  if (freq < 1 || freq > 10000) {
    return 0;
  }
  return freq;
}