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

// Define TM1637 display pins
#define CLK 8  // Clock pin
#define DIO 7  // Data In/Out pin

// Initialize TM1637TinyDisplay
TM1637TinyDisplay display(CLK, DIO);

// OLED display setup
#define SCREEN_WIDTH 128 // OLED display width, in pixels
#define SCREEN_HEIGHT 64 // OLED display height, in pixels
#define OLED_RESET    -1 // Reset pin # (or -1 if sharing Arduino reset pin)
Adafruit_SSD1306 oled(SCREEN_WIDTH, SCREEN_HEIGHT, &Wire, OLED_RESET);

// Pin assignments
const int buttonIncrease = 2;  // Button to increase time
const int buttonDecrease = 3;  // Button to decrease time
const int buttonStart = 4;     // Button to start the game
const int relayPin = 5;        // Relay control pin
const int buzzerPin = 6;       // Buzzer control pin

// LED pins for each wire
const int ledPins[11] = {22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32}; 

// Wire pins (connected in the correct sequence order)
const int wirePins[11] = {9, 10, 11, 12, A0, A1, A2, A3, A4, A5, A6}; 
const int correctOrder[9] = {9, 10, 11, 12, A0, A1, A2, A3, A4};  // Blue, Green, Yellow, Orange, Purple, White

// 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'}
};
byte rowPins[ROWS] = {30, 31, 32, 33};
byte colPins[COLS] = {34, 35, 36, 37};
Keypad keypad = Keypad(makeKeymap(keys), rowPins, colPins, ROWS, COLS);

// Variables
int gameTime = 0;  // Initial game time in minutes
int timerSeconds = 0;
int remainingAttempts = 3;  // Set the initial attempts allowed
bool gameStarted = false;
bool stage2Active = false;
int currentStep = 0;  // Keeps track of the current expected wire in the correct order
bool wiresDisconnected[11] = {false, false, false, false, false, false, false, false, false, false, false};
bool wiresSkipped[11] = {false, false, false, false, false, false, false, false, false, false, false};  // Tracks skipped wires

// Penalty variables
int mistakes = 0;
bool gameOver = false;
bool gameWon = false;

// Timer management
unsigned long previousMillis = 0;
long interval = 1000;  // 1-second interval initially
unsigned long previousStage2Millis = 0;

// Keypad input
String enteredCode = "";

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

  // Button inputs
  pinMode(buttonIncrease, INPUT_PULLUP);
  pinMode(buttonDecrease, INPUT_PULLUP);
  pinMode(buttonStart, INPUT_PULLUP);

  // Output pins
  pinMode(relayPin, OUTPUT);
  pinMode(buzzerPin, OUTPUT);

  // Wire inputs
  for (int i = 0; i < 11; i++) {
    pinMode(wirePins[i], INPUT_PULLUP); // Using INPUT_PULLUP to detect disconnect
    pinMode(ledPins[i], OUTPUT);        // LED pins as output
    digitalWrite(ledPins[i], HIGH);     // Turn on all LEDs initially
  }

  // Initialize outputs
  digitalWrite(relayPin, LOW);
  digitalWrite(buzzerPin, LOW);

  // Initialize the TM1637TinyDisplay
  display.setBrightness(7);  // Set brightness (0 to 7)
  updateTimeDisplay();

  // Initialize OLED display
  if (!oled.begin(SSD1306_SWITCHCAPVCC, 0x3C)) {
    Serial.println(F("SSD1306 allocation failed"));
    for (;;);
  }
  oled.clearDisplay();
  oled.setTextSize(3); // Set size to 20% larger
  oled.setTextColor(SSD1306_WHITE);
  oled.setCursor((SCREEN_WIDTH - (6 * 3 * 6)) / 2, (SCREEN_HEIGHT - 3 * 8) / 2);  // Center text
  oled.setTextWrap(false);
  oled.println("AOJ");
  oled.display();
}

// Function to increase game time by 1 minute
void increaseTime() {
  if (!gameStarted) {
    gameTime++;
    updateTimeDisplay();
    Serial.print("Time set to: ");
    Serial.print(gameTime);
    Serial.println(" minutes");
  }
}

// Function to decrease game time by 1 minute
void decreaseTime() {
  if (!gameStarted && gameTime > 0) {
    gameTime--;
    updateTimeDisplay();
    Serial.print("Time set to: ");
    Serial.print(gameTime);
    Serial.println(" minutes");
  }
}

// Function to handle game timer
void startGame() {
  if (!gameStarted && gameTime > 0) {
    gameStarted = true;
    timerSeconds = gameTime * 60;  // Convert minutes to seconds
    Serial.println("Game started!");
  }
}

// Function to update the TM1637TinyDisplay during game time setting
void updateTimeDisplay() {
  int displayValue = gameTime * 100; // Display only minutes when setting time
  display.showNumberDec(displayValue, true);  // Display as MM:00
}

// Function to update the TM1637TinyDisplay during countdown
void updateCountdownDisplay() {
  int minutes = timerSeconds / 60;
  int seconds = timerSeconds % 60;
  int displayValue = (minutes * 100) + seconds;
  display.showNumberDec(displayValue, true);  // Display as MM:SS with leading zeros
}

// Function to play victory sound
void playVictorySound() {
  for (int i = 0; i < 3; i++) {
    digitalWrite(buzzerPin, HIGH);
    delay(200);
    digitalWrite(buzzerPin, LOW);
    delay(200);
  }
  oled.clearDisplay();
  oled.setTextSize(1);
  oled.setCursor(0, 0);
  oled.println("Bomb Defused!");
  oled.display();
}

// Function to play disarm sound
void playDisarmSound() {
  for (int i = 0; i < 5; i++) {
    digitalWrite(buzzerPin, HIGH);
    delay(100 * (i + 1));  // Increasing delay to create a power-down effect
    digitalWrite(buzzerPin, LOW);
    delay(100);
  }
}

// Function to check wire disconnection
void checkWires() {
  if (gameOver || gameWon) return;

  for (int i = 0; i < 11; i++) {
    if (!digitalRead(wirePins[i]) && !wiresDisconnected[i]) {
      wiresDisconnected[i] = true;

      if (currentStep < 9 && wirePins[i] == correctOrder[currentStep]) {
        Serial.println("Correct wire disconnected.");
        flashLED(ledPins[i], 3);
        currentStep++;

        while (currentStep < 9 && wiresSkipped[findWireIndex(correctOrder[currentStep])]) {
          currentStep++;
        }

        if (currentStep == 9) {
          timerSeconds = 0;
          turnOffAllLEDs();
          playVictorySound();
          Serial.println("Game won!");
          stage2Active = true;
          timerSeconds = 600;
          previousStage2Millis = millis();
          updateCountdownDisplay();
        }
      } else {
        handleMistake(i);
        wiresSkipped[i] = true;
        blinkLED(ledPins[i], 500);
      }
    }
  }
}

// Helper function to find the index of a wire pin in the wirePins array
int findWireIndex(int pin) {
  for (int i = 0; i < 11; i++) {
    if (wirePins[i] == pin) return i;
  }
  return -1;  // Should not happen if the pins are configured correctly
}

// Function to handle mistakes
void handleMistake(int wireIndex) {
  if (wireIndex == 9 || wireIndex == 10) {
    // Detonation condition - immediate activation of relay
    Serial.println("Incorrect wire disconnected! Detonation!");
    activateRelay(3000); // Relay on for 3 seconds immediately
    gameOver = true;
    triggerGameOver();
    return;
  }

  // Apply penalty for the mistake
  mistakes++;
  switch (mistakes) {
    case 1:
      // First mistake speeds up countdown by 15%, but does not reduce time
      interval = (long)(interval * 0.85);  // 15% faster countdown
      activateRelay(500); // Relay on for 0.5 seconds
      Serial.println("First mistake: Countdown speed increased by 15%.");
      break;
    case 2:
      // Second mistake deducts 25% of remaining time
      timerSeconds = max((int)(timerSeconds * 0.75), 1);
      Serial.println("Second mistake: 25% time deduction.");
      break;
    case 3:
      // Third mistake halves the remaining time
      timerSeconds = max(timerSeconds / 2, 1);
      Serial.println("Third mistake: Timer halved.");
      break;
    case 4:
      // Fourth mistake triggers game over
      Serial.println("Fourth mistake: Game over, timer reset.");
      activateRelay(3000); // Relay on for 3 seconds immediately
      gameOver = true;
      triggerGameOver();
      break;
  }

  // Update the display after applying the penalty
  updateCountdownDisplay();
}

// Function to activate relay for a specified duration
void activateRelay(int duration) {
  digitalWrite(relayPin, HIGH);
  delay(duration);
  digitalWrite(relayPin, LOW);
}

// Function to flash an LED for a given number of times
void flashLED(int pin, int times) {
  for (int i = 0; i < times; i++) {
    digitalWrite(pin, LOW);
    delay(250);
    digitalWrite(pin, HIGH);
    delay(250);
  }
  digitalWrite(pin, LOW);  // Turn off LED after flashing
}

// Function to blink an LED every 0.5 seconds indefinitely
void blinkLED(int pin, int duration) {
  unsigned long blinkStart = millis();
  while (millis() - blinkStart < duration) {
    digitalWrite(pin, LOW);
    delay(250);
    digitalWrite(pin, HIGH);
    delay(250);
  }
  digitalWrite(pin, LOW);  // Turn off LED after blinking
}

// Function to turn off all LEDs
void turnOffAllLEDs() {
  for (int i = 0; i < 11; i++) {
    digitalWrite(ledPins[i], LOW);
  }
}

// Function to trigger game over behavior
void triggerGameOver() {
  turnOffAllLEDs();
  for (int i = 0; i < 10; i++) {
    for (int j = 0; j < 11; j++) {
      digitalWrite(ledPins[j], HIGH);
    }
    delay(250);
    for (int j = 0; j < 8; j++) {
      digitalWrite(ledPins[j], LOW);
    }
    delay(250);
  }
  oled.clearDisplay();
  oled.setCursor(0, 0);
  oled.println("Game Over");
  oled.display();
}

// Function to blink LEDs to reveal the passcode during Stage 2
void blinkStage2Sequence() {
  static unsigned long blinkStartMillis = 0;
  static int currentLED = 0;
  static int currentBlinkCount = 0;
  static bool ledOn = false;

  unsigned long currentMillis = millis();
  int blinkCounts[8] = {1, 8, 8, 3, 0, 9, 0, 2};

  // Ensure the game is not over before continuing LED blink sequence
  if (gameOver) return;

  // Check if timer has reached zero before continuing the blinking sequence
  if (timerSeconds == 0 && !gameOver) {
    // Trigger detonation immediately
    gameOver = true;
    Serial.println("Stage 2 time's up! Game lost.");
    oled.clearDisplay();
    oled.setTextSize(2);
    oled.setCursor((SCREEN_WIDTH - (8 * 6)) / 2, (SCREEN_HEIGHT - 3 * 8) / 2);  // Center text
    oled.println("DETONATED");
    oled.display();
    activateRelay(3000);  // Activate relay for 3 seconds immediately to indicate detonation
    triggerGameOver();
    return;
  }

  if (currentMillis - blinkStartMillis >= (ledOn ? 250 : 500)) {
    blinkStartMillis = currentMillis;

    if (ledOn) {
      // Turn off the LED
      digitalWrite(ledPins[currentLED], LOW);
      ledOn = false;
      currentBlinkCount++;

      // Move to next LED if done with current LED's blinks
      if (currentBlinkCount >= blinkCounts[currentLED]) {
        currentBlinkCount = 0;
        currentLED++;
        if (currentLED >= 11) {
          currentLED = 0;
          blinkStartMillis = currentMillis + 3000; // Short delay before repeating the sequence
        }
      }
    } else {
      // Turn on the LED if there are blinks remaining
      if (blinkCounts[currentLED] > 0) {
        digitalWrite(ledPins[currentLED], HIGH);
        ledOn = true;
      } else {
        // If no blinks are required for the current LED, skip to the next one
        currentLED++;
        if (currentLED >= 11) {
          currentLED = 0;
          blinkStartMillis = currentMillis + 3000; // Short delay before repeating the sequence
        }
      }
    }
  }
}

// Function to update timer, OLED display, and keypad during LED blinking in Stage 2
void updateStage2State() {
  if (gameOver) {
    // Stop all actions and keep the "Game Over" message on the screen
    oled.clearDisplay();
    oled.setTextSize(2);
    oled.setCursor((SCREEN_WIDTH - (8 * 6)) / 2, (SCREEN_HEIGHT - 3 * 8) / 2);  // Center text
    oled.println("GAME OVER");
    oled.display();
    return;  // Exit to prevent further actions
  }

  unsigned long currentMillis = millis();
  if (currentMillis - previousMillis >= interval) {
    previousMillis = currentMillis;
    if (timerSeconds > 0) {
      timerSeconds--;
      updateCountdownDisplay();
    } else {
      timerSeconds = 0;
      Serial.println("Stage 2 time's up! Game lost.");
      gameOver = true;

      // Trigger game over and display "DETONATED" on OLED
      oled.clearDisplay();
      oled.setTextSize(2);
      oled.setCursor((SCREEN_WIDTH - (8 * 6)) / 2, (SCREEN_HEIGHT - 3 * 8) / 2);  // Center text
      oled.println("DETONATED");
      oled.display();

      activateRelay(3000);  // Activate relay for 3 seconds immediately to indicate detonation
      triggerGameOver();  // Trigger game over sequence like flashing LEDs
      return;  // Exit to prevent further actions
    }
  }

  // Handle keypad input only if game is not over
  char key = keypad.getKey();
  if (key) {
    if (key == '*') {
      if (enteredCode.length() > 0) {
        enteredCode.remove(enteredCode.length() - 1);
      }
    } else if (key == '#') {
      if (enteredCode.length() == 8) {
        if (enteredCode == "18830902") {
          Serial.println("Code correct! Bomb disarmed!");
          playDisarmSound();
          gameWon = true;
          stage2Active = false;
          oled.clearDisplay();
          delay(500); // Add a small delay to ensure display is fully cleared
          oled.setTextSize(2);
          oled.setCursor((SCREEN_WIDTH - (8 * 6)) / 2, (SCREEN_HEIGHT - 3 * 8) / 2);  // Center text
          oled.println("Disarmed");
          oled.display();
          return; // Ensure no other content is updated after "Disarmed" is shown
        } else {
          remainingAttempts--;  // Decrement attempts if code is wrong
          Serial.println("Incorrect code! Try again.");
          enteredCode = "";

          // Trigger game over if no attempts are left
          if (remainingAttempts == 0) {
            Serial.println("No attempts left! Game over.");
            gameOver = true;
            oled.clearDisplay();
            oled.setTextSize(2);
            oled.setCursor((SCREEN_WIDTH - (8 * 6)) / 2, (SCREEN_HEIGHT - 3 * 8) / 2);  // Center text
            oled.println("GAME OVER");
            oled.display();
            activateRelay(3000);  // Activate relay for 3 seconds to indicate detonation
            triggerGameOver();
            return;  // Exit to prevent further actions
          }
        }
      }
    } else {
      if (enteredCode.length() < 8) {
        enteredCode += key;
      }
    }
  }

  // Update OLED display to show remaining attempts and keypad input
  if (!gameWon) {
    oled.clearDisplay();
    oled.setCursor(0, 50); // Move to the bottom
    oled.println("Press * to delete");
    oled.println("Press # to enter");
    oled.setCursor(0, 20); // Display remaining attempts in the center
    oled.setTextSize(1);
    oled.print("Attempts Left: ");
    oled.print(remainingAttempts);
    oled.setCursor(0, 35); // Show entered code below attempts
    oled.print("Entered: ");
    oled.print(enteredCode);
    oled.display();
  }
}


void loop() {
  // Handle buttons before the game starts
  if (!gameStarted) {
    if (digitalRead(buttonIncrease) == LOW) {
      increaseTime();
      delay(300);  // Debounce delay
    }
    if (digitalRead(buttonDecrease) == LOW) {
      decreaseTime();
      delay(300);  // Debounce delay
    }
    if (digitalRead(buttonStart) == LOW) {
      startGame();
      delay(300);  // Debounce delay
    }
  }

  // Game running
  if (gameStarted && !gameOver && !gameWon) {
    checkWires();

    // Timer countdown (Stage 1 and Stage 2)
    unsigned long currentMillis = millis();
    if (currentMillis - previousMillis >= interval) {
      previousMillis = currentMillis;

      if (timerSeconds > 0) {
        timerSeconds--;
        updateCountdownDisplay();

        // Buzzer beeps in sync with the faster countdown interval
        digitalWrite(buzzerPin, HIGH);
        delay(50);  // Short beep
        digitalWrite(buzzerPin, LOW);
      } else if (timerSeconds == 0 && currentStep < 6) {
        Serial.println("Time's up! Game lost.");
        activateRelay(3000); // Relay on for 3 seconds when time is up
        gameOver = true;
        triggerGameOver();
      }
    }
  }

  // Stage 2 - LED sequence hints and keypad input
  if (stage2Active && !gameWon) {
    blinkStage2Sequence(); // Blink LEDs in sequence to reveal passcode
    updateStage2State();   // Update Stage 2 timer and handle input
  }
}