#include <U8g2lib.h>
#include <Bounce2.h>

#define WIDTH 128
#define HEIGHT 64
#define START_BUTTON_PIN 12    // GPIO12 for start button
#define PAUSE_PLAY_BUTTON_PIN 13   // GPIO13 for pause/play button (changed from 14)
#define RESET_BUTTON_PIN 27     // GPIO27 for reset button

U8G2_SSD1306_128X64_NONAME_F_HW_I2C u8g2(U8G2_R0, /* reset=/ U8X8_PIN_NONE, / clock=/ SCL, / data=/ SDA, / reset=/ U8X8_PIN_NONE, / address=*/ 0x3C);

bool cells[WIDTH][HEIGHT];
bool newCells[WIDTH][HEIGHT];
bool paused = true; // Start with the game paused
bool started = false; // Game not started initially
bool displayInitialMessage = true; // Flag to indicate whether to display initial message

Bounce startButton = Bounce(); // Bounce object for start button
Bounce pausePlayButton = Bounce(); // Bounce object for pause/play button
Bounce resetButton = Bounce(); // Bounce object for reset button

// RGB LED Pins
const int RED_PIN = 4;     // GPIO15 for red channel of RGB LED
const int GREEN_PIN = 16;   // GPIO32 for green channel of RGB LED
const int BLUE_PIN = 17;    // GPIO33 for blue channel of RGB LED

void setup() {
  Serial.begin(115200);
  u8g2.begin();
  u8g2.setFont(u8g2_font_ncenB08_tr); // Set font to bold
  initializeDisplay(); // Clear and initialize display
  displayMessage("John Conway's\nGame of Life"); // Display initial message
  
  randomSeed(analogRead(0));
  
  // Setup RGB LED pins
  pinMode(RED_PIN, OUTPUT);
  pinMode(GREEN_PIN, OUTPUT);
  pinMode(BLUE_PIN, OUTPUT);
  
  // Initially set RGB LED to blue (waiting state)
  digitalWrite(RED_PIN, LOW);
  digitalWrite(GREEN_PIN, LOW);
  digitalWrite(BLUE_PIN, HIGH);
  
  // Setup buttons
  pinMode(START_BUTTON_PIN, INPUT_PULLUP); // Set start button pin as input with internal pull-up resistor
  pinMode(PAUSE_PLAY_BUTTON_PIN, INPUT_PULLUP); // Set pause/play button pin as input with internal pull-up resistor (changed to GPIO13)
  pinMode(RESET_BUTTON_PIN, INPUT_PULLUP); // Set reset button pin as input with internal pull-up resistor
  
  startButton.attach(START_BUTTON_PIN);
  startButton.interval(5); // Interval in ms
  pausePlayButton.attach(PAUSE_PLAY_BUTTON_PIN);
  pausePlayButton.interval(5); // Interval in ms
  resetButton.attach(RESET_BUTTON_PIN);
  resetButton.interval(5); // Interval in ms
}

void loop() {
  startButton.update();
  pausePlayButton.update();
  resetButton.update();
  
  // Check if the start button is pressed
  if (startButton.fell() && !started) {
    initializeRandomCells();
    paused = false; // Start the game
    started = true;
    displayInitialMessage = false; // Turn off initial message display
    displayMessage("Game Started");
    Serial.println("Game Started");
    
    // Set RGB LED to green (playing state)
    digitalWrite(RED_PIN, LOW);
    digitalWrite(GREEN_PIN, HIGH);
    digitalWrite(BLUE_PIN, LOW);
  }
  
  // Check if the pause/play button is pressed
  if (pausePlayButton.fell()) {
    paused = !paused; // Toggle the paused state
    
    // Display debug information to serial monitor
    if (paused) {
      Serial.println("Paused");
      
      // Set RGB LED to yellow (paused state)
      digitalWrite(RED_PIN, HIGH);
      digitalWrite(GREEN_PIN, HIGH);
      digitalWrite(BLUE_PIN, LOW);
    } else {
      Serial.println("Unpaused");
      
      // Set RGB LED to green (playing state)
      digitalWrite(RED_PIN, LOW);
      digitalWrite(GREEN_PIN, HIGH);
      digitalWrite(BLUE_PIN, LOW);
    }
  }
  
  // Check if the reset button is pressed
  if (resetButton.fell()) {
    initializeDisplay();
    paused = true;
    started = false;
    displayInitialMessage = true; // Allow displaying initial message again
    displayMessage("John Conway's\nGame of Life");
    Serial.println("Waiting for start button");
    
    // Set RGB LED to blue (waiting state)
    digitalWrite(RED_PIN, LOW);
    digitalWrite(GREEN_PIN, LOW);
    digitalWrite(BLUE_PIN, HIGH);
  }
  
  // Update the game if it's not paused and started
  if (!paused && started) {
    updateGame();
    drawGame(); // Ensure the game is redrawn after updating
  }
  
  // Draw "John Conway's\nGame of Life" message if it's the initial display or waiting for start button
  if (displayInitialMessage && !started) {
    u8g2.drawStr(0, 10, "John Conway's");
    u8g2.drawStr(0, 20, "Game of Life");
    u8g2.sendBuffer();
  }
  
  // Draw game grid when paused
  if (paused && started) {
    drawGame();
  }
  
  // Delay to control frame rate (adjust as needed)
  delay(50);
}

void initializeRandomCells() {
  for (int x = 0; x < WIDTH; x++) {
    for (int y = 0; y < HEIGHT; y++) {
      cells[x][y] = random(2); // Randomly initialize cells to be alive or dead
    }
  }
}

void updateGame() {
  // Copy cells to newCells
  memcpy(newCells, cells, sizeof(cells));
  
  // Update newCells based on Conway's Game of Life rules
  for (int x = 0; x < WIDTH; x++) {
    for (int y = 0; y < HEIGHT; y++) {
      int liveNeighbors = countLiveNeighbors(x, y);
      
      if (cells[x][y]) {
        if (liveNeighbors < 2 || liveNeighbors > 3) {
          newCells[x][y] = false; // Any live cell with fewer than two live neighbors dies, as if by underpopulation.
        }
      } else {
        if (liveNeighbors == 3) {
          newCells[x][y] = true; // Any dead cell with exactly three live neighbors becomes a live cell, as if by reproduction.
        }
      }
    }
  }
  
  // Copy newCells back to cells
  memcpy(cells, newCells, sizeof(cells));
}

int countLiveNeighbors(int x, int y) {
  int count = 0;
  
  for (int i = -1; i <= 1; i++) {
    for (int j = -1; j <= 1; j++) {
      if (i == 0 && j == 0) continue;
      int nx = x + i;
      int ny = y + j;
      
      // Wrap around edges (toroidal grid)
      if (nx < 0) nx = WIDTH - 1;
      if (nx >= WIDTH) nx = 0;
      if (ny < 0) ny = HEIGHT - 1;
      if (ny >= HEIGHT) ny = 0;
      
      count += cells[nx][ny];
    }
  }
  
  return count;
}

void drawGame() {
  u8g2.clearBuffer();
  
  // Draw cells based on their state
  for (int x = 0; x < WIDTH; x++) {
    for (int y = 0; y < HEIGHT; y++) {
      if (cells[x][y]) {
        u8g2.drawPixel(x, y);
      }
    }
  }
  
  u8g2.sendBuffer();
}

void displayMessage(const char* message) {
  // Display message on serial monitor only
  Serial.println(message);
}

void initializeDisplay() {
  u8g2.clearBuffer();
  u8g2.sendBuffer();
}