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

// Function prototypes
int measureDistance();
bool isLargeBag(int distance);
bool isSmallBag(int distance);

// Pin Definitions
const int trigPin = 6;           // TRIG pin of ultrasonic sensor
const int echoPin = 7;           // ECHO pin of ultrasonic sensor
const int actuatorPin = 3;       // Actuator connected to digital pin 3
const int emergencyStopPin = 4;  // Emergency stop button connected to pin 4
const int buzzerPin = 5;         // Optional buzzer for system halt alert

// Constants
const int detectionDistance = 100;            // Max distance threshold for bag detection in cm
const int largeBagMinDistance = 30;           // Min distance for large bags
const int largeBagMaxDistance = 35;           // Max distance for large bags
const int smallBagMinDistance = 15;           // Min distance for small bags
const int smallBagMaxDistance = 20;           // Max distance for small bags
const int maxLargeBagsAllowed = 2;            // Maximum number of large bags allowed
const int maxSmallBagsAllowed = 2;            // Maximum number of small bags allowed

// Variables
bool emergencyStop = false;
unsigned long lastBagTime = 0;
int largeBagCount = 0;
int smallBagCount = 0;
unsigned long startTime;

void setup() {
  // Initialize pins
  pinMode(trigPin, OUTPUT);
  pinMode(echoPin, INPUT);
  pinMode(actuatorPin, OUTPUT);
  pinMode(emergencyStopPin, INPUT_PULLUP); // Using internal pull-up resistor
  pinMode(buzzerPin, OUTPUT);

  // Initialize the OLED display
  if (!display.begin(SSD1306_SWITCHCAPVCC, 0x3C)) {
    Serial.println(F("SSD1306 allocation failed"));
    for (;;);
  }
  display.display();
  delay(2000); // Pause for 2 seconds
  display.clearDisplay();
  display.setTextSize(1);      // Normal 1:1 pixel scale
  display.setTextColor(SSD1306_WHITE); // Draw white text
  display.setCursor(0, 0);     // Start at top-left corner
  display.display();

  // Initialize serial communication for debugging
  Serial.begin(9600);

  // Start the timer
  startTime = millis();
}

void loop() {
  // Read ultrasonic sensor and emergency stop button
  int distance = measureDistance();
  int stopButtonState = digitalRead(emergencyStopPin);

  // Check emergency stop
  if (stopButtonState == LOW) {
    emergencyStop = true;
  }

  // If emergency stop is not triggered, proceed with normal operation
  if (!emergencyStop) {
    unsigned long currentTime = millis();
    unsigned long elapsedLargeBagTime = currentTime - lastBagTime;
    unsigned long elapsedSmallBagTime = currentTime - lastBagTime;

    if (distance <= detectionDistance) {
      // Bag detected, activate actuator
      digitalWrite(actuatorPin, HIGH);
      delay(500); // Hold actuator for 0.5 seconds (adjust as needed)
      digitalWrite(actuatorPin, LOW);

      // Check bag type and count
      if (isLargeBag(distance)) {
        largeBagCount++;
        lastBagTime = currentTime; // Update last large bag time
      } else if (isSmallBag(distance)) {
        smallBagCount++;
        lastBagTime = currentTime; // Update last small bag time
      }
    }

    // Reset large bag count after 6 seconds if more than 0
    if (largeBagCount > 0 && elapsedLargeBagTime >= 6000) { // 6 seconds
      largeBagCount = 0;
      lastBagTime = currentTime; // Update last bag time
    }

    // Reset small bag count after 10 seconds if more than 0
    if (smallBagCount > 0 && elapsedSmallBagTime >= 10000) { // 10 seconds
      smallBagCount = 0;
      lastBagTime = currentTime; // Update last bag time
    }

    // Check if the number of large bags or small bags exceeds the maximum allowed
    if (largeBagCount > maxLargeBagsAllowed || smallBagCount > maxSmallBagsAllowed) {
      emergencyStop = true; // Trigger emergency stop
    }

    // Display status on main screen
    display.clearDisplay();
    display.setCursor(0, 0);
    display.print("Large Bags: ");
    display.print(largeBagCount);
    display.setCursor(0, 10);
    display.print("Small Bags: ");
    display.print(smallBagCount);
    display.setCursor(0, 20);
    display.print("Elapsed Time: ");
    display.print((currentTime - startTime) / 1000); // Display time in seconds

    // Display bag allowance status on secondary screen
    display.setCursor(0, 40);
    if (largeBagCount <= 1) {
      display.print("Large Bags Allowed");
    } else if (largeBagCount == 2) {
      display.print("No More Large Bags");
    }
    display.setCursor(0, 50);
    if (smallBagCount <= 1) {
      display.print("Small Bags Allowed");
    } else if (smallBagCount == 2) {
      display.print("No More Small Bags");
    }

    display.display();

  } else {
    // If emergency stop is triggered, deactivate everything
    digitalWrite(actuatorPin, LOW);
    digitalWrite(buzzerPin, HIGH); // Sound buzzer for alert
    display.clearDisplay();
    display.setCursor(0, 0);
    display.print("System Halted");
    display.setCursor(0, 10);
    display.print("Emergency Stop");
    display.display();
  }

  delay(100); // Small delay to debounce sensor and button readings
}

int measureDistance() {
  // Clear the TRIG pin
  digitalWrite(trigPin, LOW);
  delayMicroseconds(2);

  // Set the TRIG pin HIGH for 10 microseconds
  digitalWrite(trigPin, HIGH);
  delayMicroseconds(10);
  digitalWrite(trigPin, LOW);

  // Read the ECHO pin
  long duration = pulseIn(echoPin, HIGH);

  // Calculate the distance in centimeters
  int distance = duration * 0.034 / 2;

  return distance;
}

bool isLargeBag(int distance) {
  return distance >= largeBagMinDistance && distance <= largeBagMaxDistance;
}

bool isSmallBag(int distance) {
  return distance >= smallBagMinDistance && distance <= smallBagMaxDistance;
}