#include <Adafruit_GFX.h>esp32type:All
#include <Adafruit_SSD1306.h>

// === BUTTON CONFIGURATION ===
const int buttonPinA = D9;
const int buttonPinB = D10;

// === LED CONFIGURATION ===
const int ledPin0 = D0;
const int ledPin1 = D1;
const int ledPin2 = D2;

// === MOTOR CONFIGURATION ===
const int motorA_IN1 = D3;
const int motorA_IN2 = D6;
const int motorB_IN1 = D8;
const int motorB_IN2 = D7;

// === VARIABLES ===
bool ledState0 = LOW;
bool ledState1 = LOW;
bool ledState2 = LOW;

unsigned long pressStartA = 0;
unsigned long pressStartB = 0;
bool wasPressedA = false;
bool wasPressedB = false;

const unsigned long debounceDelay = 50;   // ms
const unsigned long longPressTime = 1000; // ms
const unsigned long comboGraceTime = 120; // ms

int mode = 0;  // 0–5

// === MOTOR STATE VARIABLES ===
int speedA = 0;
int speedB = 0;
bool dirA = true; // true = forward
bool dirB = true;
bool motorAon = false;
bool motorBon = false;

// === OLED CONFIG ===
#define SCREEN_WIDTH 128
#define SCREEN_HEIGHT 32
#define OLED_RESET -1
Adafruit_SSD1306 display(SCREEN_WIDTH, SCREEN_HEIGHT, &Wire, OLED_RESET);

void setup() {
  pinMode(buttonPinA, INPUT);
  pinMode(buttonPinB, INPUT);

  pinMode(ledPin0, OUTPUT);
  pinMode(ledPin1, OUTPUT);
  pinMode(ledPin2, OUTPUT);

  pinMode(motorA_IN1, OUTPUT);
  pinMode(motorA_IN2, OUTPUT);
  pinMode(motorB_IN1, OUTPUT);
  pinMode(motorB_IN2, OUTPUT);

  analogWriteResolution(8); // 0–255 PWM

  Serial.begin(9600);
  while (!Serial);
  Serial.println("System ready. Mode: 0");
  
  // ... your existing setup code ...

  if(!display.begin(SSD1306_SWITCHCAPVCC, 0x3C)) {
    Serial.println("OLED not found!");
    while(1); // stop if OLED not found
  }
  display.clearDisplay();
  display.setTextSize(1);
  display.setTextColor(SSD1306_WHITE);
  display.display();
  Serial.println("display ready");
  
}

// --- Utility function ---
void setMotor(int in1, int in2, int duty, bool direction, bool enabled) {
  if (!enabled || duty == 0) {
    digitalWrite(in1, LOW);
    digitalWrite(in2, LOW);
    return;
  }

  int pwmValue = map(duty, 0, 100, 0, 255);
  if (direction) {
    analogWrite(in1, pwmValue);
    digitalWrite(in2, LOW);
  } else {
    digitalWrite(in1, LOW);
    analogWrite(in2, pwmValue);
  }
}

// --- Debounce + Press Logic ---
void loop() {
  Serial.println("loop");
  static int lastA = LOW, lastB = LOW;
  int stateA = digitalRead(buttonPinA);
  int stateB = digitalRead(buttonPinB);
  unsigned long now = millis();

  // Debounce
  if (stateA != lastA) {
    delay(debounceDelay);
    stateA = digitalRead(buttonPinA);
  }
  if (stateB != lastB) {
    delay(debounceDelay);
    stateB = digitalRead(buttonPinB);
  }

  // Press starts
  if (stateA == HIGH && !wasPressedA) { pressStartA = now; wasPressedA = true; }
  if (stateB == HIGH && !wasPressedB) { pressStartB = now; wasPressedB = true; }

  // Press releases
  if (stateA == LOW && wasPressedA) {
    wasPressedA = false;
    unsigned long pressDuration = now - pressStartA;

    if (wasPressedB || (abs((long)pressStartA - (long)pressStartB) < comboGraceTime)) {
      // handled by both logic
    } else if (pressDuration < longPressTime) {
      ledState0 = !ledState0;
      digitalWrite(ledPin0, ledState0);
      handleButton('A');
    }
  }

  if (stateB == LOW && wasPressedB) {
    wasPressedB = false;
    unsigned long pressDuration = now - pressStartB;

    if (wasPressedA || (abs((long)pressStartA - (long)pressStartB) < comboGraceTime)) {
      // handled by both logic
    } else if (pressDuration < longPressTime) {
      ledState1 = !ledState1;
      digitalWrite(ledPin1, ledState1);
      handleButton('B');
    }
  }

  // Both pressed logic
  static bool bothActive = false;
  static unsigned long bothStart = 0;

  if (stateA == HIGH && stateB == HIGH && !bothActive) {
    bothActive = true;
    bothStart = now;
  }

  if (bothActive && stateA == LOW && stateB == LOW) {
    bothActive = false;
    unsigned long pressTime = now - bothStart;

    if (pressTime >= longPressTime) {
      // Long press both → change mode
      mode = (mode + 1) % 6;
      ledState2 = ledState2= ledState0=HIGH;
      digitalWrite(ledPin0, ledState0);
      digitalWrite(ledPin1, ledState1);
      digitalWrite(ledPin2, ledState2);
      Serial.print("Mode changed to ");
      Serial.println(mode);
    } else {
      // Short press both → special function
      ledState2 = !ledState2;
      digitalWrite(ledPin2, ledState2);
      handleBoth();
    }
  }

  // Update motors continuously
  setMotor(motorA_IN1, motorA_IN2, speedA, dirA, motorAon);
  setMotor(motorB_IN1, motorB_IN2, speedB, dirB, motorBon);

  lastA = stateA;
  lastB = stateB;
  Serial.println("display update");
  updateDisplay();
}

// === MODE-SPECIFIC LOGIC ===
void handleButton(char btn) {
  switch (mode) {
    case 0: // Toggle 50% on/off
      if (btn == 'A') {
        motorAon = !motorAon;
        speedA = motorAon ? 50 : 0;
        Serial.print("Motor A ");
        Serial.println(motorAon ? "ON (50%)" : "OFF");
      } else if (btn == 'B') {
        motorBon = !motorBon;
        speedB = motorBon ? 50 : 0;
        Serial.print("Motor B ");
        Serial.println(motorBon ? "ON (50%)" : "OFF");
      }
      break;

    case 1: // Control Motor A
      if (btn == 'A') speedA = constrain(speedA + 5, 0, 100);
      if (btn == 'B') speedA = constrain(speedA - 10, 0, 100);
      motorAon = (speedA > 0);
      Serial.print("Motor A speed: "); Serial.println(speedA);
      break;

    case 2: // Control Motor B
      if (btn == 'A') speedB = constrain(speedB + 5, 0, 100);
      if (btn == 'B') speedB = constrain(speedB - 10, 0, 100);
      motorBon = (speedB > 0);
      Serial.print("Motor B speed: "); Serial.println(speedB);
      break;

    case 3: // Control A, keep B state
      if (btn == 'A') speedA = constrain(speedA + 5, 0, 100);
      if (btn == 'B') speedA = constrain(speedA - 10, 0, 100);
      motorAon = (speedA > 0);
      Serial.print("Motor A speed: "); Serial.println(speedA);
      break;

    case 4: // Control B, keep A state
      if (btn == 'A') speedB = constrain(speedB + 5, 0, 100);
      if (btn == 'B') speedB = constrain(speedB - 10, 0, 100);
      motorBon = (speedB > 0);
      Serial.print("Motor B speed: "); Serial.println(speedB); 
      break;

    case 5: // Everything off
      motorAon = motorBon = false;
      speedA = speedB = 0;
      Serial.println("All motors OFF");
      break;
  }
}

void handleBoth() {
  switch (mode) {
    case 0: // Both pressed → change global speed
      static int globalSpeed = 10;
      globalSpeed += 10;
      if (globalSpeed > 100) globalSpeed = 10;
      if (motorAon) speedA = globalSpeed;
      if (motorBon) speedB = globalSpeed;
      Serial.print("Speed changed to "); Serial.println(globalSpeed);
      break;

    case 1: // Reverse motor A
      dirA = !dirA;
      Serial.print("Motor A reversed. Now ");
      Serial.println(dirA ? "forward" : "reverse");
      break;

    case 2: // Reverse motor B
      dirB = !dirB;
      Serial.print("Motor B reversed. Now ");
      Serial.println(dirB ? "forward" : "reverse");
      break;

    case 3: // Reverse motor A only
      dirA = !dirA;
      Serial.print("Motor A reversed. Now ");
      Serial.println(dirA ? "forward" : "reverse");
      break;

    case 4: // Reverse motor B only
      dirB = !dirB;
      Serial.print("Motor B reversed. Now ");
      Serial.println(dirB ? "forward" : "reverse");
      break;
  }
}



void updateDisplay() {
      Serial.print("inside display ");
  display.clearDisplay();

  // Mode info
  display.setCursor(0,0);
  display.print("Mode "); display.print(mode); display.print(": ");
  switch(mode) {
    case 0: display.print("Dual Toggle"); break;
    case 1: display.print("Motor A Ctrl"); break;
    case 2: display.print("Motor B Ctrl"); break;
    case 3: display.print("Motor A+Hold B"); break;
    case 4: display.print("Motor B+Hold A"); break;
    case 5: display.print("All Off"); break;
  }

  // Motor info
  display.setCursor(0,10);
  display.print("A: "); display.print(speedA); display.print("% ");
  display.print(dirA ? "F" : "R");

  display.setCursor(70,10);
  display.print("B: "); display.print(speedB); display.print("% ");
  display.print(dirB ? "F" : "R");

  // Button states
  display.setCursor(0,22);
  display.print("BtnA:"); display.print(digitalRead(buttonPinA) ? "1 " : "0 ");
  display.print("BtnB:"); display.print(digitalRead(buttonPinB) ? "1" : "0");

  display.display();
  
      Serial.print("leaving display ");
}