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

const int correctButtonPin = 8;  // Pin for the correct button
const int incorrectButtonPin = 9;  // Pin for the incorrect button
const int buzzerPin = 12;  // Pin for the buzzer
const int resetButtonPin = 11;  // Pin for the reset button
const int blueButtonPin = 10;  // Pin for the blue button
const int CLK = 2;  // TM1637 clock pin
const int DIO = 3;  // TM1637 data pin
const int correctLEDPin = 42;  // Pin for the green LED indicating correct answer
const int incorrectLEDPin = 43;  // Pin for the red LED indicating incorrect answer

// Correct answers LED pins
const int correctLEDPins[10] = {22, 23, 24, 25, 26, 27, 28, 29, 30, 31};
// Incorrect answers LED pins
const int incorrectLEDPins[10] = {32, 33, 34, 35, 36, 37, 38, 39, 40, 41};

// Define the notes and durations for the winner tune
const int winnerMelody[] = {
  659, 0, 659, 0, 659, 0, 523, 659, 784
};
const int winnerNoteDurations[] = {
  200, 200, 200, 200, 200, 200, 400, 400, 800
};

// Define the notes and durations for the loser tune
const int loserMelody[] = {
  330, 0, 330, 0, 330, 0, 294, 0, 262
};
const int loserNoteDurations[] = {
  300, 150, 300, 150, 300, 150, 500, 150, 1000
};

enum Difficulty { BASIC, MEDIUM, HARD };
enum GameState { MENU, CATEGORY, IN_GAME, GAME_OVER };

int correctAnswer;
int displayedAnswer;
bool isAnswerCorrect;
int correctCount = 0;  // Counter for correct answers
int incorrectCount = 0;  // Counter for incorrect answers
bool timedOut = false;
Difficulty currentDifficulty;
GameState currentState = MENU;

LiquidCrystal_I2C lcd(0x27, 20, 4);  // Set the LCD address to 0x27 for a 20 chars and 4 line display
TM1637Display display(CLK, DIO);

void setup() {
  pinMode(correctButtonPin, INPUT_PULLUP);  // Use internal pull-up resistor
  pinMode(incorrectButtonPin, INPUT_PULLUP);  // Use internal pull-up resistor
  pinMode(resetButtonPin, INPUT_PULLUP);  // Use internal pull-up resistor
  pinMode(blueButtonPin, INPUT_PULLUP);  // Use internal pull-up resistor
  pinMode(buzzerPin, OUTPUT);  // Set buzzer pin as output
  pinMode(correctLEDPin, OUTPUT);  // Set the correct LED pin as output
  pinMode(incorrectLEDPin, OUTPUT);  // Set the incorrect LED pin as output

  for (int i = 0; i < 10; i++) {
    pinMode(correctLEDPins[i], OUTPUT);
    pinMode(incorrectLEDPins[i], OUTPUT);
    digitalWrite(correctLEDPins[i], LOW);
    digitalWrite(incorrectLEDPins[i], LOW);
  }

  Serial.begin(9600);
  lcd.begin(20, 4);  // Initialize the LCD with 20 columns and 4 rows
  lcd.backlight();
  display.setBrightness(0x0f);  // Set the brightness of the display
  randomSeed(analogRead(0));  // Seed the random number generator with an unconnected analog pin
}

void loop() {
  static GameState previousState = MENU;  // Track the previous state to detect changes

  if (previousState != currentState) {
    lcd.clear();  // Clear the display only when the state changes
    previousState = currentState;  // Update the previous state
  }

  switch (currentState) {
    case MENU:
      displayMenu();
      handleMenuInput();
      break;
    case CATEGORY:
      displayCategory();
      handleCategoryInput();
      break;
    case IN_GAME:
      generateQuestion();
      handleGameInput();
      break;
    case GAME_OVER:
      handleGameOver();
      break;
  }
}

void displayMenu() {
  lcd.setCursor(0, 0);
  lcd.print("Quizmaster Math Ed.");
  lcd.setCursor(0, 1);
  lcd.print("Green: Start");
  lcd.setCursor(0, 2);
  lcd.print("Red: Quit");
}

void handleMenuInput() {
  if (digitalRead(correctButtonPin) == LOW) {
    delay(50);  // Debounce delay
    if (digitalRead(correctButtonPin) == LOW) {  // Check again after delay
      currentState = CATEGORY;
    }
  } else if (digitalRead(incorrectButtonPin) == LOW) {
    delay(50);  // Debounce delay
    if (digitalRead(incorrectButtonPin) == LOW) {  // Check again after delay
      // Quit the program
      while (true) {
        lcd.clear();
        lcd.setCursor(0, 0);
        lcd.print("Quitting...");
        delay(1000); // Wait for 1 second before quitting
      }
    }
  }
}


void displayCategory() {
  lcd.setCursor(0, 0);
  lcd.print("Select Difficulty");
  lcd.setCursor(0, 1);
  lcd.print("Green: Basic");
  lcd.setCursor(0, 2);
  lcd.print("Red: Medium");
  lcd.setCursor(0, 3);
  lcd.print("Blue: Hard");
}

void handleCategoryInput() {
  if (digitalRead(correctButtonPin) == LOW) {
    delay(50);  // Debounce delay
    if (digitalRead(correctButtonPin) == LOW) {  // Check again after delay
      currentDifficulty = BASIC;
      currentState = IN_GAME;
    }
  } else if (digitalRead(incorrectButtonPin) == LOW) {
    delay(50);  // Debounce delay
    if (digitalRead(incorrectButtonPin) == LOW) {  // Check again after delay
      currentDifficulty = MEDIUM;
      currentState = IN_GAME;
    }
  } else if (digitalRead(blueButtonPin) == LOW) {
    delay(50);  // Debounce delay
    if (digitalRead(blueButtonPin) == LOW) {  // Check again after delay
      currentDifficulty = HARD;
      currentState = IN_GAME;
    }
  }
}

void generateQuestion() {
  int num1, num2;
  int operation = random(0, 4);  // Randomly select an operation (0 = addition, 1 = subtraction, 2 = multiplication, 3 = division)
  String question;

  switch (currentDifficulty) {
    case BASIC:
      num1 = random(1, 10);  // Generate a random number between 1 and 9
      num2 = random(1, 10);  // Generate another random number between 1 and 9
      operation = random(0, 2);  // Only addition and subtraction
      break;
    case MEDIUM:
      num1 = random(1, 100);  // Generate a random number between 1 and 99
      num2 = random(1, 100);  // Generate another random number between 1 and 99
      break;
    case HARD:
      num1 = random(1, 1000);  // Generate a random number between 1 and 999
      num2 = random(1, 1000);  // Generate another random number between 1 and 999
      break;
  }

  switch (operation) {
    case 0:
      question = String(num1) + " + " + String(num2) + " = ?";
      correctAnswer = num1 + num2;
      break;
    case 1:
      question = String(num1) + " - " + String(num2) + " = ?";
      correctAnswer = num1 - num2;
      break;
    case 2:
      question = String(num1) + " * " + String(num2) + " = ?";
      correctAnswer = num1 * num2;
      break;
    case 3:
      while (num2 == 0) {
        num2 = random(1, 100);  // Ensure no division by zero
      }
      question = String(num1) + " / " + String(num2) + " = ?";
      correctAnswer = num1 / num2;
      break;
  }

  displayedAnswer = random(correctAnswer - 10, correctAnswer + 10);  // Generate a random displayed answer close to the correct answer
  isAnswerCorrect = (displayedAnswer == correctAnswer);

  lcd.clear();
  lcd.setCursor(0, 0);
  lcd.print("Question:");
  lcd.setCursor(0, 1);
  lcd.print(question);

  // Display the question on Serial Monitor for debugging
  Serial.println(question);
  Serial.print("Correct answer: ");
  Serial.println(correctAnswer);

  lcd.setCursor(0, 2);
  lcd.print("Ans: " + String(displayedAnswer));
  
  // Display the answer on Serial Monitor for debugging
  Serial.print("Displayed answer: ");
  Serial.println(displayedAnswer);
}

void handleGameInput() {
  if (startUserResponseTimer()) {
    // Wait for user to press correct or incorrect button
    boolean answered = false;
    while (!answered && !timedOut) {
      if (digitalRead(correctButtonPin) == LOW) {
        delay(50);  // Debounce delay
        if (digitalRead(correctButtonPin) == LOW) {  // Check again after delay
          answered = true;
          checkAnswer(isAnswerCorrect);
        }
      } else if (digitalRead(incorrectButtonPin) == LOW) {
        delay(50);  // Debounce delay
        if (digitalRead(incorrectButtonPin) == LOW) {  // Check again after delay
          answered = true;
          checkAnswer(!isAnswerCorrect);
        }
      }
    }
  }

  if (timedOut) {
    lcd.clear();
    lcd.setCursor(0, 0);
    lcd.print("Time's up!");
    //tone(buzzerPin, 300, 500);  // Sound for timeout
    delay(1000);  // Display the result for 1 seconds
    noTone(buzzerPin);
    timedOut = false;  // Reset timeout flag
    incrementIncorrectBarGraph();
  }

  /*// Display correct count after showing the result of the answer
  lcd.clear();
  lcd.setCursor(0, 0);
  lcd.print("Correct: " + String(correctCount));
  lcd.setCursor(0, 1);
  lcd.print("Incorrect: " + String(incorrectCount));
  delay(3000);  // Wait for 3 seconds before generating the next question*/

  // Check for win/lose condition
  if (correctCount == 10) {
    lcd.clear();
    lcd.setCursor(0, 0);
    lcd.print("You Win!");
    playWinnerTune();
    currentState = GAME_OVER;
  } else if (incorrectCount == 10) {
    lcd.clear();
    lcd.setCursor(0, 0);
    lcd.print("You Lose!");
    playLoserTune();
    currentState = GAME_OVER;
  }
}

bool startUserResponseTimer() {
  timedOut = false;
  unsigned long startTime = millis();
  while (millis() - startTime < 5000) {  // 10 seconds countdown
    int timeLeft = 5 - (millis() - startTime) / 1000;
    display.showNumberDec(timeLeft, true);  // Display the countdown

    if (digitalRead(correctButtonPin) == LOW || digitalRead(incorrectButtonPin) == LOW) {
      return true;  // Exit the loop if button is pressed
    }
    delay(50);  // Check buttons every 50ms for better responsiveness
  }

  // If the loop exits without a button press, mark the answer as incorrect and increment the incorrect bar graph
  timedOut = true;  // Set timeout flag
  checkAnswer(false);  // Mark as incorrect
  return false;
}

void checkAnswer(boolean userSaidCorrect) {
  if (userSaidCorrect) {
    lcd.clear();
    lcd.setCursor(0, 0);
    lcd.print("Correct!");
    Serial.println("Correct!");
    correctCount++;  // Increment the correct answer count
    incrementCorrectBarGraph();
    // Custom sound for correct answer
    tone(buzzerPin, 1000, 200);  // 1000 Hz tone for 200 ms
    delay(200);
    tone(buzzerPin, 1500, 200);  // 1500 Hz tone for 200 ms
    delay(200);
    digitalWrite(correctLEDPin, HIGH);
    delay(1000);  // Display the result for 1 seconds
    digitalWrite(correctLEDPin, LOW);  // Turn off the LED
    noTone(buzzerPin);
  } else {
    lcd.clear();
    lcd.setCursor(0, 0);
    lcd.print("Wrong.");
    lcd.setCursor(0, 1);
    lcd.print("Answer: " + String(correctAnswer));
    Serial.print("Wrong. The correct answer is ");
    Serial.println(correctAnswer);
    incorrectCount++;  // Increment the incorrect answer count
    incrementIncorrectBarGraph();
    // Custom sound for incorrect answer
    tone(buzzerPin, 400, 300);  // 400 Hz tone for 300 ms
    delay(200);
    tone(buzzerPin, 300, 300);  // 300 Hz tone for 300 ms
    delay(200);
    digitalWrite(incorrectLEDPin, HIGH);
    delay(1000);  // Display the result for 1 seconds
    digitalWrite(incorrectLEDPin, LOW);  // Turn off the LED
    noTone(buzzerPin);
  }
  delay(1000);  // Display the result for 1 seconds
}

void incrementCorrectBarGraph() {
  if (correctCount <= 10) {
    digitalWrite(correctLEDPins[10 - correctCount], HIGH);  // Turn on the next LED from the end
  }
}

void incrementIncorrectBarGraph() {
  if (incorrectCount <= 10) {
    digitalWrite(incorrectLEDPins[10 - incorrectCount], HIGH);  // Turn on the next LED from the end
  }
}

// Function to play the winner tune
void playWinnerTune() {
  for (int i = 0; i < sizeof(winnerMelody) / sizeof(winnerMelody[0]); i++) {
    tone(buzzerPin, winnerMelody[i], winnerNoteDurations[i]); 
    delay(winnerNoteDurations[i] + 30); // Add a small delay between notes for better separation
  }
}

// Function to play the loser tune
void playLoserTune() {
  for (int i = 0; i < sizeof(loserMelody) / sizeof(loserMelody[0]); i++) {
    tone(buzzerPin, loserMelody[i], loserNoteDurations[i]); 
    delay(loserNoteDurations[i] + 30); // Add a small delay between notes for better separation
  }
}

void handleGameOver() {
  lcd.setCursor(0, 0);
  lcd.print("Press Grey to Reset");

  while (true) {
    if (digitalRead(resetButtonPin) == LOW) {
      delay(50);  // Debounce delay
      if (digitalRead(resetButtonPin) == LOW) {  // Check again after delay
        resetGame();
        break;
      }
    }
  }
}

void resetGame() {
  correctCount = 0;
  incorrectCount = 0;
  for (int i = 0; i < 10; i++) {
    digitalWrite(correctLEDPins[i], LOW);
    digitalWrite(incorrectLEDPins[i], LOW);
  }
  currentState = MENU;
}