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

// Screen dimensions
#define SCREEN_WIDTH 128
#define SCREEN_HEIGHT 64

// OLED display setup
#define OLED_RESET -1
Adafruit_SSD1306 display(SCREEN_WIDTH, SCREEN_HEIGHT, &Wire, OLED_RESET);

// Button pins
#define UP_BUTTON_PIN     PA0
#define DOWN_BUTTON_PIN   PA1
#define OK_BUTTON_PIN     PA2
#define CANCEL_BUTTON_PIN PA3

// IBT_2 Motor Driver pins
#define RPWM_PIN PB0
#define LPWM_PIN PB1

// Motor control variables
int motorSpeed = 0;
bool motorDirection = true; // true for forward, false for reverse
bool motorEnabled = false;

void setup() {
  Serial.begin(115200);

  // Initialize OLED display
  if(!display.begin(SSD1306_SWITCHCAPVCC, 0x3C)) {
    Serial.println(F("SSD1306 allocation failed"));
    for(;;);
  }

  // Button pin setup with internal pull-up resistors
  pinMode(UP_BUTTON_PIN, INPUT_PULLUP);
  pinMode(DOWN_BUTTON_PIN, INPUT_PULLUP);
  pinMode(OK_BUTTON_PIN, INPUT_PULLUP);
  pinMode(CANCEL_BUTTON_PIN, INPUT_PULLUP);

  // Motor driver pin setup
  pinMode(RPWM_PIN, OUTPUT);
  pinMode(LPWM_PIN, OUTPUT);

  updateDisplay();
}

void loop() {
  bool needsUpdate = false;

  // Button handling
  if (digitalRead(UP_BUTTON_PIN) == LOW) {
    Serial.println("Up button pressed");
    motorSpeed += 5;
    if (motorSpeed > 255) motorSpeed = 255;
    needsUpdate = true;
    delay(100); // Simple debounce
  }

  if (digitalRead(DOWN_BUTTON_PIN) == LOW) {
    Serial.println("Down button pressed");
    motorSpeed -= 5;
    if (motorSpeed < 0) motorSpeed = 0;
    needsUpdate = true;
    delay(100); // Simple debounce
  }

  if (digitalRead(OK_BUTTON_PIN) == LOW) {
    Serial.println("OK button pressed");
    motorEnabled = !motorEnabled;
    needsUpdate = true;
    delay(200); // Simple debounce
  }

  if (digitalRead(CANCEL_BUTTON_PIN) == LOW) {
    Serial.println("Cancel button pressed");
    motorDirection = !motorDirection;
    needsUpdate = true;
    delay(200); // Simple debounce
  }

  // Only update the display if a button was pressed
  if (needsUpdate) {
    updateDisplay();
  }

  // Update motor driver
  if (motorEnabled) {
    if (motorDirection) {
      analogWrite(RPWM_PIN, motorSpeed);
      analogWrite(LPWM_PIN, 0);
    } else {
      analogWrite(RPWM_PIN, 0);
      analogWrite(LPWM_PIN, motorSpeed);
    }
  } else {
    analogWrite(RPWM_PIN, 0);
    analogWrite(LPWM_PIN, 0);
  }
}

void updateDisplay() {
  display.clearDisplay();
  display.setTextColor(SSD1306_WHITE);

  // --- TOP YELLOW SECTION (Labels) ---
  display.setTextSize(1);
  display.setCursor(5, 4);
  display.print("SPEED");
  display.setCursor(50, 4);
  display.print("DIR");
  display.setCursor(95, 4);
  display.print("STATE");
  display.drawLine(0, 15, 127, 15, SSD1306_WHITE); // Separator line

  // --- BOTTOM BLUE SECTION (Values) ---
  display.setTextSize(2);
  
  // Display Speed Value with padding for alignment
  display.setCursor(0, 25);
  if (motorSpeed < 10) {
    display.print("  ");
  } else if (motorSpeed < 100) {
    display.print(" ");
  }
  display.print(motorSpeed);
  
  // Display Direction Value
  display.setCursor(45, 25);
  display.print(motorDirection ? "FWD" : "REV");
  
  // Display State Value
  display.setCursor(95, 25);
  display.print(motorEnabled ? "ON" : "OFF");

  display.display();
}