#include <Wire.h>
#include <Adafruit_GFX.h>
#include <Adafruit_SH110X.h>
#include <Keypad.h>

// OLED Display Setup
#define SCREEN_WIDTH 128
#define SCREEN_HEIGHT 64
Adafruit_SH1107 display = Adafruit_SH1107(SCREEN_WIDTH, SCREEN_HEIGHT, &Wire, 0x3C); // Adjust I2C address if needed

// Keypad Setup
const byte ROWS = 4;
const byte COLS = 4;
char keys[ROWS][COLS] = {
  {'1', '2', '3', 'A'},
  {'4', '5', '6', 'B'},
  {'7', '8', '9', 'C'},
  {'*', '0', '#', 'D'}
};
byte rowPins[ROWS] = {9, 8, 7, 6}; // Connect to the row pins of the keypad
byte colPins[COLS] = {5, 4, 3, 2}; // Connect to the column pins of the keypad
Keypad keypad = Keypad(makeKeymap(keys), rowPins, colPins, ROWS, COLS);

// Relay Setup
#define RELAY_PIN 10
bool relayActivated = false;

// Timer Variables
unsigned long timerDuration = 0; // Timer duration in seconds
unsigned long startTime = 0;
bool timerActive = false;
bool stopRelayActivation = false; // Flag to disable relay activation if D is pressed

void setup() {
  // Initialize relay
  pinMode(RELAY_PIN, OUTPUT);
  digitalWrite(RELAY_PIN, LOW);

  // Initialize OLED display
  if (!display.begin(0x3C, true)) { // Ensure the I2C address is correct
    triggerRelay(); // Trigger relay if OLED is not detected
    while (true);   // Halt program
  }
  display.clearDisplay();
  display.display();
  display.setTextSize(1);
  display.setTextColor(SH110X_WHITE);

  // Initial message
  displayMessage("Set timer using keypad\nPress '#' to start");
}

void loop() {
  // Check connections during the countdown
  if (timerActive && !checkConnections()) {
    triggerRelay(); // Trigger relay if any component is disconnected
    displayMessage("Connection lost!");
    timerActive = false;
    return;
  }

  char key = keypad.getKey(); // Read key input

  // Key handling when the timer is NOT active
  if (!timerActive) {
    if (key) {
      if (key >= '0' && key <= '9') {
        timerDuration = timerDuration * 10 + (key - '0'); // Append digit
        displayTimer(timerDuration);
      } else if (key == '#') { // Start timer
        if (timerDuration > 0) {
          startTime = millis();
          timerActive = true;
          stopRelayActivation = false; // Ensure relay can activate
          displayMessage("Timer started!");
        } else {
          displayMessage("Set a valid time!");
        }
      } else if (key == '*') { // Clear input
        timerDuration = 0;
        displayMessage("Timer cleared!");
      } else if (key == 'A') { // Increase by 10 seconds
        timerDuration += 10;
        displayTimer(timerDuration);
      } else if (key == 'B') { // Decrease by 10 seconds
        if (timerDuration >= 10) timerDuration -= 10;
        else timerDuration = 0;
        displayTimer(timerDuration);
      }
    }
  } else { // Timer is active
    if (key == 'D') { // Stop and prevent relay activation
      timerActive = false;
      stopRelayActivation = true;
      timerDuration = 0;
      deactivateRelay();
      displayMessage("Timer stopped!");
    } else if (key == 'A') { // Increase by 10 seconds
      timerDuration += 10;
      displayTimer(timerDuration - elapsedSeconds());
    } else if (key == 'B') { // Decrease by 10 seconds
      if (timerDuration > 10) timerDuration -= 10;
      else timerDuration = 0;
      displayTimer(timerDuration - elapsedSeconds());
    }

    // Handle countdown logic
    unsigned long elapsed = elapsedSeconds();
    if (elapsed >= timerDuration) {
      if (!stopRelayActivation) { // Only activate relay if not stopped
        activateRelay();
        displayMessage("Timer done!");
        delay(5000); // Keep relay active for 5 seconds
        deactivateRelay();
      } else {
        displayMessage("Timer canceled!");
      }
      timerActive = false;
      timerDuration = 0;
      displayMessage("Set timer using keypad\nPress '#' to start");
    } else {
      displayTimer(timerDuration - elapsed);
    }
  }
}

// Function to calculate elapsed time in seconds
unsigned long elapsedSeconds() {
  return (millis() - startTime) / 1000;
}

// Function to convert seconds into HH:MM:SS format and display on the OLED
void displayTimer(unsigned long seconds) {
  unsigned long hours = seconds / 3600;
  unsigned long minutes = (seconds % 3600) / 60;
  unsigned long remainingSeconds = seconds % 60;

  display.clearDisplay();
  display.setCursor(0, 0);
  display.print("Timer: ");
  if (hours < 10) display.print("0");
  display.print(hours);
  display.print(":");
  if (minutes < 10) display.print("0");
  display.print(minutes);
  display.print(":");
  if (remainingSeconds < 10) display.print("0");
  display.print(remainingSeconds);
  display.display();
}

// Function to display a custom message on the OLED
void displayMessage(const char* message) {
  display.clearDisplay();
  display.setCursor(0, 0);
  display.println(message);
  display.display();
}

// Function to activate the relay
void activateRelay() {
  digitalWrite(RELAY_PIN, HIGH);
  relayActivated = true;
}

// Function to deactivate the relay
void deactivateRelay() {
  digitalWrite(RELAY_PIN, LOW);
  relayActivated = false;
}

// Function to trigger the relay in case of disconnection
void triggerRelay() {
  digitalWrite(RELAY_PIN, HIGH);
  relayActivated = true;
}

// Function to check if all components are connected
bool checkConnections() {
  // Check OLED
  for (int i = 0; i < 3; i++) { // Retry up to 3 times
    Wire.beginTransmission(0x3C);
    if (Wire.endTransmission() == 0) return true; // Connection is good
    delay(10);
  }

  return false; // If all retries fail, assume disconnection
}