/*
 * Automatic Load Shedding Board - Corrected Version
 * Arduino Uno + I2C LCD 20x4 + DS3231 RTC + 4 Buttons + 3 Channel Relay
 */

#include <Wire.h>
#include <LiquidCrystal_I2C.h>
#include <RTClib.h>

// LCD Configuration
LiquidCrystal_I2C lcd(0x27, 20, 4);

// RTC Object
RTC_DS3231 rtc;

// Pins
#define BTN_UP 2
#define BTN_DOWN 3
#define BTN_ENTER 4
#define BTN_BACK 5
#define RELAY_1 6
#define RELAY_2 7
#define RELAY_3 8
#define BACKLIGHT_PIN 10 // Connect to LCD LED jumper pin

enum MenuState {
  HOME_SCREEN, MAIN_MENU, TIME_DATE_MENU, SET_TIME, SET_DATE,
  LOAD_PARAMS_MENU, LOAD_SETTINGS, SETTINGS_MENU,
  AUTO_MODE_SETTING, BRIGHTNESS_SETTING, RESET_SETTING, ABOUT_SCREEN
};

MenuState currentState = HOME_SCREEN;
int menuIndex = 0;
int subMenuIndex = 0;
bool autoMode = true;
int brightness = 255;
int currentLoadEdit = 0; 
int loadEditField = 0; // 0:Enabled, 1:OnH, 2:OnM, 3:OffH, 4:OffM

struct LoadConfig {
  char name[12];
  int onHour, onMinute, offHour, offMinute;
  bool enabled, state;
};

LoadConfig loads[3] = {
  {"Load 1", 6, 0, 18, 0, true, false},
  {"Load 2", 7, 0, 19, 0, true, false},
  {"Load 3", 8, 0, 20, 0, true, false}
};

int tempHour, tempMinute, tempSecond, tempDay, tempMonth, tempYear, editingField;
unsigned long lastDebounce = 0;
const unsigned long debounceDelay = 200;

void setup() {
  lcd.init();
  lcd.backlight();
  pinMode(BACKLIGHT_PIN, OUTPUT);
  analogWrite(BACKLIGHT_PIN, brightness);
  
  if (!rtc.begin()) {
    lcd.print("RTC Error!");
    while (1);
  }

  pinMode(BTN_UP, INPUT_PULLUP);
  pinMode(BTN_DOWN, INPUT_PULLUP);
  pinMode(BTN_ENTER, INPUT_PULLUP);
  pinMode(BTN_BACK, INPUT_PULLUP);
  
  pinMode(RELAY_1, OUTPUT);
  pinMode(RELAY_2, OUTPUT);
  pinMode(RELAY_3, OUTPUT);
  
  // Relays are active LOW
  digitalWrite(RELAY_1, HIGH);
  digitalWrite(RELAY_2, HIGH);
  digitalWrite(RELAY_3, HIGH);
  
  showSplashScreen();
}

void loop() {
  handleButtons();
  if (autoMode) checkSchedule();
  updateDisplay();
  delay(50);
}

void updateDisplay() {
  switch (currentState) {
    case HOME_SCREEN:        displayHomeScreen(); break;
    case MAIN_MENU:          displayMainMenu(); break;
    case TIME_DATE_MENU:     displayTimeDateMenu(); break;
    case SET_TIME:           displaySetTime(); break;
    case SET_DATE:           displaySetDate(); break;
    case LOAD_PARAMS_MENU:   displayLoadParamsMenu(); break;
    case LOAD_SETTINGS:      displayLoadSettings(currentLoadEdit); break;
    case SETTINGS_MENU:      displaySettingsMenu(); break;
    case AUTO_MODE_SETTING:  displayAutoModeSetting(); break;
    case BRIGHTNESS_SETTING: displayBrightnessSetting(); break;
    case RESET_SETTING:      displayResetSetting(); break;
    case ABOUT_SCREEN:       displayAboutScreen(); break;
  }
}

// --- LOGIC FUNCTIONS ---

void checkSchedule() {
  DateTime now = rtc.now();
  int currentMins = (now.hour() * 60) + now.minute();
  
  for (int i = 0; i < 3; i++) {
    if (!loads[i].enabled) {
      setRelay(i, false);
      continue;
    }
    
    int onMins = (loads[i].onHour * 60) + loads[i].onMinute;
    int offMins = (loads[i].offHour * 60) + loads[i].offMinute;
    bool shouldBeOn = false;

    if (onMins < offMins) {
      // Normal range (e.g., 08:00 to 17:00)
      if (currentMins >= onMins && currentMins < offMins) shouldBeOn = true;
    } else {
      // Overnight range (e.g., 22:00 to 04:00)
      if (currentMins >= onMins || currentMins < offMins) shouldBeOn = true;
    }
    setRelay(i, shouldBeOn);
  }
}

void setRelay(int relayNum, bool state) {
  loads[relayNum].state = state;
  digitalWrite(RELAY_1 + relayNum, state ? LOW : HIGH);
}

// --- BUTTON HANDLING ---

void handleButtons() {
  if (millis() - lastDebounce < debounceDelay) return;

  if (digitalRead(BTN_ENTER) == LOW) { lastDebounce = millis(); handleEnter(); }
  if (digitalRead(BTN_BACK) == LOW)  { lastDebounce = millis(); handleBack(); }
  if (digitalRead(BTN_UP) == LOW)    { lastDebounce = millis(); handleUp(); }
  if (digitalRead(BTN_DOWN) == LOW)  { lastDebounce = millis(); handleDown(); }
}

void handleEnter() {
  switch (currentState) {
    case HOME_SCREEN: currentState = MAIN_MENU; break;
    case MAIN_MENU:
      if (menuIndex == 0) currentState = TIME_DATE_MENU;
      else if (menuIndex == 1) currentState = LOAD_PARAMS_MENU;
      else if (menuIndex == 2) currentState = SETTINGS_MENU;
      else currentState = ABOUT_SCREEN;
      subMenuIndex = 0;
      break;
    case TIME_DATE_MENU:
      if (subMenuIndex == 0) { 
        DateTime n = rtc.now(); tempHour = n.hour(); tempMinute = n.minute(); tempSecond = 0;
        editingField = 0; currentState = SET_TIME; 
      }
      else if (subMenuIndex == 1) { 
        DateTime n = rtc.now(); tempDay = n.day(); tempMonth = n.month(); tempYear = n.year();
        editingField = 0; currentState = SET_DATE; 
      }
      else currentState = MAIN_MENU;
      break;
    case SET_TIME:
      editingField++;
      if (editingField > 2) {
        DateTime now = rtc.now();
        rtc.adjust(DateTime(now.year(), now.month(), now.day(), tempHour, tempMinute, tempSecond));
        currentState = TIME_DATE_MENU;
      }
      break;
    case SET_DATE:
      editingField++;
      if (editingField > 2) {
        DateTime now = rtc.now();
        rtc.adjust(DateTime(tempYear, tempMonth, tempDay, now.hour(), now.minute(), now.second()));
        currentState = TIME_DATE_MENU;
      }
      break;
    case LOAD_PARAMS_MENU:
      currentLoadEdit = subMenuIndex;
      loadEditField = 0;
      currentState = LOAD_SETTINGS;
      break;
    case LOAD_SETTINGS:
      loadEditField++;
      if (loadEditField > 4) currentState = LOAD_PARAMS_MENU;
      break;
    case SETTINGS_MENU:
      if (subMenuIndex == 0) currentState = AUTO_MODE_SETTING;
      else if (subMenuIndex == 1) currentState = BRIGHTNESS_SETTING;
      else if (subMenuIndex == 2) currentState = RESET_SETTING;
      else currentState = MAIN_MENU;
      break;
    case RESET_SETTING: performSystemReset(); currentState = MAIN_MENU; break;
  }
}

void handleBack() {
  if (currentState == MAIN_MENU) currentState = HOME_SCREEN;
  else if (currentState == LOAD_SETTINGS) currentState = LOAD_PARAMS_MENU;
  else currentState = MAIN_MENU;
}

void handleUp() {
  if (currentState == MAIN_MENU && menuIndex > 0) menuIndex--;
  else if ((currentState == TIME_DATE_MENU || currentState == LOAD_PARAMS_MENU || currentState == SETTINGS_MENU) && subMenuIndex > 0) subMenuIndex--;
  else if (currentState == SET_TIME) {
    if (editingField == 0) tempHour = (tempHour + 1) % 24;
    else if (editingField == 1) tempMinute = (tempMinute + 1) % 60;
  }
  else if (currentState == SET_DATE) {
    if (editingField == 0) tempDay = (tempDay % 31) + 1;
    else if (editingField == 1) tempMonth = (tempMonth % 12) + 1;
    else if (editingField == 2) tempYear++;
  }
  else if (currentState == LOAD_SETTINGS) handleLoadEdit(currentLoadEdit, true);
  else if (currentState == BRIGHTNESS_SETTING) {
    brightness = min(brightness + 25, 255);
    analogWrite(BACKLIGHT_PIN, brightness);
  }
}

void handleDown() {
  if (currentState == MAIN_MENU && menuIndex < 3) menuIndex++;
  else if ((currentState == TIME_DATE_MENU || currentState == LOAD_PARAMS_MENU || currentState == SETTINGS_MENU) && subMenuIndex < 2) subMenuIndex++;
  else if (currentState == SET_TIME) {
    if (editingField == 0) tempHour = (tempHour + 23) % 24;
    else if (editingField == 1) tempMinute = (tempMinute + 59) % 60;
  }
  else if (currentState == LOAD_SETTINGS) handleLoadEdit(currentLoadEdit, false);
  else if (currentState == BRIGHTNESS_SETTING) {
    brightness = max(brightness - 25, 0);
    analogWrite(BACKLIGHT_PIN, brightness);
  }
}

// --- DISPLAY RENDERING (Simplified for space) ---

void displayHomeScreen() {
  DateTime now = rtc.now();
  lcd.clear();
  lcd.setCursor(4, 0); lcd.print("LOAD CONTROL");
  lcd.setCursor(1, 1);
  lcd.print(now.day()); lcd.print("/"); lcd.print(now.month()); lcd.print("/"); lcd.print(now.year());
  lcd.setCursor(6, 2);
  if (now.hour() < 10) lcd.print("0"); lcd.print(now.hour()); lcd.print(":");
  if (now.minute() < 10) lcd.print("0"); lcd.print(now.minute());
  lcd.setCursor(0, 3);
  for(int i=0; i<3; i++) {
    lcd.print("L"); lcd.print(i+1); lcd.print(loads[i].state ? ":ON " : ":OFF");
  }
}

void displayLoadSettings(int ln) {
  lcd.clear();
  lcd.print(loads[ln].name);
  lcd.setCursor(0,1); lcd.print(loadEditField==0 ? ">En: " : " En: "); lcd.print(loads[ln].enabled ? "YES" : "NO");
  lcd.setCursor(0,2); lcd.print(loadEditField==1 || loadEditField==2 ? ">ON: " : " ON: "); 
  lcd.print(loads[ln].onHour); lcd.print(":"); lcd.print(loads[ln].onMinute);
  lcd.setCursor(0,3); lcd.print(loadEditField==3 || loadEditField==4 ? ">OFF: " : " OFF: "); 
  lcd.print(loads[ln].offHour); lcd.print(":"); lcd.print(loads[ln].offMinute);
}

// Helper for Load Editing
void handleLoadEdit(int ln, bool inc) {
  int amt = inc ? 1 : -1;
  if (loadEditField == 0) loads[ln].enabled = !loads[ln].enabled;
  else if (loadEditField == 1) loads[ln].onHour = (loads[ln].onHour + amt + 24) % 24;
  else if (loadEditField == 2) loads[ln].onMinute = (loads[ln].onMinute + amt + 60) % 60;
  else if (loadEditField == 3) loads[ln].offHour = (loads[ln].offHour + amt + 24) % 24;
  else if (loadEditField == 4) loads[ln].offMinute = (loads[ln].offMinute + amt + 60) % 60;
}

// Add the rest of your UI functions (Splash, Main Menu, etc.) from original code.
// They remain mostly the same but ensure they call lcd.clear() at the start.

void showSplashScreen() {
  lcd.clear();
  lcd.setCursor(3, 1); lcd.print("LOAD SHEDDING");
  lcd.setCursor(5, 2); lcd.print("SYSTEM V1.1");
  delay(2000);
}

void displayMainMenu() {
  lcd.clear();
  const char* m[] = {"Time & Date", "Load Params", "Settings", "About"};
  for(int i=0; i<4; i++) {
    lcd.setCursor(0, i);
    lcd.print(menuIndex == i ? ">" : " ");
    lcd.print(m[i]);
  }
}

void performSystemReset() {
  for(int i=0; i<3; i++) {
    loads[i].enabled = true;
    loads[i].onHour = 6+i;
    loads[i].offHour = 18+i;
  }
  autoMode = true;
  brightness = 255;
  analogWrite(BACKLIGHT_PIN, 255);
}

void displayTimeDateMenu() {
  lcd.clear();
  lcd.setCursor(0, 0);
  lcd.print("== TIME & DATE ==");
  
  const char* menuItems[] = {"Set Time", "Set Date", "Back"};
  
  for (int i = 0; i < 3; i++) {
    lcd.setCursor(1, i + 1);
    if (i == subMenuIndex) lcd.print(">");
    else lcd.print(" ");
    lcd.print(menuItems[i]);
  }
}
void displaySetTime() {
  lcd.clear();
  lcd.setCursor(0, 0);
  lcd.print("===== SET TIME =====");
  
  lcd.setCursor(3, 2);
  if (editingField == 0) lcd.print("[");
  if (tempHour < 10) lcd.print("0");
  lcd.print(tempHour);
  if (editingField == 0) lcd.print("]");
  else lcd.print(" ");
  
  lcd.print(":");
  
  if (editingField == 1) lcd.print("[");
  if (tempMinute < 10) lcd.print("0");
  lcd.print(tempMinute);
  if (editingField == 1) lcd.print("]");
  else lcd.print(" ");
  
  lcd.print(":");
  
  if (editingField == 2) lcd.print("[");
  if (tempSecond < 10) lcd.print("0");
  lcd.print(tempSecond);
  if (editingField == 2) lcd.print("]");
  
  lcd.setCursor(1, 3);
  lcd.print("ENTER:Save BACK:Exit");
}
void displaySetDate() {
  lcd.clear();
  lcd.setCursor(0, 0);
  lcd.print("===== SET DATE =====");
  
  lcd.setCursor(1, 2);
  if (editingField == 0) lcd.print("[");
  if (tempDay < 10) lcd.print("0");
  lcd.print(tempDay);
  if (editingField == 0) lcd.print("]");
  else lcd.print(" ");
  
  lcd.print("/");
  
  if (editingField == 1) lcd.print("[");
  if (tempMonth < 10) lcd.print("0");
  lcd.print(tempMonth);
  if (editingField == 1) lcd.print("]");
  else lcd.print(" ");
  
  lcd.print("/");
  
  if (editingField == 2) lcd.print("[");
  lcd.print(tempYear);
  if (editingField == 2) lcd.print("]");
  
  lcd.setCursor(1, 3);
  lcd.print("ENTER:Save BACK:Exit");
}

void displayLoadParamsMenu() {
  lcd.clear();
  lcd.setCursor(0, 0);
  lcd.print("== LOAD PARAMETERS ==");
  
  for (int i = 0; i < 3; i++) {
    lcd.setCursor(1, i + 1);
    if (i == subMenuIndex) lcd.print(">");
    else lcd.print(" ");
    lcd.print(loads[i].name);
    lcd.print(loads[i].enabled ? " [ON]" : " [OFF]");
  }
}

void displaySettingsMenu() {
  lcd.clear();
  lcd.setCursor(0, 0);
  lcd.print("===== SETTINGS =====");
  
  const char* menuItems[] = {"Auto Mode", "Brightness", "Reset System", "Back"};
  
  for (int i = 0; i < 3 && i + subMenuIndex < 4; i++) {
    lcd.setCursor(1, i + 1);
    if (i == 0) lcd.print(">");
    else lcd.print(" ");
    lcd.print(menuItems[i + subMenuIndex]);
  }
}

void displayAutoModeSetting() {
  lcd.clear();
  lcd.setCursor(0, 0);
  lcd.print("===== AUTO MODE =====");
  
  lcd.setCursor(3, 2);
  lcd.print("Mode: ");
  lcd.print(autoMode ? "AUTO  " : "MANUAL");
  
  lcd.setCursor(1, 3);
  lcd.print("UP/DOWN:Toggle BACK");
}

void displayBrightnessSetting() {
  lcd.clear();
  lcd.setCursor(0, 0);
  lcd.print("==== BRIGHTNESS ====");
  
  lcd.setCursor(2, 2);
  lcd.print("Level: ");
  int percent = map(brightness, 0, 255, 0, 100);
  if (percent < 100) lcd.print(" ");
  if (percent < 10) lcd.print(" ");
  lcd.print(percent);
  lcd.print("%");
  
  // Draw brightness bar
  lcd.setCursor(0, 3);
  lcd.print("[");
  int bars = map(brightness, 0, 255, 0, 18);
  for (int i = 0; i < 18; i++) {
    if (i < bars) lcd.print("=");
    else lcd.print(" ");
  }
  lcd.print("]");
}

void displayResetSetting() {
  lcd.clear();
  lcd.setCursor(0, 0);
  lcd.print("===== RESET SYSTEM ==");
  
  lcd.setCursor(1, 1);
  lcd.print("Are you sure you");
  lcd.setCursor(1, 2);
  lcd.print("want to RESET?");
  
  lcd.setCursor(0, 3);
  lcd.print("ENTER:YES  BACK:NO");
}

void displayAboutScreen() {
  lcd.clear();
  lcd.setCursor(0, 0);
  lcd.print("===== ABOUT =====");
  
  lcd.setCursor(0, 1);
  lcd.print("Load Shedding Board");
  lcd.setCursor(0, 2);
  lcd.print("Version: 1.0.0");
  lcd.setCursor(0, 3);
  lcd.print("Press BACK to exit");
}