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

// ---------------- OLED ----------------
#define SCREEN_WIDTH 128
#define SCREEN_HEIGHT 64

Adafruit_SSD1306 oled(SCREEN_WIDTH, SCREEN_HEIGHT, &Wire, -1);

// ---------------- PINS ----------------
const int pirPin   = 16;
const int trigPin  = 3;
const int echoPin  = 2;
const int potPin   = 26;
const int servoPin = 15;

// ---------------- SERVO ----------------
Servo servoMotor;

// ---------------- SYSTEM ----------------
bool system_on = false;

// ---------------- SERVO FUNCTION ----------------
void setAngle(int angle) {
  servoMotor.write(angle);
}

// ---------------- ULTRASONIC FUNCTION ----------------
float measureDistance() {

  digitalWrite(trigPin, LOW);
  delayMicroseconds(2);

  digitalWrite(trigPin, HIGH);
  delayMicroseconds(10);

  digitalWrite(trigPin, LOW);

  long duration = pulseIn(echoPin, HIGH, 30000);

  // no reading protection
  if (duration == 0) {
    return 400;
  }

  float distance = duration * 0.0343 / 2.0;

  // clamp
  if (distance > 400) distance = 400;
  if (distance < 0) distance = 0;

  return distance;
}

void setup() {

  // ---------------- OLED I2C ----------------
  Wire.setSDA(0); // GP0
  Wire.setSCL(1); // GP1
  Wire.begin();

  // ---------------- OLED START ----------------
  if (!oled.begin(SSD1306_SWITCHCAPVCC, 0x3C)) {

    while (1);
  }

  oled.clearDisplay();
  oled.setTextSize(1);
  oled.setTextColor(WHITE);

  oled.setCursor(20, 25);
  oled.println("SYSTEM READY");

  oled.display();

  delay(2000);

  // ---------------- PIN MODES ----------------
  pinMode(pirPin, INPUT_PULLDOWN);

  pinMode(trigPin, OUTPUT);
  pinMode(echoPin, INPUT);

  // ---------------- SERVO ----------------
  servoMotor.attach(servoPin);

  setAngle(90);
}

// ---------------- MAIN LOOP ----------------
void loop() {

  // ---------------- PIR START ----------------
  if (!system_on) {

    if (digitalRead(pirPin) == HIGH) {

      system_on = true;

      delay(300);
    }
  }

  // ---------------- SYSTEM OFF ----------------
  if (!system_on) {

    oled.clearDisplay();

    oled.setCursor(20, 20);
    oled.println("SYSTEM OFF");

    oled.setCursor(0, 40);
    oled.println("Motion to Start");

    oled.display();

    delay(300);

    return;
  }

  // ---------------- SENSOR READINGS ----------------
  float distance = measureDistance();

  // Raspberry Pi Pico ADC = 0 - 4095
  int colorValue = analogRead(potPin);

  // ---------------- SIZE DETECTION ----------------
  String size;

  if (distance <= 120) {

    size = "Large";
  }
  else if (distance <= 250) {

    size = "Medium";
  }
  else {

    size = "Small";
  }

  // ---------------- COLOR DETECTION ----------------
  String color;

  if (colorValue < 1300) {

    color = "Dark";
  }
  else if (colorValue < 3000) {

    color = "Medium";
  }
  else {

    color = "Light";
  }

  // ---------------- SORTING LOGIC ----------------
  String result;

  if (size == "Small" && color == "Dark") {

    result = "Premium";

    setAngle(30);
  }
  else if (size == "Medium" && color == "Medium") {

    result = "Standard";

    setAngle(90);
  }
  else {

    result = "Reject";

    setAngle(150);
  }

  // ---------------- OLED DISPLAY ----------------
  oled.clearDisplay();

  oled.setCursor(0, 0);
  oled.println("Coffee Sorter");

  oled.setCursor(0, 15);
  oled.print("Dist: ");
  oled.print((int)distance);
  oled.println(" cm");

  oled.setCursor(0, 30);
  oled.print("Size: ");
  oled.println(size);

  oled.setCursor(0, 45);
  oled.print("Color: ");
  oled.println(color);

  oled.setCursor(0, 55);
  oled.println(result);

  oled.display();

  // ---------------- LOOP DELAY ----------------
  delay(500);
}