#include <SPI.h>
#include <Adafruit_GFX.h>
#include <Adafruit_ILI9341.h>

// Display pins
#define TFT_CS 3
#define TFT_DC 2

// Button pins
#define LEFT_BUTTON 4
#define RIGHT_BUTTON 5
#define ROTATE_BUTTON 6
#define PAUSE_BUTTON 7  // Added a dedicated pause button

Adafruit_ILI9341 tft = Adafruit_ILI9341(TFT_CS, TFT_DC);

// Game constants
#define GRID_WIDTH 12
#define GRID_HEIGHT 24
#define CELL_SIZE 12
#define GRID_OFFSET_X 2
#define GRID_OFFSET_Y 2

// Tetromino shapes (unchanged)
const uint8_t TETROMINOS[7][4][4] = {
  // I
  {
    {0, 0, 0, 0},
    {1, 1, 1, 1},
    {0, 0, 0, 0},
    {0, 0, 0, 0}
  },
  // J
  {
    {1, 0, 0, 0},
    {1, 1, 1, 0},
    {0, 0, 0, 0},
    {0, 0, 0, 0}
  },
  // L
  {
    {0, 0, 1, 0},
    {1, 1, 1, 0},
    {0, 0, 0, 0},
    {0, 0, 0, 0}
  },
  // O
  {
    {0, 1, 1, 0},
    {0, 1, 1, 0},
    {0, 0, 0, 0},
    {0, 0, 0, 0}
  },
  // S
  {
    {0, 1, 1, 0},
    {1, 1, 0, 0},
    {0, 0, 0, 0},
    {0, 0, 0, 0}
  },
  // T
  {
    {0, 1, 0, 0},
    {1, 1, 1, 0},
    {0, 0, 0, 0},
    {0, 0, 0, 0}
  },
  // Z
  {
    {1, 1, 0, 0},
    {0, 1, 1, 0},
    {0, 0, 0, 0},
    {0, 0, 0, 0}
  }
};

// Tetromino colors (unchanged)
const uint16_t TETROMINO_COLORS[7] = {
  ILI9341_CYAN,     // I
  ILI9341_BLUE,     // J
  ILI9341_ORANGE,   // L
  ILI9341_YELLOW,   // O
  ILI9341_GREEN,    // S
  ILI9341_PURPLE,   // T
  ILI9341_RED       // Z
};

// Game variables
uint8_t grid[GRID_HEIGHT][GRID_WIDTH] = {0};
int8_t currentPiece[4][4];
int8_t nextPiece[4][4];
uint8_t currentType, nextType;
int16_t currentX, currentY;
int16_t prevX, prevY;  // Track previous position for erasing
int8_t prevPiece[4][4]; // Track previous rotation for erasing
uint16_t score = 0;
uint16_t highScore = 0;
uint8_t level = 1;
uint16_t linesCleared = 0;
bool gameOver = false;
bool gamePaused = false;
unsigned long lastDropTime = 0;
unsigned long lastButtonPress = 0;

void setup() {
  Serial.begin(9600);
  
  // Initialize display
  tft.begin();
  tft.setRotation(0);
  tft.fillScreen(ILI9341_BLACK);
  
  // Initialize buttons
  pinMode(LEFT_BUTTON, INPUT_PULLUP);
  pinMode(RIGHT_BUTTON, INPUT_PULLUP);
  pinMode(ROTATE_BUTTON, INPUT_PULLUP);
  pinMode(PAUSE_BUTTON, INPUT_PULLUP); // Added pause button
  
  // Initialize random seed
  randomSeed(analogRead(0));
  
  // Start new game
  newGame();
}

void loop() {
  if (gameOver) {
    if (digitalRead(ROTATE_BUTTON) == LOW) {
      delay(200); // Debounce
      newGame();
    }
    return;
  }
  
  // Check for pause with dedicated button
  if (digitalRead(PAUSE_BUTTON) == LOW && millis() - lastButtonPress > 500) {
    gamePaused = !gamePaused;
    lastButtonPress = millis();
    if (gamePaused) {
      drawPauseScreen();
    } else {
      drawResumeScreen();
      tft.fillRect(tft.width()/4, tft.height()/3, tft.width()/2, 40, ILI9341_BLACK);
      lastDropTime = millis();
    }
  }
  
  if (gamePaused) {
    return;
  }
  
  // Store previous position and rotation
  prevX = currentX;
  prevY = currentY;
  memcpy(prevPiece, currentPiece, sizeof(prevPiece));
  
  // Handle input
  handleInput();
  
  // Game logic
  updateGame();
  
  // Draw game - only update changed parts
  drawGame();
  
  // Small delay to control game speed
  delay(20);
}

void newGame() {
  // Reset game state
  memset(grid, 0, sizeof(grid));
  score = 0;
  level = 1;
  linesCleared = 0;
  gameOver = false;
  gamePaused = false;
  
  // Generate first pieces
  nextType = random(7);
  generateNextPiece();
  spawnNewPiece();
  
  // Clear screen
  tft.fillScreen(ILI9341_BLACK);
  
  // Draw initial game elements
  drawGrid();
  drawNextPiece();
  drawUI();
}

void generateNextPiece() {
  currentType = nextType;
  nextType = random(7);
  
  // Copy the current piece
  memcpy(currentPiece, TETROMINOS[currentType], sizeof(currentPiece));
  
  // Copy the next piece
  memcpy(nextPiece, TETROMINOS[nextType], sizeof(nextPiece));
  
  // Update the next piece display
  drawNextPiece();
}

void spawnNewPiece() {
  currentX = GRID_WIDTH / 2 - 2;
  currentY = 0;
  prevX = currentX;
  prevY = currentY;
  memcpy(prevPiece, currentPiece, sizeof(prevPiece));
  
  // Check if game over
  if (!isValidPosition()) {
    gameOver = true;
    if (score > highScore) highScore = score;
    drawGameOver();
  }
}

void handleInput() {
  // Left movement
  if (digitalRead(LEFT_BUTTON) == LOW) {
    movePiece(-1, 0);
    delay(150);
  }
  
  // Right movement
  if (digitalRead(RIGHT_BUTTON) == LOW) {
    movePiece(1, 0);
    delay(150);
  }
  
  // Rotation - fixed to only rotate, not pause
  if (digitalRead(ROTATE_BUTTON) == LOW && millis() - lastButtonPress > 200) {
    rotatePiece();
    lastButtonPress = millis();
  }
}

void updateGame() {
  // Automatic dropping
  unsigned long currentTime = millis();
  int dropDelay = max(1000 - (level - 1) * 100, 100);
  
  if (currentTime - lastDropTime > dropDelay) {
    movePiece(0, 1);
    lastDropTime = currentTime;
  }
}

void movePiece(int dx, int dy) {
  // Try to move the piece
  currentX += dx;
  currentY += dy;
  
  // If the new position is invalid, revert the move
  if (!isValidPosition()) {
    currentX -= dx;
    currentY -= dy;
    
    // If moving down was invalid, place the piece
    if (dy > 0) {
      placePiece();
      clearLines();
      generateNextPiece();
      spawnNewPiece();
    }
  }
}

void rotatePiece() {
  // Save current piece
  int8_t oldPiece[4][4];
  memcpy(oldPiece, currentPiece, sizeof(oldPiece));
  
  // Rotate the piece 90 degrees clockwise
  for (int i = 0; i < 4; i++) {
    for (int j = 0; j < 4; j++) {
      currentPiece[i][j] = oldPiece[3 - j][i];
    }
  }
  
  // If the rotation is invalid, revert it
  if (!isValidPosition()) {
    memcpy(currentPiece, oldPiece, sizeof(currentPiece));
  }
}

bool isValidPosition() {
  for (int y = 0; y < 4; y++) {
    for (int x = 0; x < 4; x++) {
      if (currentPiece[y][x]) {
        int gridX = currentX + x;
        int gridY = currentY + y;
        
        // Check if out of bounds
        if (gridX < 0 || gridX >= GRID_WIDTH || gridY >= GRID_HEIGHT) {
          return false;
        }
        
        // Check if collides with placed pieces
        if (gridY >= 0 && grid[gridY][gridX]) {
          return false;
        }
      }
    }
  }
  return true;
}

void placePiece() {
  for (int y = 0; y < 4; y++) {
    for (int x = 0; x < 4; x++) {
      if (currentPiece[y][x]) {
        int gridX = currentX + x;
        int gridY = currentY + y;
        
        if (gridY >= 0 && gridY < GRID_HEIGHT && gridX >= 0 && gridX < GRID_WIDTH) {
          grid[gridY][gridX] = currentType + 1;
        }
      }
    }
  }
}

void clearLines() {
  uint8_t linesToClear = 0;
  
  for (int y = GRID_HEIGHT - 1; y >= 0; y--) {
    bool lineComplete = true;
    
    for (int x = 0; x < GRID_WIDTH; x++) {
      if (!grid[y][x]) {
        lineComplete = false;
        break;
      }
    }
    
    if (lineComplete) {
      // Move all lines above down
      for (int y2 = y; y2 > 0; y2--) {
        for (int x = 0; x < GRID_WIDTH; x++) {
          grid[y2][x] = grid[y2 - 1][x];
        }
      }
      
      // Clear the top line
      for (int x = 0; x < GRID_WIDTH; x++) {
        grid[0][x] = 0;
      }
      
      linesToClear++;
      y++; // Check the same line again
    }
  }
  
  if (linesToClear > 0) {
    // Update score
    uint16_t points = 0;
    switch (linesToClear) {
      case 1: points = 100; break;
      case 2: points = 300; break;
      case 3: points = 500; break;
      case 4: points = 800; break;
    }
    score += points * level;
    
    // Update lines cleared and level
    linesCleared += linesToClear;
    level = linesCleared / 10 + 1;
    
    // Redraw grid and UI
    drawGrid();
    drawUI();
  }
}

void drawGame() {
  // Erase the previous piece position
  erasePiece(prevX, prevY, prevPiece);
  
  // Draw the current piece
  drawPiece(currentX, currentY, currentPiece, currentType);
}

void erasePiece(int16_t x, int16_t y, int8_t piece[4][4]) {
  // Erase the piece at the specified position
  for (int py = 0; py < 4; py++) {
    for (int px = 0; px < 4; px++) {
      if (piece[py][px]) {
        int gridX = x + px;
        int gridY = y + py;
        
        if (gridY >= 0) { // Only erase if the cell is visible
          if (gridY < GRID_HEIGHT && gridX >= 0 && gridX < GRID_WIDTH && grid[gridY][gridX] == 0) {
            // Only erase if the cell is empty in the grid
            tft.fillRect(GRID_OFFSET_X + gridX * CELL_SIZE, 
                         GRID_OFFSET_Y + gridY * CELL_SIZE, 
                         CELL_SIZE, CELL_SIZE, ILI9341_BLACK);
          }
        }
      }
    }
  }
}

void drawGrid() {
   // Draw the grid background
   tft.fillRect(GRID_OFFSET_X, GRID_OFFSET_Y, 
                GRID_WIDTH * CELL_SIZE, GRID_HEIGHT * CELL_SIZE, 
                ILI9341_BLACK);
   
   // Draw placed pieces
   for (int y = 0; y < GRID_HEIGHT; y++) {
     for (int x = 0; x < GRID_WIDTH; x++) {
       if (grid[y][x]) {
         uint8_t colorIndex = grid[y][x] - 1;
         drawCell(x, y, TETROMINO_COLORS[colorIndex]);
       }
     }
   }
   
   // Draw grid border
   tft.drawRect(GRID_OFFSET_X - 1, GRID_OFFSET_Y - 1, 
                GRID_WIDTH * CELL_SIZE + 2, GRID_HEIGHT * CELL_SIZE + 2, 
                ILI9341_RED);
 }

void drawPiece(int16_t x, int16_t y, int8_t piece[4][4], uint8_t type) {
  // Draw the piece at the specified position
  for (int py = 0; py < 4; py++) {
    for (int px = 0; px < 4; px++) {
      if (piece[py][px]) {
        int gridX = x + px;
        int gridY = y + py;
        
        if (gridY >= 0) { // Only draw if the cell is visible
          drawCell(gridX, gridY, TETROMINO_COLORS[type]);
        }
      }
    }
  }
}

void drawCell(int16_t x, int16_t y, uint16_t color) {
  tft.fillRect(GRID_OFFSET_X + x * CELL_SIZE, 
               GRID_OFFSET_Y + y * CELL_SIZE, 
               CELL_SIZE, CELL_SIZE, color);
  
  // Draw cell border
  tft.drawRect(GRID_OFFSET_X + x * CELL_SIZE, 
               GRID_OFFSET_Y + y * CELL_SIZE, 
               CELL_SIZE, CELL_SIZE, ILI9341_BLACK);
}

void drawNextPiece() {
  // Clear next piece area
  tft.fillRect(160, 220, 80, 60, ILI9341_BLACK);
  
  // Draw next piece title
  tft.setTextColor(ILI9341_WHITE);
  tft.setTextSize(2);
  tft.setCursor(160, 200);
  tft.print("NEXT:");
  
  // Draw the next piece - centered in the preview area
  int centerX = 175;
  int centerY = 230;
  
  for (int y = 0; y < 4; y++) {
    for (int x = 0; x < 4; x++) {
      if (nextPiece[y][x]) {
        tft.fillRect(centerX + x * (CELL_SIZE-2), 
                     centerY + y * (CELL_SIZE-2), 
                     CELL_SIZE-2, CELL_SIZE-2, 
                     TETROMINO_COLORS[nextType]);
        
        // Draw cell border
        tft.drawRect(centerX + x * (CELL_SIZE-2), 
                     centerY + y * (CELL_SIZE-2), 
                     CELL_SIZE-2, CELL_SIZE-2, ILI9341_BLACK);
      }
    }
  }
}

void drawUI() {
  // Clear UI area
  tft.fillRect(160, 0, 80, 160, ILI9341_BLACK);
  
  // Draw game title
  tft.setTextColor(ILI9341_YELLOW);
  tft.setTextSize(2);
  tft.setCursor(160, 2);
  tft.print("TETRIS");
  
  // Draw score
  tft.setTextColor(ILI9341_GREEN);
  tft.setTextSize(2);
  tft.setCursor(160, 22);
  tft.print("SCORE:");
  tft.setCursor(160, 42);
  tft.setTextColor(ILI9341_WHITE);
  tft.print(score);
  
  // Draw level
  tft.setCursor(160, 62);
   tft.setTextColor(ILI9341_GREEN);
  tft.print("LEVEL:");
  tft.setCursor(160, 82);
  tft.setTextColor(ILI9341_WHITE);
  tft.print(level);
  
  // Draw lines
  tft.setCursor(160, 102);
  tft.setTextColor(ILI9341_GREEN);
  tft.print("LINES:");
  tft.setCursor(160, 122);
  tft.setTextColor(ILI9341_WHITE);
  tft.print(linesCleared);
  
  // Draw high score
  tft.setCursor(160, 142);
  tft.setTextColor(ILI9341_GREEN);
  tft.print("HIGH:");
  tft.setCursor(160, 162);
  tft.setTextColor(ILI9341_WHITE);
  tft.print(highScore);
  
  // Draw next piece
  drawNextPiece();
}
void drawResumeScreen(){
   tft.setTextColor(ILI9341_BLACK);
  tft.setTextSize(2);
  tft.setCursor(30, 90);
  tft.print("PAUSED");
}
void drawPauseScreen() {
  tft.fillRect(GRID_OFFSET_X, GRID_OFFSET_Y + 80, 
               GRID_WIDTH * CELL_SIZE, 40, ILI9341_BLACK);
  tft.setTextColor(ILI9341_WHITE);
  tft.setTextSize(2);
  tft.setCursor(30, 90);
  tft.print("PAUSED");
}

void drawGameOver() {
  tft.fillRect(GRID_OFFSET_X, GRID_OFFSET_Y + 80, 
               GRID_WIDTH * CELL_SIZE, 60, ILI9341_BLACK);
  tft.setTextColor(ILI9341_RED);
  tft.setTextSize(2);
  tft.setCursor(GRID_OFFSET_X + 10, GRID_OFFSET_Y + 90);
  tft.print("GAME OVER");
  
  tft.setTextColor(ILI9341_WHITE);
  tft.setTextSize(1);
  tft.setCursor(GRID_OFFSET_X + 15, GRID_OFFSET_Y + 115);
  tft.print("Score: ");
  tft.print(score);
  
  tft.setCursor(GRID_OFFSET_X + 10, GRID_OFFSET_Y + 130);
  tft.print("Press ROTATE");
  tft.setCursor(GRID_OFFSET_X + 20, GRID_OFFSET_Y + 145);
  tft.print("to restart");
}