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

// OLED display width and height
#define SCREEN_WIDTH 128
#define SCREEN_HEIGHT 64

// OLED reset pin (you can change it to another pin or set to -1 if not used)
#define OLED_RESET -1

// Create display object
Adafruit_SSD1306 display(SCREEN_WIDTH, SCREEN_HEIGHT, &Wire, OLED_RESET);

// Keypad constants
const byte ROWS = 4;
const byte COLS = 4;

// Array to represent keys on keypad
char hexaKeys[ROWS][COLS] = {
  {'1', '2', '3', 'A'},
  {'4', '5', '6', 'B'},
  {'7', '8', '9', 'C'},
  {'*', '0', '#', 'D'},
};

// Connections to the arduino
byte rowPins[ROWS] = {9, 8, 7, 6};
byte colPins[COLS] = {5, 4, 3, 2};

// Create keypad object
Keypad customKeypad = Keypad(makeKeymap(hexaKeys), rowPins, colPins, ROWS, COLS);

// Game Constants
const int BUTTON_ROWS = 3;
const int BUTTON_COLS = 3;

// Player Constants
const int PLAYER_NONE = 0;
const int PLAYER_1 = 1;
const int PLAYER_2 = 2;

// Game Variables
int board[BUTTON_ROWS][BUTTON_COLS];
int currentPlayer;
bool gameOver;
int player1Points = 0;
int player2Points = 0;

void setup() {
  // Initialize OLED display
  if (!display.begin(SSD1306_SWITCHCAPVCC, 0x3C)) {
    Serial.println(F("SSD1306 allocation failed"));
    for (;;);
  }
  display.display();
  delay(2000);
  display.clearDisplay();
  display.display();
  
  // Initialize Game Variables
  currentPlayer = PLAYER_1;
  gameOver = false;
  resetBoard();
  updateDisplay();
}

void loop() {
  // Check for Button Press
  char keypress = customKeypad.getKey();

  if (keypress == '0') {
    setup();
  }

  if (keypress == '#') {
    resetGame();
  }

  if (!gameOver) {
    if (keypress <= '9' && keypress >= '1') {
      int number = keypress - '1';

      int i = number / 3;
      int j = number % 3;

      if (board[i][j] == PLAYER_NONE) {
        board[i][j] = currentPlayer;
        updateDisplay();
        checkWinCondition();
        switchPlayers();
      }
      delay(200); // Button Rebound
    }
  }
}

void resetGame() {
  resetBoard();
  currentPlayer = PLAYER_1;
  gameOver = false;
  updateDisplay();
}

void resetBoard() {
  for (int i = 0; i < BUTTON_ROWS; i++) {
    for (int j = 0; j < BUTTON_COLS; j++) {
      board[i][j] = PLAYER_NONE;
    }
  }
}

void updateDisplay() {
  display.clearDisplay();
  display.setTextSize(1);
  display.setTextColor(SSD1306_WHITE);
  display.setCursor(0, 0);
  display.print("Player ");
  display.print(currentPlayer == PLAYER_1 ? "1" : "2");
  
  for (int i = 0; i < BUTTON_ROWS; i++) {
    for (int j = 0; j < BUTTON_COLS; j++) {
      display.setCursor(j * 20, (i + 1) * 16);
      display.print("|");
      display.print(symbolForPlayer(board[i][j]));
    }
    display.print("|");
  }
  display.display();
}

char symbolForPlayer(int player) {
  if (player == PLAYER_1) {
    return 'X';
  } else if (player == PLAYER_2) {
    return 'O';
  }
  return ' ';
}

void switchPlayers() {
  currentPlayer = (currentPlayer == PLAYER_1) ? PLAYER_2 : PLAYER_1;
}

void checkWinCondition() {
  // Check Rows
  for (int i = 0; i < BUTTON_ROWS; i++) {
    if (board[i][0] != PLAYER_NONE && board[i][0] == board[i][1] && board[i][1] == board[i][2]) {
      gameOver = true;
      updatePoints(board[i][0]);
      displayWinMessage(board[i][0]);
      delay(3000);
      displayPoints();
      delay(3000);
      resetGame();
      return;
    }
  }

  // Check Columns
  for (int j = 0; j < BUTTON_COLS; j++) {
    if (board[0][j] != PLAYER_NONE && board[0][j] == board[1][j] && board[1][j] == board[2][j]) {
      gameOver = true;
      updatePoints(board[0][j]);
      displayWinMessage(board[0][j]);
      delay(3000);
      displayPoints();
      delay(3000);
      resetGame();
      return;
    }
  }

  // Check Diagonals
  if (board[0][0] != PLAYER_NONE && board[0][0] == board[1][1] && board[1][1] == board[2][2]) {
    gameOver = true;
    updatePoints(board[0][0]);
    displayWinMessage(board[0][0]);
    delay(3000);
    displayPoints();
    delay(3000);
    resetGame();
    return;
  }

  if (board[0][2] != PLAYER_NONE && board[0][2] == board[1][1] && board[1][1] == board[2][0]) {
    gameOver = true;
    updatePoints(board[0][2]);
    displayWinMessage(board[0][2]);
    delay(3000);
    displayPoints();
    delay(3000);
    resetGame();
    return;
  }

  // Check Draw
  bool draw = true;
  for (int i = 0; i < BUTTON_ROWS; i++) {
    for (int j = 0; j < BUTTON_COLS; j++) {
      if (board[i][j] == PLAYER_NONE) {
        draw = false;
        break;
      }
    }
    if (!draw) {
      break;
    }
  }
  if (draw) {
    gameOver = true;
    display.clearDisplay();
    display.setTextSize(2);
    display.setCursor((SCREEN_WIDTH - 12 * 7) / 2, (SCREEN_HEIGHT - 16) / 2); // Centered
    display.print("Draw!");
    display.display();
    delay(3000);
    resetGame();
  }
}

void updatePoints(int player) {
  if (player == PLAYER_1) {
    player1Points += 50;
  } else if (player == PLAYER_2) {
    player2Points += 50;
  }
}

void displayWinMessage(int player) {
  display.clearDisplay();
  display.setTextSize(2);
  display.setTextColor(SSD1306_WHITE);
  display.setCursor((SCREEN_WIDTH - 12 * 7) / 2, (SCREEN_HEIGHT - 16) / 2); // Centered
  display.print("Player ");
  display.print(player == PLAYER_1 ? "1" : "2");
  display.print(" Wins!");
  display.display();
}

void displayPoints() {
  display.clearDisplay();
  display.setTextSize(1);
  display.setTextColor(SSD1306_WHITE);
  display.setCursor(0, 0);
  display.print("Points:");
  display.setCursor(0, 16);
  display.print("Player 1: ");
  display.print(player1Points);
  display.setCursor(0, 32);
  display.print("Player 2: ");
  display.print(player2Points);
  display.display();
}