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

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

#define BUZZER_PIN 2       // Interchanged with BUTTON_LEFT_PIN
#define BUTTON_LEFT_PIN 3  // Interchanged with BUZZER_PIN
#define BUTTON_RIGHT_PIN 4
#define BUTTON_ROTATE_PIN 5
#define BUTTON_DOWN_PIN 6
#define LED_R_PIN 9
#define LED_G_PIN 10
#define LED_B_PIN 11

Bounce buttonLeft = Bounce();
Bounce buttonRight = Bounce();
Bounce buttonRotate = Bounce();
Bounce buttonDown = Bounce();

const uint16_t TETROMINOES[7][4] = {
  {0x0F00, 0x2222, 0x0F00, 0x2222},  // I
  {0x0E20, 0x88E0, 0x2E00, 0xE880},  // J
  {0x0E80, 0x44E0, 0xE200, 0xE440},  // L
  {0x6600, 0x6600, 0x6600, 0x6600},  // O
  {0x0E40, 0x4C40, 0x4E00, 0x4640},  // T
  {0x06C0, 0x8C40, 0x6C00, 0x46C0},  // S
  {0x0C60, 0x4C80, 0xC600, 0xC680}   // Z
};

uint8_t board[10][20] = {0};
uint8_t currentTetromino, rotation;
int8_t x, y;
unsigned long lastUpdateTime = 0;
const unsigned long updateInterval = 500; // Game update interval in ms

void setup() {
  pinMode(BUZZER_PIN, OUTPUT);
  pinMode(BUTTON_LEFT_PIN, INPUT_PULLUP);
  pinMode(BUTTON_RIGHT_PIN, INPUT_PULLUP);
  pinMode(BUTTON_ROTATE_PIN, INPUT_PULLUP);
  pinMode(BUTTON_DOWN_PIN, INPUT_PULLUP);
  pinMode(LED_R_PIN, OUTPUT);
  pinMode(LED_G_PIN, OUTPUT);
  pinMode(LED_B_PIN, OUTPUT);

  buttonLeft.attach(BUTTON_LEFT_PIN);
  buttonLeft.interval(5);
  buttonRight.attach(BUTTON_RIGHT_PIN);
  buttonRight.interval(5);
  buttonRotate.attach(BUTTON_ROTATE_PIN);
  buttonRotate.interval(5);
  buttonDown.attach(BUTTON_DOWN_PIN);
  buttonDown.interval(5);

  if (!display.begin(SSD1306_SWITCHCAPVCC, 0x3C)) {
    for (;;); // Don't proceed, loop forever
  }
  display.clearDisplay();
  display.display();
  
  showStartScreen();
  delay(2000);  // Display start screen for 2 seconds

  showLoadingScreen();
  delay(2000);  // Display loading screen for 2 seconds

  resetGame();
}

void showStartScreen() {
  display.clearDisplay();
  display.setTextSize(2);
  display.setTextColor(WHITE);
  display.setCursor(10, 20);
  display.print("NaNTetris");
  display.display();
  playTetrisTheme();
}

void showLoadingScreen() {
  display.clearDisplay();
  display.setTextSize(1);
  display.setCursor(20, 30);
  display.print("Loading...");
  display.display();
}

void resetGame() {
  memset(board, 0, sizeof(board));
  spawnTetromino();
}

void spawnTetromino() {
  currentTetromino = random(7);
  rotation = 0;
  x = 4;
  y = 0;
  if (!isValidPosition()) {
    gameOver();
  } else {
    setLEDColor(0, 255, 0); // Green
  }
}

bool isValidPosition() {
  for (int8_t i = 0; i < 4; i++) {
    for (int8_t j = 0; j < 4; j++) {
      if (TETROMINOES[currentTetromino][rotation] & (0x8000 >> (i + j * 4))) {
        int8_t newX = x + i;
        int8_t newY = y + j;
        if (newX < 0 || newX >= 10 || newY >= 20 || (newY >= 0 && board[newX][newY])) {
          return false;
        }
      }
    }
  }
  return true;
}

void drawBoard() {
  display.clearDisplay();
  for (uint8_t i = 0; i < 10; i++) {
    for (uint8_t j = 0; j < 20; j++) {
      if (board[i][j]) {
        display.fillRect(i * 6, j * 3, 5, 2, WHITE);
      }
    }
  }
  // Draw current piece
  for (int8_t i = 0; i < 4; i++) {
    for (int8_t j = 0; j < 4; j++) {
      if (TETROMINOES[currentTetromino][rotation] & (0x8000 >> (i + j * 4))) {
        display.fillRect((x + i) * 6, (y + j) * 3, 5, 2, WHITE);
      }
    }
  }
  display.display();
}

void lockTetromino() {
  for (int8_t i = 0; i < 4; i++) {
    for (int8_t j = 0; j < 4; j++) {
      if (TETROMINOES[currentTetromino][rotation] & (0x8000 >> (i + j * 4))) {
        int8_t newX = x + i;
        int8_t newY = y + j;
        if (newY >= 0 && newY < 20) {
          board[newX][newY] = 1;
        }
      }
    }
  }
  tone(BUZZER_PIN, 1000, 100); // Short beep
}

void clearLines() {
  for (int8_t j = 19; j >= 0; j--) {
    bool fullLine = true;
    for (int8_t i = 0; i < 10; i++) {
      if (!board[i][j]) {
        fullLine = false;
        break;
      }
    }
    if (fullLine) {
      for (int8_t k = j; k > 0; k--) {
        for (int8_t i = 0; i < 10; i++) {
          board[i][k] = board[i][k - 1];
        }
      }
      j++;
      tone(BUZZER_PIN, 500, 200); // Long beep
      setLEDColor(255, 0, 0); // Red
    }
  }
}

void gameOver() {
  showGameOverScreen();
  delay(2000);  // Display game over screen for 2 seconds
  showStartScreen();
  delay(2000);  // Display start screen for 2 seconds before restarting
  resetGame();
}

void showGameOverScreen() {
  display.clearDisplay();
  display.setTextSize(1);
  display.setCursor(20, 30);
  display.print("Game Over");
  display.display();
  tone(BUZZER_PIN, 300, 500); // Play a longer tone
  setLEDColor(128, 0, 128); // Purple
}

void setLEDColor(uint8_t red, uint8_t green, uint8_t blue) {
  analogWrite(LED_R_PIN, red);
  analogWrite(LED_G_PIN, green);
  analogWrite(LED_B_PIN, blue);
}

void playTetrisTheme() {
  // Play a simple melody emulating the Tetris theme
  tone(BUZZER_PIN, 400, 300);
  delay(300);
  tone(BUZZER_PIN, 600, 300);
  delay(300);
  tone(BUZZER_PIN, 800, 300);
  delay(300);
  tone(BUZZER_PIN, 600, 300);
  delay(300);
  tone(BUZZER_PIN, 400, 300);
  delay(300);
  noTone(BUZZER_PIN);
}

void loop() {
  buttonLeft.update();
  buttonRight.update();
  buttonRotate.update();
  buttonDown.update();

  // Handle inputs
  if (buttonLeft.fell()) {
    x--;
    if (!isValidPosition()) x++;
  }
  if (buttonRight.fell()) {
    x++;
    if (!isValidPosition()) x--;
  }
  if (buttonRotate.fell()) {
    rotation = (rotation + 1) % 4;
    if (!isValidPosition()) rotation = (rotation - 1) % 4;
  }
  if (buttonDown.fell()) {
    y++;
    if (!isValidPosition()) y--;
  }

  // Update game state
  if (millis() - lastUpdateTime > updateInterval) {
    y++;
    if (!isValidPosition()) {
      y--;
      lockTetromino();
      clearLines();
      spawnTetromino();
    }
    lastUpdateTime = millis();
  }

  drawBoard();
}