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

// Define the size of the OLED screen
#define SCREEN_WIDTH 128
#define SCREEN_HEIGHT 64

// Define the OLED reset pin
#define OLED_RESET    -1

// Create an instance of the display
Adafruit_SSD1306 display(SCREEN_WIDTH, SCREEN_HEIGHT, &Wire, OLED_RESET);

// Define pins for rotary encoder and buttons
#define ENCODER_CLK 2
#define ENCODER_DT  3
#define SELECT_PIN  8
#define BACK_PIN    12

// Define pins for servos
#define SERVO_WATER_PIN 11
#define SERVO_ORANGE_PIN 10
#define SERVO_LEMON_PIN 6
#define SERVO_MINT_PIN 5

// Create Servo objects
Servo servoWater;
Servo servoOrange;
Servo servoLemon;
Servo servoMint;

// Define drinks
String boissons[] = {"Orange", "Citron", "Menthe", "Eau"};
int currentSelection = 0;
int newSelection = 0; // To track the new selection before confirmation
bool isConfirming = false;
bool isServingAnother = false; // Flag for serving another drink
int lastEncoderState = LOW;
unsigned long lastEncoderChange = 0;
const unsigned long debounceDelay = 50; // Debounce delay in milliseconds

void setup() {
  // Initialize Serial for debugging
  Serial.begin(9600);

  // Initialize the OLED display
  if (!display.begin(SSD1306_SWITCHCAPVCC, 0x3C)) { // Default I2C address for many OLED displays
    Serial.println(F("SSD1306 allocation failed"));
    for(;;); // Infinite loop if initialization fails
  }

  display.clearDisplay();
  display.setTextSize(1);      
  display.setTextColor(SSD1306_WHITE); 
  display.setCursor(0,0);     
  display.display();

  // Initialize pins
  pinMode(ENCODER_CLK, INPUT);
  pinMode(ENCODER_DT, INPUT);
  pinMode(SELECT_PIN, INPUT_PULLUP);
  pinMode(BACK_PIN, INPUT_PULLUP);

  // Attach servos to their pins
  servoWater.attach(SERVO_WATER_PIN);
  servoOrange.attach(SERVO_ORANGE_PIN);
  servoLemon.attach(SERVO_LEMON_PIN);
  servoMint.attach(SERVO_MINT_PIN);

  displayMenu();
}

void loop() {
  int encoderState = digitalRead(ENCODER_CLK);

  // Check if the encoder has been turned
  if (encoderState != lastEncoderState && (millis() - lastEncoderChange) > debounceDelay) {
    if (!isConfirming && !isServingAnother) { // Only update selection if not in confirmation or serving another mode
      if (digitalRead(ENCODER_DT) == encoderState) { // Invert the direction here
        newSelection++;
        if (newSelection >= 4) newSelection = 0; // Carrousel: Wrap around if end of list
      } else {
        newSelection--;
        if (newSelection < 0) newSelection = 3; // Carrousel: Wrap around if beginning of list
      }
      displayMenu(); // Update display to show new selection
      lastEncoderChange = millis(); // Update last change time
    }
  }
  lastEncoderState = encoderState;

  // Handle selection
  if (digitalRead(SELECT_PIN) == LOW) {
    if (isConfirming) {
      dispenseDrink(); // Dispense the selected drink
      isConfirming = false;
      isServingAnother = true; // Set flag to ask if user wants to serve another drink
      displayServeAnother(); // Ask if the user wants to serve another drink
      delay(500); // Debounce
    } else if (isServingAnother) {
      isServingAnother = false;
      displayMenu(); // Return to menu after serving another or choosing not to
      delay(500); // Debounce
    } else {
      // Enter confirmation mode
      currentSelection = newSelection; // Confirm the new selection
      isConfirming = true;
      displayConfirm();
      delay(500); // Debounce
    }
  }

  // Handle back
  if (digitalRead(BACK_PIN) == LOW) {
    if (isConfirming || isServingAnother) {
      isConfirming = false;
      isServingAnother = false;
      displayMenu(); // Return to menu
      delay(500); // Debounce
    }
  }
}

void displayMenu() {
  display.clearDisplay();
  display.setCursor(0, 0);
  for (int i = 0; i < 4; i++) {
    if (i == newSelection) {
      display.print("> "); // Indicate selection
    } else {
      display.print("  ");
    }
    display.println(boissons[i]);
  }
  display.display();
}

void displayConfirm() {
  display.clearDisplay();
  display.setCursor(0, 0);
  display.println("Confirmer?");
  display.println(boissons[currentSelection]);
  display.println("Select = OK");
  display.println("Back = Cancel");
  display.display();
}

void displayServeAnother() {
  display.clearDisplay();
  display.setCursor(0, 0);
  display.println("Voulez-vous");
  display.println("se resservir?");
  display.println("Select = Oui");
  display.println("Back = Non");
  display.display();
}

void dispenseDrink() {
  switch (currentSelection) {
    case 0: // Orange
      dispense(servoOrange, 30); // Example duration for orange
      dispense(servoWater, 10);   // Example duration for water
      break;
    case 1: // Citron
      dispense(servoLemon, 30);  // Example duration for lemon
      dispense(servoWater, 10);   // Example duration for water
      break;
    case 2: // Menthe
      dispense(servoMint, 30);   // Example duration for mint
      dispense(servoWater, 10);   // Example duration for water
      break;
    case 3: // Eau
      dispense(servoWater, 40);  // Larger amount of water for pure water drink
      break;
  }
}

void dispense(Servo& servo, int duration) {
  servo.write(90); // Open valve (adjust angle as needed)
  delay(duration * 10); // Wait for the specified duration
  servo.write(0); // Close valve (adjust angle as needed)
}