#include <TM1637TinyDisplay.h>
#include <Keypad.h>
#include <Wire.h>
#include <U8g2lib.h>

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

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

// Initialize U8g2 display for a common 128x64 I2C OLED
U8G2_SSD1306_128X64_NONAME_F_HW_I2C oled(U8G2_R0, /* reset=*/ U8X8_PIN_NONE);

// 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};

// 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;
int timerSeconds = 0;
int remainingAttempts = 3;
bool awaitingReconnection = false;
bool gameStarted = false;
bool stage2Active = false;
int currentStep = 0;
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};

int mistakes = 0;
bool gameOver = false;
bool gameWon = false;

unsigned long previousMillis = 0;
long interval = 1000;
unsigned long previousStage2Millis = 0;

String enteredCode = "";

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

  pinMode(buttonIncrease, INPUT_PULLUP);
  pinMode(buttonDecrease, INPUT_PULLUP);
  pinMode(buttonStart, INPUT_PULLUP);

  pinMode(relayPin, OUTPUT);
  pinMode(buzzerPin, OUTPUT);

  for (int i = 0; i < 11; i++) {
    pinMode(wirePins[i], INPUT_PULLUP);
    pinMode(ledPins[i], OUTPUT);
    digitalWrite(ledPins[i], HIGH);
  }

  digitalWrite(relayPin, LOW);
  digitalWrite(buzzerPin, LOW);

  display.setBrightness(7);
  updateTimeDisplay();

  oled.begin();
  oled.clearBuffer();
  oled.setFont(u8g2_font_ncenB08_tr);
  oled.drawStr(0, 10, "   Airsoft Online Japan");
  oled.drawStr(0, 20, " ");
  oled.drawStr(0, 30, "     BAKUDAN Mk.III ");
  oled.drawStr(0, 40, " ");
  oled.drawStr(0, 50, "        Created by Ben ");
  oled.sendBuffer();
}

void increaseTime() {
  if (!gameStarted) {
    gameTime++;
    updateTimeDisplay();
    Serial.print("Time set to: ");
    Serial.print(gameTime);
    Serial.println(" minutes");
  }
}

void decreaseTime() {
  if (!gameStarted && gameTime > 0) {
    gameTime--;
    updateTimeDisplay();
    Serial.print("Time set to: ");
    Serial.print(gameTime);
    Serial.println(" minutes");
  }
}

void startGame() {
  if (!gameStarted && gameTime > 0) {
    gameStarted = true;
    timerSeconds = gameTime * 60;
    Serial.println("Game started!");
  }
}

void updateTimeDisplay() {
  int displayValue = gameTime * 100;
  display.showNumberDec(displayValue, true);
}

void updateCountdownDisplay() {
  int minutes = timerSeconds / 60;
  int seconds = timerSeconds % 60;
  int displayValue = (minutes * 100) + seconds;
  display.showNumberDec(displayValue, true);
}

void playVictorySound() {
  for (int i = 0; i < 3; i++) {
    digitalWrite(buzzerPin, HIGH);
    delay(200);
    digitalWrite(buzzerPin, LOW);
    delay(200);
  }
  oled.clearBuffer();
  oled.drawStr(0, 10, "Bomb Defused!");
  oled.sendBuffer();
}

void playDisarmSound() {
  for (int i = 0; i < 5; i++) {
    digitalWrite(buzzerPin, HIGH);
    delay(100 * (i + 1));
    digitalWrite(buzzerPin, LOW);
    delay(100);
  }
}

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("Circuit Unstable");
          oled.clearBuffer();
          oled.drawStr(0, 10, "Circuit Unstable");
          oled.sendBuffer();
          awaitingReconnection = true;
        }

      } else {
        handleMistake(i);
        wiresSkipped[i] = true;
        blinkLED(ledPins[i], 500);
      }
    }
  }
}

int findWireIndex(int pin) {
  for (int i = 0; i < 11; i++) {
    if (wirePins[i] == pin) return i;
  }
  return -1;
}

void handleMistake(int wireIndex) {
  if (wireIndex == 9 || wireIndex == 10) {
    Serial.println("Incorrect wire disconnected! Detonation!");
    activateRelay(3000);
    gameOver = true;
    triggerGameOver();
    return;
  }

  mistakes++;
  switch (mistakes) {
    case 1:
      interval = (long)(interval * 0.85);
      activateRelay(500);
      Serial.println("First mistake: Countdown speed increased by 15%.");
      break;
    case 2:
      timerSeconds = max((int)(timerSeconds * 0.75), 1);
      Serial.println("Second mistake: 25% time deduction.");
      break;
    case 3:
      timerSeconds = max(timerSeconds / 2, 1);
      Serial.println("Third mistake: Timer halved.");
      break;
    case 4:
      Serial.println("Fourth mistake: Game over, timer reset.");
      activateRelay(3000);
      gameOver = true;
      triggerGameOver();
      break;
  }

  updateCountdownDisplay();
}

void activateRelay(int duration) {
  digitalWrite(relayPin, HIGH);
  delay(duration);
  digitalWrite(relayPin, LOW);
}

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);
}

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);
}

void turnOffAllLEDs() {
  for (int i = 0; i < 11; i++) {
    digitalWrite(ledPins[i], LOW);
  }
}

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.clearBuffer();
  oled.drawStr(32, 32, "Game Over");
  oled.sendBuffer();
}

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};

  if (gameOver) return;

  if (timerSeconds == 0 && !gameOver) {
    gameOver = true;
    Serial.println("Stage 2 time's up! Game lost.");
    oled.clearBuffer();
    oled.drawStr(32, 32, "DETONATED");
    oled.sendBuffer();
    activateRelay(3000);
    triggerGameOver();
    return;
  }

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

    if (ledOn) {
      digitalWrite(ledPins[currentLED], LOW);
      ledOn = false;
      currentBlinkCount++;

      if (currentBlinkCount >= blinkCounts[currentLED]) {
        currentBlinkCount = 0;
        currentLED++;
        if (currentLED >= 11) {
          currentLED = 0;
          blinkStartMillis = currentMillis + 3000;
        }
      }
    } else {
      if (blinkCounts[currentLED] > 0) {
        digitalWrite(ledPins[currentLED], HIGH);
        ledOn = true;
      } else {
        currentLED++;
        if (currentLED >= 11) {
          currentLED = 0;
          blinkStartMillis = currentMillis + 3000;
        }
      }
    }
  }
}

void updateStage2State() {
  if (gameOver) {
    oled.clearBuffer();
    oled.drawStr(32, 32, "GAME OVER");
    oled.sendBuffer();
    return;
  }

  unsigned long currentMillis = millis();
  if (currentMillis - previousMillis >= interval) {
    previousMillis = currentMillis;
    if (timerSeconds > 0) {
      timerSeconds--;
      updateCountdownDisplay();
    } else {
      timerSeconds = 0;
      gameOver = true;
      oled.clearBuffer();
      oled.drawStr(32, 32, "DETONATED");
      oled.sendBuffer();
      activateRelay(3000);
      triggerGameOver();
      return;
    }
  }

  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.clearBuffer();
          oled.drawStr(32, 32, "Disarmed");
          oled.sendBuffer();
          return;
        } else {
          remainingAttempts--;
          Serial.println("Incorrect code! Try again.");
          enteredCode = "";

          if (remainingAttempts == 0) {
            Serial.println("No attempts left! Game over.");
            gameOver = true;
            oled.clearBuffer();
            oled.drawStr(32, 32, "DETONATED");
            oled.sendBuffer();
            activateRelay(3000);
            triggerGameOver();
            return;
          }
        }
      }
    } else {
      if (enteredCode.length() < 8) {
        enteredCode += key;
      }
    }
  }

  if (!gameWon) {
    oled.clearBuffer();
    oled.setCursor(0, 50);
    oled.drawStr(0, 50, "Press * to delete");
    oled.drawStr(0, 60, "Press # to enter");
    oled.setCursor(0, 20);
    oled.print("Attempts Left: ");
    oled.print(remainingAttempts);
    oled.setCursor(0, 35);
    oled.print("Entered: ");
    oled.print(enteredCode);
    oled.sendBuffer();
  }
}

void loop() {
  if (!gameStarted) {
    if (digitalRead(buttonIncrease) == LOW) {
      increaseTime();
      delay(300);
    }
    if (digitalRead(buttonDecrease) == LOW) {
      decreaseTime();
      delay(300);
    }
    if (digitalRead(buttonStart) == LOW) {
      startGame();
      delay(300);
    }
  }

  if (gameStarted && !gameOver && !gameWon) {
    checkWires();
    if (awaitingReconnection) {
      bool a2Connected = digitalRead(wirePins[6]) == HIGH;
      bool a3Connected = digitalRead(wirePins[7]) == HIGH;

      if (a2Connected && a3Connected) {
        awaitingReconnection = false;
        oled.clearBuffer();
        oled.drawStr(0, 10, "Stabilised");
        oled.sendBuffer();

        stage2Active = true;
        timerSeconds = 600;  // Set timer for stage 2
        previousStage2Millis = millis();
        previousMillis = millis();
        updateCountdownDisplay();
      }
    }


    unsigned long currentMillis = millis();
    if (currentMillis - previousMillis >= interval) {
      previousMillis = currentMillis;

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

        digitalWrite(buzzerPin, HIGH);
        delay(50);
        digitalWrite(buzzerPin, LOW);
      } else if (timerSeconds == 0 && currentStep < 6) {
        Serial.println("Time's up! Game lost.");
        activateRelay(3000);
        gameOver = true;
        triggerGameOver();
      }
    }
  }

  if (stage2Active && !gameWon) {
    blinkStage2Sequence();
    updateStage2State();
  }
}