#include <Arduino.h>
#include <Wire.h>
#include <Adafruit_GFX.h>
#include <Adafruit_SSD1306.h>
#include <HX711.h>
#include <Servo.h>
#include <Encoder.h>

// Pins
#define ENCODER_CLK 2
#define ENCODER_DT 3
#define BUTTON_CALIBRATE 4
#define BUTTON_POUR 5
#define SERVO_PIN 6
#define HX711_DOUT 7
#define HX711_SCK 8


#define LONG_PRESS_DURATION 1000  // 1 second for long press
#define CALIBRATION_WEIGHT 100  // 100 grams calibration weight

#define SCREEN_WIDTH 128 // OLED display width, in pixels
#define SCREEN_HEIGHT 64 // OLED display height, in pixels

// Declaration for an SSD1306 display connected to I2C (SDA, SCL pins)
#define OLED_RESET     -1 // Reset pin # (or -1 if sharing Arduino reset pin)
Adafruit_SSD1306 display(SCREEN_WIDTH, SCREEN_HEIGHT, &Wire, OLED_RESET);

HX711 scale;
Servo gateServo;
Encoder encoder(ENCODER_CLK, ENCODER_DT);

long totalAmount = 1000;  // Total coffee in grams
long selectedAmount = 200;  // Amount to pour
long consumedAmount = 0;
long currentWeight = 0;
bool inMenu = false;
unsigned long buttonPressTime = 0;

// Creatign a menu list
enum MenuState { IDLE, MAIN_MENU, SET_TOTAL, SET_AMOUNT, BACK };
MenuState menuState = IDLE;

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

  // Initialize components
  pinMode(BUTTON_CALIBRATE, INPUT_PULLUP);
  pinMode(BUTTON_POUR, INPUT_PULLUP);
  scale.begin(HX711_DOUT, HX711_SCK);
  gateServo.attach(SERVO_PIN);
  gateServo.write(0);  // Close gate initially

  // SSD1306_SWITCHCAPVCC = generate display voltage from 3.3V internally
  if (!display.begin(SSD1306_SWITCHCAPVCC, 0x3C)) {
    Serial.println(F("SSD1306 allocation failed"));
    for (;;); // Don't proceed, loop forever
  }
  display.clearDisplay();
  display.setTextSize(1);      // Normal 1:1 pixel scale
  display.setTextColor(WHITE, BLACK); // Draw white text
  display.display();

  // Calibration (initially tare the scale)
  scale.tare();

  // Show main screen
  showMainScreen();
}

void loop() {
  // Check if user has pressed Calibrate button and we are not in menu
  if (digitalRead(BUTTON_CALIBRATE) == LOW) {
    if (menuState == IDLE)
      menuState = MAIN_MENU; // Set the menu state to main menu
    showMenu();
    delay(300);
  }
  
  // Handle menu navigation with encoder
  if (inMenu) {
    handleMenu();
  }

  // Pour coffee and check for scale calibration
  if (digitalRead(BUTTON_POUR) == LOW && !inMenu) {
    unsigned long cTime = millis(); // Note current button press time
    while (digitalRead(BUTTON_POUR) == LOW) {
      if (millis() - cTime > LONG_PRESS_DURATION) {
        calibrateScale(); // If user has pressed the button long enough, calibrate the scale
        break;
      }
    }
    pourCoffee();
  }
}
void showMessage(byte x, byte y, char *message) {
  display.setCursor(x, y);
  display.print(message);
  display.display();
}
void showMainScreen() {
  display.clearDisplay();
  display.setCursor(0, 0);
  display.print("Total: ");
  display.print(totalAmount);
  display.print("g");

  display.setCursor(0, 10);
  display.print("Selected: ");
  display.print(selectedAmount);
  display.print("g");

  display.setCursor(0, 20);
  display.print("Weight: ");
  display.print(currentWeight);
  display.print(" g  ");

  display.setCursor(0, 30);
  display.print("Consumed: ");
  display.print(consumedAmount);
  display.print(" g ");

  display.display();
}

void updateMainScreen() {
  display.setCursor(0, 20);
  display.print("Weight: ");
  display.print(currentWeight);
  display.print("g");
  display.display();
}

void pourCoffee() {
  scale.tare(); // clear scale
  long targetWeight = currentWeight + selectedAmount;
  gateServo.write(90);  // Open the gate
  while (scale.get_units() < targetWeight) {
    currentWeight = scale.get_units();
    updateMainScreen();
  }
  gateServo.write(0);  // Close the gate
  consumedAmount += selectedAmount;
  updateMainScreen();
  display.setCursor(0, 30);
  display.print("Consumed: ");
  display.print(consumedAmount);
  display.print(" g  ");
}

void calibrateScale() {
  display.clearDisplay();
  display.setCursor(0, 0);
  display.print("Place weight:");
  display.display();

  delay(2000);  // Wait for user to place weight

  scale.set_scale();
  scale.tare();  // Tare again to reset
  scale.set_scale(CALIBRATION_WEIGHT / scale.get_units());  // Adjust calibration factor
  showMainScreen();
}
void displayCursor(byte x, byte y) {
  display.setCursor(x, y); // Set curosor position
  display.print(">");
  display.display();
}
void showMenu() {
  inMenu = false;
  // menuState = SET_TOTAL;  // Set the menu state to set total
  // Reset encoder position value
  // encoder.write(0);
  display.clearDisplay();
  display.setCursor(3, 0);
  display.print("1. Set Total");
  display.setCursor(3, 10);
  display.print("2. Set Amount");
  display.setCursor(3, 20);
  display.print("3. Back");
  display.display();
  switch (menuState)
  {
    case MAIN_MENU:
      inMenu = true;
      delay(500);
      encoder.write(totalAmount); // Set the encoder value to total amount
      display.clearDisplay();
      menuState = SET_TOTAL;
      
      displayCursor(0, 0);
      break;
    case SET_TOTAL:
      inMenu = true;
      encoder.write(selectedAmount);  // Set the encoder value to selected amount
      delay(500);
      display.clearDisplay();
      menuState = SET_AMOUNT;
      displayCursor(0, 10);
      break;
    case SET_AMOUNT:
      menuState = BACK;
      displayCursor(0, 20);
      break;
    case BACK:
      menuState = IDLE;
      showMainScreen();
      break;
    default:
      break;
  }

}

void handleMenu() {
  long encoderPosition = encoder.read();

  switch (menuState) {
    case SET_TOTAL:
      display.setCursor(0, 0);
      display.print("Set Total: ");
      totalAmount = encoderPosition; // Adjust the total amount with encoder
      display.print(totalAmount);  // print total amount
      display.print("  g ");
      break;
    case SET_AMOUNT:
      display.setCursor(0, 10);
      display.print("Pour Amount: ");
      selectedAmount = encoderPosition;  // Adjust the amount with encoder
      display.print(selectedAmount);  // Adjust the amount
      display.print("  g ");      break;
    case BACK:
      inMenu = false;
      showMainScreen();
      break;
  }
  display.display();
}