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

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

// Pin Definitions
const int SERVO_PIN = 9;   
const int PIR_PIN = 13;
const int BUZZER_PIN = 12;
const int RED_LED = 11;
const int GREEN_LED = 10;

// Security Logic
const String CORRECT_PIN = "1234"; 
String inputString = "";
int failedAttempts = 0;
const int MAX_ATTEMPTS = 2; 

// Keypad Configuration
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'}
};

// Reverted to standard Analog pins so it matches your wiring!
byte rowPins[ROWS] = {5, 4, 3, 2}; 
byte colPins[COLS] = {A3, A2, A1, A0}; 

Keypad keypad = Keypad(makeKeymap(keys), rowPins, colPins, ROWS, COLS);
Servo lockServo;

void setup() {
  Serial.begin(115200);

  pinMode(PIR_PIN, INPUT);
  pinMode(BUZZER_PIN, OUTPUT);
  pinMode(RED_LED, OUTPUT);
  pinMode(GREEN_LED, OUTPUT);
  
  lockServo.attach(SERVO_PIN);
  lockServo.write(0); // Ensure locked (0 degrees)

  delay(500); // Give OLED time to power up

  if(!display.begin(SSD1306_SWITCHCAPVCC, 0x3C)) {
    Serial.println(F("OLED failed"));
    for(;;); 
  }

  showWelcome();
}

void loop() {
  // 1. Security Check (PIR)
  if (digitalRead(PIR_PIN) == HIGH) {
    securityAlert();
    // Wait until motion stops so the keypad doesn't freeze
    while(digitalRead(PIR_PIN) == HIGH) {
      delay(50); 
    }
    updateEntryScreen(); 
  }

  // 2. PIN Entry Check
  char key = keypad.getKey();
  if (key) {
    handleKeypad(key);
  }
}

void showWelcome() {
  display.clearDisplay();
  display.setTextColor(WHITE);
  display.setTextSize(1);
  display.setCursor(20, 15);
  display.println("SMART BANK SAFE");
  display.setCursor(30, 35);
  display.println("System Ready");
  display.display();
  delay(2000);
  updateEntryScreen();
}

void updateEntryScreen() {
  display.clearDisplay();
  display.setTextSize(1);
  display.setCursor(0, 0);
  display.println("ENTER PIN:");
  display.setCursor(0, 30);
  display.setTextSize(2);
  
  // This explicitly shows the numbers typed on the screen
  String displayChars = inputString;
  for(int i = inputString.length(); i < 4; i++) {
    displayChars += "_";
  }
  
  display.println(displayChars);
  display.display();
}

void handleKeypad(char key) {
  // Only accept numbers 0-9
  if (key >= '0' && key <= '9') {
    inputString += key;
    tone(BUZZER_PIN, 1000, 50); 
    updateEntryScreen();
  } 
  // '*' acts as a Clear button
  else if (key == '*') {
    inputString = "";
    updateEntryScreen();
  }

  // Check the password when 4 digits are entered
  if (inputString.length() == 4) {
    if (inputString == CORRECT_PIN) {
      grantAccess();
    } else {
      denyAccess();
    }
    inputString = ""; // Reset for next try
  }
}

void grantAccess() {
  failedAttempts = 0;
  
  // SUCCESS: Green LED ON, Red LED OFF
  digitalWrite(GREEN_LED, HIGH);
  digitalWrite(RED_LED, LOW); 
  
  // SUCCESS: OLED Message
  display.clearDisplay();
  display.setTextSize(2);
  display.setCursor(20, 20);
  display.println("GRANTED");
  display.display();
  
  // SUCCESS: Servo Moves to 90 degrees
  lockServo.write(90); 
  tone(BUZZER_PIN, 1500, 500);
  
  delay(4000); // Door stays open for 4 seconds
  
  // Relock System
  lockServo.write(0); 
  digitalWrite(GREEN_LED, LOW);
  updateEntryScreen();
}

void denyAccess() {
  failedAttempts++;
  
  // FAIL: Red LED and Buzzer ON
  digitalWrite(RED_LED, HIGH);
  tone(BUZZER_PIN, 400, 800); 
  
  // FAIL: OLED Message
  display.clearDisplay();
  display.setTextSize(2);
  display.setCursor(25, 20);
  display.println("DENIED");
  display.display();
  
  delay(2000);
  digitalWrite(RED_LED, LOW);

  if (failedAttempts >= MAX_ATTEMPTS) {
    lockoutAlarm();
  } else {
    updateEntryScreen();
  }
}

void securityAlert() {
  display.clearDisplay();
  display.setCursor(15, 10);
  display.setTextSize(2);
  display.println("WARNING!");
  display.setCursor(15, 40);
  display.setTextSize(1);
  display.println("Motion Detected");
  display.display();
  
  digitalWrite(RED_LED, HIGH);
  tone(BUZZER_PIN, 800, 100); 
  delay(100);
  digitalWrite(RED_LED, LOW);
}

void lockoutAlarm() {
  display.clearDisplay();
  display.setTextSize(2);
  display.setCursor(10, 20);
  display.println("LOCKED!");
  display.display();

  for(int i = 0; i < 5; i++) {
    digitalWrite(RED_LED, HIGH);
    tone(BUZZER_PIN, 1200);
    delay(300);
    digitalWrite(RED_LED, LOW);
    noTone(BUZZER_PIN);
    delay(300);
  }
  
  failedAttempts = 0; 
  updateEntryScreen();
}