#include <Adafruit_GFX.h>         // Core graphics library
#include <Adafruit_SSD1306.h>     // OLED library
#include <SPI.h>

// OLED display dimensions
#define SCREEN_WIDTH 128
#define SCREEN_HEIGHT 64

// SPI pins for OLED
#define OLED_MOSI 19  // SPI Data (SDA)
#define OLED_CLK 18   // SPI Clock (SCK)
#define OLED_DC 17    // Data/Command
#define OLED_CS 22    // Chip Select
#define OLED_RESET 16 // Reset

// Button pin definitions
#define ADD_TIME_BUTTON 3 // Add 15 minutes (5 seconds for demo) 
#define CLEAR_BUTTON 4    
#define START_BUTTON 7    
#define SWITCH_MODE_BUTTON 9 // Switch mode button (timer only vs. strict focus)
#define SPEAKER_PIN 15    // Speaker pin for beep
#define FSR_PIN A0        // Force Sensing Resistor pin

// Time variables
int totalTime = 0;      // Total countdown time in seconds
bool timerRunning = false;

// Force Sensing Resistor variables
int initialFSRReading = 0;
bool monitorFSR = false;
unsigned long lastBeepTime = 0;
unsigned long weightRemovedTime = 0;
bool weightRemoved = false;
bool firstBeepDone = false; // Track first beep after weight removal
const int beepInterval = 3000; 
const int weightThreshold = 40; // Minimum weight threshold to start timer in this mode

// Mode variables
enum Mode { TIMER_ONLY, STRICT_FOCUS };
Mode currentMode = TIMER_ONLY; 
bool modeSwitched = false;

// Timer tracking variables
unsigned long previousMillis = 0;

// OLED Display object
Adafruit_SSD1306 display(SCREEN_WIDTH, SCREEN_HEIGHT, &SPI, OLED_DC, OLED_RESET, OLED_CS);

// Function prototypes
void displayTime(int timeInSeconds);
void updateTimer();
void playBeep();
void monitorFSRChange();
void switchMode();
void displayMode();
bool checkWeight();

void setup() {
  // Initialize buttons as input with pull-up resistors
  pinMode(ADD_TIME_BUTTON, INPUT_PULLUP);
  pinMode(CLEAR_BUTTON, INPUT_PULLUP);
  pinMode(START_BUTTON, INPUT_PULLUP);
  pinMode(SWITCH_MODE_BUTTON, INPUT_PULLUP);

  // Initialize speaker pin
  pinMode(SPEAKER_PIN, OUTPUT);

  // Initialize OLED display
  if (!display.begin(SSD1306_SWITCHCAPVCC)) {
    while (true);
  }
  
  // Display "Hello" at startup
  display.clearDisplay();
  display.setTextSize(2);
  display.setTextColor(SSD1306_WHITE);
  display.setCursor(20, 20);
  display.println("Hello");
  display.display();
  delay(3000); 

  // Display "Enter Time" and default mode
  display.clearDisplay();
  displayMode();
  display.setCursor(0, 20);
  display.println("Enter Time...");
  display.display();
}

void loop() {
  if (digitalRead(SWITCH_MODE_BUTTON) == LOW) {
    switchMode(); 
    delay(200);  
  }

  if (!timerRunning) {
    // Check for button presses
    if (digitalRead(ADD_TIME_BUTTON) == LOW) {
      totalTime += 5; 
      displayTime(totalTime);
      delay(200); 
    }

    if (digitalRead(CLEAR_BUTTON) == LOW) {
      totalTime = 0; 
      monitorFSR = false; 
      display.clearDisplay();
      displayMode();
      display.setCursor(0, 20);
      display.println("Enter Time...");
      display.display();
      delay(200); 
    }

    if (digitalRead(START_BUTTON) == LOW && totalTime > 0) {
      if (currentMode == STRICT_FOCUS) {
        if (!checkWeight()) {
          display.clearDisplay();
          displayMode();
          display.setCursor(0, 20);
          display.println("Place phone");
          display.setCursor(0, 40);
          display.println("on sensor");
          display.display();
          delay(200); 
          return; // Exit loop to wait for correct weight
        }
      }
      timerRunning = true; // Start countdown
      previousMillis = millis(); // Starting time
      displayTime(totalTime);
      if (currentMode == STRICT_FOCUS) {
        initialFSRReading = analogRead(FSR_PIN); 
        monitorFSR = true; 
        weightRemoved = false; // Reset weight status
        firstBeepDone = false; // Reset first beep tracking
      }
      delay(200);
    }
  } else {
    // Timer is running
    updateTimer();
    if (monitorFSR) {
      monitorFSRChange();
    }
  }
}

void displayMode() {
  display.setTextSize(1);
  display.setCursor(0, 0);
  if (currentMode == TIMER_ONLY) {
    display.println("Timer Only");
  } else {
    display.println("Strict Focus");
  }
  display.display();
}

void switchMode() {
  if (currentMode == TIMER_ONLY) {
    currentMode = STRICT_FOCUS;
  } else {
    currentMode = TIMER_ONLY;
  }
  display.clearDisplay();
  displayMode();
  display.setCursor(0, 20);
  display.println("Enter Time...");
  display.display();
}

bool checkWeight() {
  int fsrReading = analogRead(FSR_PIN);
  return fsrReading >= weightThreshold; // Return true if weight is sufficient
}

void displayTime(int timeInSeconds) {
  int hours = timeInSeconds / 3600;
  int minutes = (timeInSeconds % 3600) / 60;
  int seconds = timeInSeconds % 60;

  // Clear and display time on OLED
  display.clearDisplay();
  displayMode();
  display.setTextSize(2);
  display.setCursor(0, 20);
  display.print(hours);
  display.print(":");
  if (minutes < 10) display.print("0");
  display.print(minutes);
  display.print(":");
  if (seconds < 10) display.print("0");
  display.println(seconds);
  display.display();
}

void updateTimer() {
  unsigned long currentMillis = millis();
  if (currentMillis - previousMillis >= 1000) { // 1-second interval
    previousMillis = currentMillis;
    totalTime--;
    displayTime(totalTime);

    if (totalTime <= 0) {
      timerRunning = false; // Stop timer when it hits 0
      monitorFSR = false;  
      display.clearDisplay();
      displayMode();
      display.setTextSize(2);
      display.setCursor(10, 20);
      display.println("Done!");
      display.display();
      playBeep();
    }
  }
}

void monitorFSRChange() {
  int currentFSRReading = analogRead(FSR_PIN);
  int difference = abs(currentFSRReading - initialFSRReading);

  if (difference >= 30) { // If weight changes significantly
    if (!weightRemoved) {
      weightRemoved = true;
      weightRemovedTime = millis();
      firstBeepDone = false;
    }

    unsigned long currentMillis = millis();
    if (!firstBeepDone && currentMillis - weightRemovedTime >= 1000) { // 1-second delay for first beep
      playBeep();
      firstBeepDone = true; // First beep done
      lastBeepTime = currentMillis; // Start 3-second timer
    } else if (firstBeepDone && currentMillis - lastBeepTime >= beepInterval) {
      lastBeepTime = currentMillis;
      playBeep(); // Beep every 3 seconds after first beep
    }
  } else {
    weightRemoved = false; 
    firstBeepDone = false; 
    lastBeepTime = millis();
  }
}

void playBeep() {
  for (int i = 0; i < 250; i++) {
    analogWrite(SPEAKER_PIN, 1);
    delayMicroseconds(1136);
    analogWrite(SPEAKER_PIN, 0);
    delayMicroseconds(1136);
  }
}