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

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

#define BUZZER_PIN 3
#define SWITCH_PIN 2

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};
byte colPins[COLS] = {5, 4, 3, 10};

Keypad customKeypad = Keypad(makeKeymap(keys), rowPins, colPins, ROWS, COLS);

const String correctCode = "7355608";
String inputCode = "";

bool systemOn = false;
bool bombArmed = false;
bool defusing = false;
unsigned long bombTimer = 0;
int beepDelay = 1000;
bool beepState = false;

void resetSystem();
void armBomb();
void updateBomb();
void explode();

void setup() {
pinMode(BUZZER_PIN, OUTPUT);
pinMode(SWITCH_PIN, INPUT_PULLUP);

if (!display.begin(SSD1306_SWITCHCAPVCC, 0x3C)) {
for (;;);
}
display.clearDisplay();
display.display();

Serial.begin(9600);
}

void loop() {
// Read switch state
systemOn = digitalRead(SWITCH_PIN) == LOW; // Switch ON = LOW (GND)

if (!systemOn) {
resetSystem();
return;
}

char key = customKeypad.getKey(); // FIXED: Added parentheses

if (key) { // If a key was pressed
Serial.print("Key Pressed: ");
Serial.println(key);

// Add key to input
if (!bombArmed) {
inputCode += key;

// Update display
display.clearDisplay();
display.setTextSize(2);
display.setCursor(10, 20);
display.print(inputCode);
display.display();

if (inputCode == correctCode) {
armBomb();
} else if (inputCode.length() >= correctCode.length()) {
inputCode = ""; // Reset input if wrong
}
}
}

if (bombArmed) {
updateBomb();
}
}

void resetSystem() {
display.clearDisplay();
display.setTextSize(2);
display.setCursor(10, 20);
display.print("OFF");
display.display();

inputCode = "";
bombArmed = false;
defusing = false;
bombTimer = 0;
noTone(BUZZER_PIN);
}

void armBomb() {
bombArmed = true;
bombTimer = millis();
display.clearDisplay();
display.setTextSize(2);
display.setCursor(10, 20);
display.print("Bomb Armed!");
display.display();

tone(BUZZER_PIN, 1000, 200);
delay(200);
}

void updateBomb() {
unsigned long elapsed = millis() - bombTimer;

if (elapsed > 5000) {
explode();
} else {
if (elapsed % beepDelay < 100) {
tone(BUZZER_PIN, 1000, 50);
}
}
}

void explode() {
display.clearDisplay();
display.setTextSize(2);
display.setCursor(10, 20);
display.print("BOOM!");
display.display();

tone(BUZZER_PIN, 500, 3000);
delay(3000);
resetSystem();
}