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

// Pin Definitions
#define SCREEN_WIDTH 128
#define SCREEN_HEIGHT 64
#define OLED_RESET -1
Adafruit_SSD1306 display(SCREEN_WIDTH, SCREEN_HEIGHT, &Wire, OLED_RESET);

const int relayPin = 5;
const int buzzerPin = 4;

const int setButtonPin = A3;
const int upButtonPin = A2;
const int downButtonPin = A1;
const int startButtonPin = A0;

// Constants for button debounce
const unsigned long debounceDelay = 50;

// Constants for timer control
const int timerIncrement = 30;

int timerSeconds = 0;
bool setupMode = false;
bool timerRunning = false;
unsigned long lastButtonCheckTime = 0;

void setup() {
  pinMode(relayPin, OUTPUT);
  pinMode(buzzerPin, OUTPUT);

  pinMode(setButtonPin, INPUT_PULLUP);
  pinMode(upButtonPin, INPUT_PULLUP);
  pinMode(downButtonPin, INPUT_PULLUP);
  pinMode(startButtonPin, INPUT_PULLUP);

  display.begin(SSD1306_SWITCHCAPVCC, 0x3C);

  displaySplashScreen(); // Display splash screen on startup
}

void loop() {
  handleButtons();

  if (setupMode) {
    displaySetupMode();
  } else if (timerRunning) {
    updateTimerDisplay();
  } else {
    displayMainScreen();
  }
}

void handleButtons() {
  unsigned long currentMillis = millis();

  if (digitalRead(setButtonPin) == LOW && currentMillis - lastButtonCheckTime >= debounceDelay) {
    setupMode = !setupMode;
    lastButtonCheckTime = currentMillis;
    if (setupMode) {
      displaySetupMode();
    } else {
      displayMainScreen();
    }
  }

  if (setupMode) {
    adjustTimer();
  }
}

void displaySplashScreen() {
  display.clearDisplay();

  // Implementing a more sophisticated splash screen
  display.setTextSize(2);
  display.setTextColor(SSD1306_WHITE);
  display.setCursor(30, 0);
  display.print("Welcome!");

  display.setTextSize(1);
  display.setCursor(20, 25);
  display.print("to the UV Cure Station");

  display.display();
  delay(3000); // Display the splash screen for 3 seconds
  display.clearDisplay();
}

void displayMainScreen() {
  display.clearDisplay();

  // Implementing a more sophisticated main screen
  display.setTextSize(2);
  display.setTextColor(SSD1306_WHITE);
  display.setCursor(10, 0);
  display.print("UV Cure Station");

  display.setTextSize(1);
  display.setCursor(10, 25);
  display.print("Press Set to enter setup");

  display.setTextSize(1);
  display.setCursor(10, 40);
  display.print("Press Start to begin curing");

  display.display();
}

void displaySetupMode() {
  display.clearDisplay();

  // Draw square around
  display.drawRect(25, 30, 95, 15, SSD1306_WHITE);

  display.setTextSize(2);
  display.setTextColor(SSD1306_WHITE);
  display.setCursor(30, 0);
  display.print("Setup");
  display.drawLine(0, 25, display.width() - 1, 25, SSD1306_WHITE);
  display.setTextSize(1);
  display.setCursor(30, 50);
  display.print("Press Start");
  display.setTextSize(1);
  display.setCursor(30, 35);
  display.print("Seconds: " + String(timerSeconds));
  display.display();
}

void adjustTimer() {
  // Update timerSeconds when up or down buttons are pressed
  if (digitalRead(upButtonPin) == LOW) {
    timerSeconds += timerIncrement;
    lastButtonCheckTime = millis();
    displaySetupMode();
  }

  if (digitalRead(downButtonPin) == LOW) {
    timerSeconds = max(0, timerSeconds - timerIncrement);
    lastButtonCheckTime = millis();
    displaySetupMode();
  }

  if (digitalRead(startButtonPin) == LOW) {
    // Start button pressed, exit setup mode and start timer
    setupMode = false;
    timerRunning = true;
    digitalWrite(relayPin, HIGH); // Turn on the relay
    delay(200); // Debounce

    digitalWrite(buzzerPin, HIGH);
    delay(100); // Buzzer sound for 0.1 seconds
    digitalWrite(buzzerPin, LOW);
  }
}

void updateTimerDisplay() {
  display.clearDisplay();
  display.setTextSize(4);
  display.setTextColor(SSD1306_WHITE);
  display.setCursor(20, 15);
  display.print(timerSeconds / 60); // Minutes
  display.print(":");
  if (timerSeconds % 60 < 10) {
    display.print("0");
  }
  display.print(timerSeconds % 60); // Seconds
  display.display();

  // Update the timer every second
  delay(1000);

  if (--timerSeconds < 0) {
    // Timer done, turn off relay, activate buzzer, and reset
    digitalWrite(relayPin, LOW);
    digitalWrite(buzzerPin, HIGH);
    delay(2000); // Buzzer sound for 2 seconds
    digitalWrite(buzzerPin, LOW);
    timerRunning = false;
    timerSeconds = 0;

    setupMode = true; // Set up mode after the timer is done
  }
}