#include <Adafruit_NeoPixel.h>
#include <Adafruit_SSD1306.h>
#include <Wire.h>

// Constants
#define SCREEN_WIDTH 128
#define SCREEN_HEIGHT 64
#define OLED_RESET -1
#define OLED_I2C_ADDRESS 0x3C  // Common I2C address for SSD1306

// Define input pins for buttons
#define PIN_SELECT 10
#define PIN_LEFT 11
#define PIN_RIGHT 12
#define PIN_UP 13
#define PIN_DOWN 14

// Define LED strip pins
#define PIN_STRIP1 15
#define PIN_STRIP2 16
#define PIN_STRIP3 17

// Constants for LEDs
#define NUM_LEDS_PER_STRIP 3

// Objects
Adafruit_NeoPixel strip1 = Adafruit_NeoPixel(NUM_LEDS_PER_STRIP, PIN_STRIP1, NEO_GRB + NEO_KHZ800);
Adafruit_NeoPixel strip2 = Adafruit_NeoPixel(NUM_LEDS_PER_STRIP, PIN_STRIP2, NEO_GRB + NEO_KHZ800);
Adafruit_NeoPixel strip3 = Adafruit_NeoPixel(4, PIN_STRIP3, NEO_GRB + NEO_KHZ800);
Adafruit_SSD1306 display(SCREEN_WIDTH, SCREEN_HEIGHT, &Wire, OLED_RESET);

// Variables
int board[3][3]; // Game board
int currentX = 0, currentY = 0; // Current selected position
int currentPlayer = 1; // Current player (1 or 2)
int player1Wins = 0, player2Wins = 0; // Win counters

// Function prototypes
void setup();
void loop();
void displayWelcomeMessage();
void displayTurn();
void handleInput();
void selectSquare();
void updateLEDs();
void checkWin();
void displayWinner(int winner);
void resetGame();
void rainbowCycle(Adafruit_NeoPixel& strip, int pixel);
uint32_t Wheel(byte WheelPos);
int getLedIndex(int x, int y);

// Setup function
void setup() {
  Serial.begin(9600);
  
  // Initialize I2C with custom pins
  Wire.begin(19, 18);

  // Initialize the OLED display
  if (!display.begin(SSD1306_SWITCHCAPVCC, OLED_I2C_ADDRESS)) {
    Serial.println(F("SSD1306 allocation failed"));
    for (;;);
  }
  display.clearDisplay();
  display.display();

  // Initialize the WS2812 LEDs
  strip1.begin();
  strip2.begin();
  strip3.begin();
  strip1.show();
  strip2.show();
  strip3.show();

  // Initialize input pins
  pinMode(PIN_SELECT, INPUT_PULLUP);
  pinMode(PIN_LEFT, INPUT_PULLUP);
  pinMode(PIN_RIGHT, INPUT_PULLUP);
  pinMode(PIN_UP, INPUT_PULLUP);
  pinMode(PIN_DOWN, INPUT_PULLUP);

  // Initialize game
  resetGame();
  displayWelcomeMessage();
  delay(2000);
  displayTurn();
}

// Main loop function
void loop() {
  handleInput();
  updateLEDs();
  delay(100); // Add delay to debounce button presses
}

// Display welcome message
void displayWelcomeMessage() {
  display.clearDisplay();
  display.setTextSize(2);
  display.setTextColor(SSD1306_WHITE);
  display.setCursor(0, 0);
  display.println("Tic-Tac-Toe");
  display.display();
}

// Display current player's turn
void displayTurn() {
  display.clearDisplay();
  display.setTextSize(1);
  display.setCursor(0, 0);
  if (currentPlayer == 1) {
    display.println("Player 1 Turn");
    Serial.println(F("P1 Turn"));
  } else {
    display.println("Player 2 Turn");
    Serial.println(F("P2 Turn"));
  }
  display.display();
}

// Handle button inputs
void handleInput() {
  static unsigned long lastDebounceTime = 0;
  unsigned long debounceDelay = 200; // Adjust as needed

  if ((millis() - lastDebounceTime) > debounceDelay) {
    if (digitalRead(PIN_LEFT) == LOW) {
      Serial.println(F("Left Pressed"));
      currentX = (currentX - 1 + 3) % 3; // Wrap around between 0 and 2
      lastDebounceTime = millis();

      // Print current position after button press
      Serial.print(F("Current Position: ("));
      Serial.print(currentX);
      Serial.print(F(", "));
      Serial.print(currentY);
      Serial.println(F(")"));
    } 
    if (digitalRead(PIN_RIGHT) == LOW) {
      Serial.println(F("Right Pressed"));
      currentX = (currentX + 1) % 3; // Wrap around between 0 and 2
      lastDebounceTime = millis();

      // Print current position after button press
      Serial.print(F("Current Position: ("));
      Serial.print(currentX);
      Serial.print(F(", "));
      Serial.print(currentY);
      Serial.println(F(")"));
    } 
    if (digitalRead(PIN_UP) == LOW) {
      Serial.println(F("Up Pressed"));
      currentY = (currentY - 1 + 3) % 3; // Wrap around between 0 and 2
      lastDebounceTime = millis();

      // Print current position after button press
      Serial.print(F("Current Position: ("));
      Serial.print(currentX);
      Serial.print(F(", "));
      Serial.print(currentY);
      Serial.println(F(")"));
    } 
    if (digitalRead(PIN_DOWN) == LOW) {
      Serial.println(F("Down Pressed"));
      currentY = (currentY + 1) % 3; // Wrap around between 0 and 2
      lastDebounceTime = millis();

      // Print current position after button press
      Serial.print(F("Current Position: ("));
      Serial.print(currentX);
      Serial.print(F(", "));
      Serial.print(currentY);
      Serial.println(F(")"));
    } 
    if (digitalRead(PIN_SELECT) == LOW) {
      Serial.println(F("Select Pressed"));
      selectSquare();
      lastDebounceTime = millis();
    }
  }
}

// Select the current square
void selectSquare() {
  if (board[currentY][currentX] == 0) {
    board[currentY][currentX] = currentPlayer;
    int ledIndex = getLedIndex(currentX, currentY);
    if (currentPlayer == 1) {
      if (ledIndex < NUM_LEDS_PER_STRIP) {
        strip1.setPixelColor(ledIndex, strip1.Color(255, 0, 0)); // Red
      } else if (ledIndex < 2 * NUM_LEDS_PER_STRIP) {
        strip2.setPixelColor(ledIndex - NUM_LEDS_PER_STRIP, strip2.Color(255, 0, 0)); // Red
      } else {
        strip3.setPixelColor(ledIndex - 2 * NUM_LEDS_PER_STRIP, strip3.Color(255, 0, 0)); // Red
      }
    } else {
      if (ledIndex < NUM_LEDS_PER_STRIP) {
        strip1.setPixelColor(ledIndex, strip1.Color(0, 0, 255)); // Blue
      } else if (ledIndex < 2 * NUM_LEDS_PER_STRIP) {
        strip2.setPixelColor(ledIndex - NUM_LEDS_PER_STRIP, strip2.Color(0, 0, 255)); // Blue
      } else {
        strip3.setPixelColor(ledIndex - 2 * NUM_LEDS_PER_STRIP, strip3.Color(0, 0, 255)); // Blue
      }
    }
    strip1.show();
    strip2.show();
    strip3.show();
    checkWin();
    currentPlayer = 3 - currentPlayer; // Switch player
    displayTurn();
  }
}

// Update LEDs
void updateLEDs() {
  for (int i = 0; i < 3; i++) {
    for (int j = 0; j < 3; j++) {
      int ledIndex = getLedIndex(i, j);
      if (board[j][i] == 0 && i == currentX && j == currentY) {
        if (ledIndex < NUM_LEDS_PER_STRIP) {
          rainbowCycle(strip1, ledIndex);
        } else if (ledIndex < 2 * NUM_LEDS_PER_STRIP) {
          rainbowCycle(strip2, ledIndex - NUM_LEDS_PER_STRIP);
        } else {
          rainbowCycle(strip3, ledIndex - 2 * NUM_LEDS_PER_STRIP);
        }
      } else {
        if (ledIndex < NUM_LEDS_PER_STRIP) {
          if (board[j][i] == 0) {
            strip1.setPixelColor(ledIndex, strip1.Color(0, 0, 0));
          } else if (board[j][i] == 1) {
            strip1.setPixelColor(ledIndex, strip1.Color(255, 0, 0)); // Red for Player 1
          } else if (board[j][i] == 2) {
            strip1.setPixelColor(ledIndex, strip1.Color(0, 0, 255)); // Blue for Player 2
          }
        } else if (ledIndex < 2 * NUM_LEDS_PER_STRIP) {
          if (board[j][i] == 0) {
            strip2.setPixelColor(ledIndex - NUM_LEDS_PER_STRIP, strip2.Color(0, 0, 0));
          } else if (board[j][i] == 1) {
            strip2.setPixelColor(ledIndex - NUM_LEDS_PER_STRIP, strip2.Color(255, 0, 0)); // Red for Player 1
          } else if (board[j][i] == 2) {
            strip2.setPixelColor(ledIndex - NUM_LEDS_PER_STRIP, strip2.Color(0, 0, 255)); // Blue for Player 2
          }
        } else {
          if (board[j][i] == 0) {
            strip3.setPixelColor(ledIndex - 2 * NUM_LEDS_PER_STRIP, strip3.Color(0, 0, 0));
          } else if (board[j][i] == 1) {
            strip3.setPixelColor(ledIndex - 2 * NUM_LEDS_PER_STRIP, strip3.Color(255, 0, 0)); // Red for Player 1
          } else if (board[j][i] == 2) {
            strip3.setPixelColor(ledIndex - 2 * NUM_LEDS_PER_STRIP, strip3.Color(0, 0, 255)); // Blue for Player 2
          }
        }
      }
    }
  }
  strip1.show();
  strip2.show();
  strip3.show();
}

// Check for a win condition
void checkWin() {
  int winner = 0;
  // Check rows
  for (int i = 0; i < 3; i++) {
    if (board[i][0] != 0 && board[i][0] == board[i][1] && board[i][1] == board[i][2]) {
      winner = board[i][0];
    }
  }
  // Check columns
  for (int i = 0; i < 3; i++) {
    if (board[0][i] != 0 && board[0][i] == board[1][i] && board[1][i] == board[2][i]) {
      winner = board[0][i];
    }
  }
  // Check diagonals
  if (board[0][0] != 0 && board[0][0] == board[1][1] && board[1][1] == board[2][2]) {
    winner = board[0][0];
  }
  if (board[0][2] != 0 && board[0][2] == board[1][1] && board[1][1] == board[2][0]) {
    winner = board[0][2];
  }

  if (winner != 0) {
    displayWinner(winner);
    resetGame();
    displayTurn();
  }
}

// Display the winner
void displayWinner(int winner) {
  display.clearDisplay();
  display.setTextSize(2);
  display.setTextColor(SSD1306_WHITE);
  display.setCursor(0, 0);
  if (winner == 1) {
    display.println("Player 1 Wins!");
    player1Wins++;
  } else {
    display.println("Player 2 Wins!");
    player2Wins++;
  }
  display.display();
  delay(2000);
}

// Reset the game board
void resetGame() {
  for (int i = 0; i < 3; i++) {
    for (int j = 0; j < 3; j++) {
      board[i][j] = 0;
    }
  }
  currentX = 0;
  currentY = 0;
  currentPlayer = 1;
}

// Get LED index from board coordinates
int getLedIndex(int x, int y) {
  return y * 3 + x;
}

// Rainbow cycle animation for a single LED
void rainbowCycle(Adafruit_NeoPixel& strip, int pixel) {
  static uint16_t j = 0;
  j++;
  if (j >= 256 * 5) j = 0; // Reset the color wheel position
  uint32_t color = Wheel(((pixel * 256 / strip.numPixels()) + j) & 255);
  strip.setPixelColor(pixel, color);
  strip.show();
}

// Input a value 0 to 255 to get a color value.
// The colours are a transition r - g - b - back to r.
uint32_t Wheel(byte WheelPos) {
  WheelPos = 255 - WheelPos;
  if (WheelPos < 85) {
    return strip1.Color(255 - WheelPos * 3, 0, WheelPos * 3);
  }
  if (WheelPos < 170) {
    WheelPos -= 85;
    return strip1.Color(0, WheelPos * 3, 255 - WheelPos * 3);
  }
  WheelPos -= 170;
  return strip1.Color(WheelPos * 3, 255 - WheelPos * 3, 0);
}