#include <Wire.h>
#include <LiquidCrystal_I2C.h>
#include <Keypad.h>
#include <Servo.h>
#include <EEPROM.h>

/* Locking mechanism definitions */
#define SERVO_PIN 6
#define SERVO_LOCK_POS 0 // Adjust this value according to your servo's locked position
#define SERVO_UNLOCK_POS 90 // Adjust this value according to your servo's unlocked position

Servo lockServo;

/* Display */
LiquidCrystal_I2C lcd(0x27, 16, 2); // Set the I2C address to 0x27 and dimensions to 16x2

/* Keypad setup */
const byte KEYPAD_ROWS = 4;
const byte KEYPAD_COLS = 4;
byte rowPins[KEYPAD_ROWS] = {5, 4, 3, 2};
byte colPins[KEYPAD_COLS] = {A3, A2, A1, A0};
char keys[KEYPAD_ROWS][KEYPAD_COLS] = {
  {'1', '2', '3', 'A'},
  {'4', '5', '6', 'B'},
  {'7', '8', '9', 'C'},
  {'*', '0', '#', 'D'}
};
Keypad keypad = Keypad(makeKeymap(keys), rowPins, colPins, KEYPAD_ROWS, KEYPAD_COLS);

/* EEPROM addresses */
#define CODE_ADDRESS 0
#define LOCK_STATE_ADDRESS 10

bool locked = true; // Initial lock state

String inputSecretCode() {
  lcd.setCursor(5, 1);
  lcd.print("[____]");
  lcd.setCursor(6, 1);
  String result = "";
  while (result.length() < 4) {
    char key = keypad.getKey();
    if (key >= '0' && key <= '9') {
      lcd.print('*');
      result += key;
    }
  }
  return result;
}

void setNewCode() {
  lcd.clear();
  lcd.setCursor(0, 0);
  lcd.print("Set Up Your Code...");
  lcd.setCursor(0, 1);
  String newCode = inputSecretCode();
  lcd.clear();
  lcd.setCursor(1, 0);
  lcd.print("Confirm Code:");
  lcd.setCursor(0, 1);
  String confirmCode = inputSecretCode();
  int attempts = 0;
  while (newCode != confirmCode && attempts < 5) {
    lcd.clear();
    lcd.setCursor(0, 0);
    lcd.print("Code Mismatch!");
    lcd.setCursor(0, 1);
    lcd.print("Retry :");
    newCode = inputSecretCode();
    lcd.clear();
    lcd.setCursor(0, 0);
    lcd.print("Confirm New Code:");
    confirmCode = inputSecretCode();
    attempts++;
  }
  if (newCode.equals(confirmCode)) {
    // Convert the code to integers before storing in EEPROM
    int code[4];
    code[0] = newCode[0] - '0';
    code[1] = newCode[1] - '0';
    code[2] = newCode[2] - '0';
    code[3] = newCode[3] - '0';
    // Store the code in EEPROM
    for (int i = 0; i < 4; i++) {
      EEPROM.write(CODE_ADDRESS + i, code[i]);
    }
    lcd.clear();
    lcd.setCursor(4, 0);
    lcd.print("Code Set!");
    delay(1000);
  } else {
    lcd.clear();
    lcd.setCursor(0, 0);
    lcd.print("Max Attempts Reached!");
    delay(2000);
  }
}

void lock() {
  Serial.println("Locking...");
  lockServo.write(SERVO_LOCK_POS);
  locked = true;
  EEPROM.write(LOCK_STATE_ADDRESS, 1); // Lock state: Locked
}

void unlock() {
  Serial.println("Unlocking...");
  lockServo.write(SERVO_UNLOCK_POS);
  locked = false;
  EEPROM.write(LOCK_STATE_ADDRESS, 0); // Lock state: Unlocked
}

void setup() {
  Serial.begin(9600); // Initialize Serial communication
  lcd.init(); // Initialize the LCD
  lcd.backlight(); // Turn on backlight
  lockServo.attach(SERVO_PIN); // Attach the servo
  // Make sure the physical lock is sync with the EEPROM state
  locked = EEPROM.read(LOCK_STATE_ADDRESS);
  if (locked) {
    lock();
  } else {
    unlock();
  }
  // Check if the initial code is set
  if (EEPROM.read(CODE_ADDRESS) == 0xFF && EEPROM.read(CODE_ADDRESS + 1) == 0xFF &&
      EEPROM.read(CODE_ADDRESS + 2) == 0xFF && EEPROM.read(CODE_ADDRESS + 3) == 0xFF) {
    lcd.setCursor(0, 1);
    lcd.print("Set Initial Code:");
    setNewCode();
  }
  // Display startup message
  lcd.setCursor(0, 0);
  lcd.print("God Is My Light");
  delay(1000);
  lcd.setCursor(0, 0);
  lcd.print("Electronic Safe");
  delay(1000);
}

void loop() {
  // Logic for locked state
  if (locked) {
    lcd.clear();
    lcd.setCursor(0, 0);
    lcd.print("Locked Put Code:");
    String userCode = inputSecretCode();
    // Check if entered code is correct
    int code[4];
    for (int i = 0; i < 4; i++) {
      code[i] = EEPROM.read(CODE_ADDRESS + i);
    }
    if (userCode.toInt() == code[0] * 1000 + code[1] * 100 + code[2] * 10 + code[3]) {
      unlock();
      lcd.clear();
      lcd.print("    Unlocked!");
      delay(1000);
    } else {
      lcd.clear();
      lcd.print("    Wrong Code!");
      delay(1000);
    }
  }
  // Logic for unlocked state
  else {
    // Check if the display needs to be updated
    static bool displayNeedsUpdate = true;
    if (displayNeedsUpdate) {
      lcd.clear();
      lcd.setCursor(3, 0);
      lcd.print("# To Lock");
      lcd.setCursor(1, 1);
      lcd.print("A For New code");
      displayNeedsUpdate = false; // Set to false since the display is now up-to-date
    }

    char key = keypad.getKey();
    // Check if # key is pressed to lock
    if (key == '#') {
      lock();
      displayNeedsUpdate = true; // Set to true to update the display next loop iteration
    }
    // Check if A key is pressed to set new code
    else if (key == 'A') {
      setNewCode();
      displayNeedsUpdate = true; // Set to true to update the display next loop iteration
    }
  }
}