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

// LCD setup
LiquidCrystal_I2C lcd(0x27, 16, 2);

// Keypad setup
const byte ROWS = 4;
const byte COLS = 3;
char keys[ROWS][COLS] = {
  {'1','2','3'},
  {'4','5','6'},
  {'7','8','9'},
  {'*','0','#'}
};
byte rowPins[ROWS] = {9, 8, 7, 6};
byte colPins[COLS] = {5, 4, 3};
Keypad keypad = Keypad(makeKeymap(keys), rowPins, colPins, ROWS, COLS);

// RGB LED pins
int redPin = 10;
int greenPin = 11;
int bluePin = 12;

// Buzzer pin
int buzzer = 2;

// Push button pin
int buttonPin = 13;
int lastButtonState = HIGH;
int gameMode = 1; // 1 = multiplication, 2 = division

// Game variables
int num1, num2, correctAnswer;
int userAnswer = 0;

// ---------- Setup ----------
void setup() {
  lcd.init();
  lcd.backlight();

  pinMode(redPin, OUTPUT);
  pinMode(greenPin, OUTPUT);
  pinMode(bluePin, OUTPUT);
  pinMode(buzzer, OUTPUT);
  pinMode(buttonPin, INPUT_PULLUP);

  randomSeed(analogRead(A1)); // seed random numbers
  generateQuestion();
}

// ---------- Main Loop ----------
void loop() {
  // Check button press to change mode
  int buttonState = digitalRead(buttonPin);
  if (buttonState == LOW && lastButtonState == HIGH) {
    gameMode = (gameMode == 1) ? 2 : 1; // toggle
    generateQuestion();
    delay(200); // debounce
  }
  lastButtonState = buttonState;

  // Handle keypad
  char key = keypad.getKey();

  if (key) {
    // Number key
    if (key >= '0' && key <= '9') {
      userAnswer = userAnswer * 10 + (key - '0');
      lcd.setCursor(0, 1);
      lcd.print("Ans: ");
      lcd.print(userAnswer);
      setColor(0, 0, 255); // Blue while typing
    }

    else if (key == '#') {   // submit answer
      if (userAnswer != 0) {
        if (userAnswer == correctAnswer) {
          showCorrect();
        } else {
          showWrong();
        }
      }
    }

    else if (key == '*') {   // reset answer
      userAnswer = 0;
      lcd.setCursor(0, 1);
      lcd.print("Ans:       ");
      setColor(0, 0, 255);
      delay(300);
      setColor(0, 0, 0);
    }
  }
}

// ---------- Functions ----------

// Make a new question
void generateQuestion() {
  lcd.clear();
  num1 = random(1, 10);
  num2 = random(1, 10);

  if (gameMode == 1) { // Multiplication
    correctAnswer = num1 * num2;
    lcd.setCursor(0, 0);
    lcd.print(num1);
    lcd.print(" x ");
    lcd.print(num2);
    lcd.print(" = ?");
  } else { // Division
    correctAnswer = num1;       // ensure integer division
    int product = num1 * num2;  // dividend
    lcd.setCursor(0, 0);
    lcd.print(product);
    lcd.print(" / ");
    lcd.print(num2);
    lcd.print(" = ?");
  }

  lcd.setCursor(10, 0);
  if (gameMode == 1) lcd.print("MUL");
  else lcd.print("DIV");

  userAnswer = 0;
  lcd.setCursor(0, 1);
  lcd.print("Ans: ");
}

// When correct
void showCorrect() {
  lcd.setCursor(0, 1);
  lcd.print("Correct!     ");
  setColor(0, 255, 0);
  tone(buzzer, 1000, 300);
  delay(1000);
  setColor(0, 0, 0);
  generateQuestion();
}

// When wrong
void showWrong() {
  lcd.setCursor(0, 1);
  lcd.print("Wrong!       ");
  setColor(255, 0, 0);
  tone(buzzer, 300, 600);
  delay(1000);
  setColor(0, 0, 0);
  generateQuestion();
}

// Set RGB LED color
void setColor(int r, int g, int b) {
  analogWrite(redPin, r);
  analogWrite(greenPin, g);
  analogWrite(bluePin, b);
}