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

// OLED display parameters
#define SCREEN_WIDTH 128
#define SCREEN_HEIGHT 64
#define OLED_RESET    -1
Adafruit_SSD1306 display(SCREEN_WIDTH, SCREEN_HEIGHT, &Wire, OLED_RESET);

// Shift register parameters
#define NUMBER_OF_SHIFT_REGISTERS 4
#define DATA_PIN 12
#define CLOCK_PIN 5
#define LATCH_PIN 13
ShiftRegister74HC595<NUMBER_OF_SHIFT_REGISTERS> sr(DATA_PIN, CLOCK_PIN, LATCH_PIN);

// Button pins
#define BUTTON_UP     17
#define BUTTON_DOWN   18
#define BUTTON_LEFT   15
#define BUTTON_RIGHT  16
#define BUTTON_SELECT 14

// Game state variables
int currentPlayer = 1;
int gameBoard[3][3] = {0}; // 0: empty, 1: Player 1, 2: Player 2
int cursorX = 0;
int cursorY = 0;
int player1Wins = 0;
int player2Wins = 0;

// Initialization
void setup() {
  // Initialize serial for debugging
  Serial.begin(115200);

  // Initialize the display
  if(!display.begin(SSD1306_SWITCHCAPVCC, 0x3C)) {
    Serial.println(F("SSD1306 allocation failed"));
    for(;;);
  }
  display.display();
  delay(2000);
  display.clearDisplay();
  display.setTextSize(1);
  display.setTextColor(SSD1306_WHITE);
  display.setCursor(0, 0);
  display.print("Welcome to Tic-Tac-Toe");
  display.display();
  delay(2000);

  // Initialize shift registers
  sr.setAllLow();

  // Initialize buttons
  pinMode(BUTTON_UP, INPUT_PULLUP);
  pinMode(BUTTON_DOWN, INPUT_PULLUP);
  pinMode(BUTTON_LEFT, INPUT_PULLUP);
  pinMode(BUTTON_RIGHT, INPUT_PULLUP);
  pinMode(BUTTON_SELECT, INPUT_PULLUP);

  // Initialize game state
  resetGame();
}

// Reset game state
void resetGame() {
  memset(gameBoard, 0, sizeof(gameBoard));
  cursorX = 0;
  cursorY = 0;
  currentPlayer = 1;
  updateDisplay();
}

// Update OLED display
void updateDisplay() {
  display.clearDisplay();
  display.setCursor(0, 0);
  display.print("Player ");
  display.print(currentPlayer);
  display.println(" Turn");

  for (int y = 0; y < 3; y++) {
    for (int x = 0; x < 3; x++) {
      display.setCursor(x * 40, y * 20 + 20);
      if (gameBoard[y][x] == 0) {
        display.print("_");
      } else if (gameBoard[y][x] == 1) {
        display.print("X");
      } else {
        display.print("O");
      }
    }
  }
  display.display();
}

// Main loop
void loop() {
  // Handle button input
  if (digitalRead(BUTTON_UP) == LOW) {
    cursorY = max(0, cursorY - 1);
    delay(200); // Debounce delay
  }
  if (digitalRead(BUTTON_DOWN) == LOW) {
    cursorY = min(2, cursorY + 1);
    delay(200);
  }
  if (digitalRead(BUTTON_LEFT) == LOW) {
    cursorX = max(0, cursorX - 1);
    delay(200);
  }
  if (digitalRead(BUTTON_RIGHT) == LOW) {
    cursorX = min(2, cursorX + 1);
    delay(200);
  }
  if (digitalRead(BUTTON_SELECT) == LOW) {
    if (gameBoard[cursorY][cursorX] == 0) {
      gameBoard[cursorY][cursorX] = currentPlayer;
      if (checkWin()) {
        if (currentPlayer == 1) {
          player1Wins++;
        } else {
          player2Wins++;
        }
        displayWinner();
        delay(2000);
        resetGame();
      } else if (isDraw()) {
        displayDraw();
        delay(2000);
        resetGame();
      } else {
        currentPlayer = 3 - currentPlayer; // Toggle between 1 and 2
      }
    }
    delay(200);
  }
  updateDisplay();
  updateLEDs();
}

void updateLEDs() {
  // Clear all LEDs
  sr.setAllLow();

  // Set LED colors based on game state
  for (int y = 0; y < 3; y++) {
    for (int x = 0; x < 3; x++) {
      int ledIndex = y * 3 + x;
      if (gameBoard[y][x] == 1) {
        // Player 1 (Red-Purple)
        sr.set(ledIndex * 3, HIGH);     // Red
        sr.set(ledIndex * 3 + 2, HIGH); // Blue
      } else if (gameBoard[y][x] == 2) {
        // Player 2 (Blue-Green)
        sr.set(ledIndex * 3 + 1, HIGH); // Green
        sr.set(ledIndex * 3 + 2, HIGH); // Blue
      }

      // Blinking effect for selected square
      if (cursorX == x && cursorY == y) {
        static bool blinkState = false;
        if (blinkState) {
          if (currentPlayer == 1) {
            sr.set(ledIndex * 3, HIGH);     // Red
            sr.set(ledIndex * 3 + 2, HIGH); // Blue
          } else {
            sr.set(ledIndex * 3 + 1, HIGH); // Green
            sr.set(ledIndex * 3 + 2, HIGH); // Blue
          }
        }
        blinkState = !blinkState;
        
  Serial.println(ledIndex);
      }
    }
  }
}

// Check for win
bool checkWin() {
  // Implement win checking logic
  return false;
}

// Check for draw
bool isDraw() {
  for (int y = 0; y < 3; y++) {
    for (int x = 0; x < 3; x++) {
      if (gameBoard[y][x] == 0) {
        return false;
      }
    }
  }
  return true;
}

// Display winner
void displayWinner() {
  display.clearDisplay();
  display.setCursor(0, 0);
  display.print("Player ");
  display.print(currentPlayer);
  display.println(" Wins!");
  display.setCursor(0, 20);
  display.print("P1 Wins: ");
  display.print(player1Wins);
  display.setCursor(0, 40);
  display.print("P2 Wins: ");
  display.print(player2Wins);
  display.display();
}

// Display draw
void displayDraw() {
  display.clearDisplay();
  display.setCursor(0, 0);
  display.println("It's a Draw!");
  display.display();
}