#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;

enum MenuState { MAIN_MENU, SET_TOTAL, SET_AMOUNT, BACK };
MenuState menuState = MAIN_MENU;

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); // Draw white text
  display.display();

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

  // Show main screen
  showMainScreen();
}

void loop() {
  currentWeight = scale.get_units();  // Get the current weight
  if (!inMenu) {
    updateMainScreen();
  }

  // Handle long press for menu
  if (digitalRead(BUTTON_POUR) == LOW) {
    if (buttonPressTime == 0) {
      buttonPressTime = millis();
    } else if (millis() - buttonPressTime > LONG_PRESS_DURATION) {
      showMenu();
    }
  } else {
    buttonPressTime = 0;
  }

  // Pour coffee
  if (digitalRead(BUTTON_POUR) == LOW && !inMenu) {
    pourCoffee();
  }

  // Calibrate scale
  if (digitalRead(BUTTON_CALIBRATE) == LOW) {
    calibrateScale();
  }

  // Handle menu navigation with encoder
  if (inMenu) {
    handleMenu();
  }
}

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() {
  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();
}

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 showMenu() {
  inMenu = true;
  menuState = SET_TOTAL;
  display.clearDisplay();
  display.setCursor(0, 0);
  display.print("1. Set Total");
  display.setCursor(0, 10);
  display.print("2. Set Amount");
  display.setCursor(0, 20);
  display.print("3. Back");
  display.display();
}

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

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