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

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

// Define the pin numbers and PWM channels for the two motors
#define MOTOR1_PWM_PIN 16
#define MOTOR2_PWM_PIN 17
#define PWM_CHANNEL_1 0
#define PWM_CHANNEL_2 1
#define PWM_FREQ 1000
#define PWM_RESOLUTION 8 // 8 bits: 0-255

#define ENVELOPE 0xAA
#define TIMEOUT 200 // milliseconds

unsigned long lastReceiveTime = 0;
uint8_t prevMotor1 = 0, prevMotor2 = 0;

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

  ledcSetup(PWM_CHANNEL_1, PWM_FREQ, PWM_RESOLUTION);
  ledcSetup(PWM_CHANNEL_2, PWM_FREQ, PWM_RESOLUTION);
  ledcAttachPin(MOTOR1_PWM_PIN, PWM_CHANNEL_1);
  ledcAttachPin(MOTOR2_PWM_PIN, PWM_CHANNEL_2);

  // OLED initialization
  if(!display.begin(SSD1306_SWITCHCAPVCC, 0x3C)) {
    Serial.println(F("SSD1306 allocation failed"));
    for(;;);
  }
  display.clearDisplay();
  display.setTextColor(SSD1306_WHITE);
  display.setTextSize(1);
  display.setCursor(0, 0);
  display.println("Init...");
  display.display();
}

void updateOLED(uint8_t m1, uint8_t m2) {
  display.clearDisplay();
  display.setTextSize(2);
  display.setCursor(0, 0);
  display.print("M1: ");
  display.println(m1);
  display.setCursor(0, 30);
  display.print("M2: ");
  display.println(m2);
  display.display();
}

void loop() {
  static enum { WAIT_HEADER, WAIT_M1, WAIT_M2 } state = WAIT_HEADER;
  static uint8_t m1_val, m2_val;

  while (Serial.available()) {
    uint8_t b = Serial.read();
    if (state == WAIT_HEADER) {
      if (b == ENVELOPE) state = WAIT_M1;
    } else if (state == WAIT_M1) {
      m1_val = b;
      state = WAIT_M2;
    } else if (state == WAIT_M2) {
      m2_val = b;
      // Only update if changed
      if (m1_val != prevMotor1) {
        ledcWrite(PWM_CHANNEL_1, m1_val);
        prevMotor1 = m1_val;
      }
      if (m2_val != prevMotor2) {
        ledcWrite(PWM_CHANNEL_2, m2_val);
        prevMotor2 = m2_val;
      }
      updateOLED(prevMotor1, prevMotor2);
      lastReceiveTime = millis();
      state = WAIT_HEADER;
    }
  }

  // Safety: Stop both motors if no update in TIMEOUT ms
  if (millis() - lastReceiveTime > TIMEOUT) {
    if (prevMotor1 != 0) {
      ledcWrite(PWM_CHANNEL_1, 0);
      prevMotor1 = 0;
    }
    if (prevMotor2 != 0) {
      ledcWrite(PWM_CHANNEL_2, 0);
      prevMotor2 = 0;
    }
    updateOLED(0, 0);
  }
}