#include <MD_MAX72xx.h>
#include <SPI.h>

// Define MAX7219 specifications
#define MAX_DEVICES 4
#define CLK_PIN 11
#define DATA_PIN 12
#define CS_PIN 10

MD_MAX72XX matrix = MD_MAX72XX(MD_MAX72XX::FC16_HW, DATA_PIN, CLK_PIN, CS_PIN, MAX_DEVICES);

#define BUTTON_PIN 2

#define WIDTH 8
#define HEIGHT 32
#define HIDDEN_STEPS 3  // How many ticks fully hidden before reappearing

int blockPos = 0;
int blockWidth = 4;
int blockRow = 0;
bool blockMoving = true;
int stack[HEIGHT][WIDTH];
int stackHeight = 0;
bool gameOver = false;
unsigned long lastMove = 0;
int moveDelay = 200;
bool moveRight = true;
bool testMode = false;
int testRow = 0;
const bool ROTATE_DISPLAY = true;

bool blockHidden = false;
int hiddenCount = 0;
bool needsRedraw = false;

// Shadow buffer for flicker-free updates
bool displayBuffer[HEIGHT][WIDTH];

void updateDisplay() {
  bool desired[HEIGHT][WIDTH];
  for (int r = 0; r < HEIGHT; r++)
    for (int c = 0; c < WIDTH; c++)
      desired[r][c] = false;

  if (testMode) {
    for (int c = 0; c < WIDTH; c++)
      desired[testRow][c] = true;
  } else {
    // Stack
    for (int r = 0; r < HEIGHT; r++)
      for (int c = 0; c < WIDTH; c++)
        if (stack[r][c]) desired[r][c] = true;

    // Moving block — draw only the visible portion (clipped to 0..WIDTH-1)
    if (!gameOver && blockMoving) {
      for (int i = 0; i < blockWidth; i++) {
        int c = blockPos + i;
        if (c >= 0 && c < WIDTH)
          desired[blockRow][c] = true;
      }
    }
  }

  // Diff and flush in one SPI batch
  matrix.control(MD_MAX72XX::UPDATE, MD_MAX72XX::OFF);
  for (int r = 0; r < HEIGHT; r++) {
    for (int c = 0; c < WIDTH; c++) {
      if (desired[r][c] != displayBuffer[r][c]) {
        displayBuffer[r][c] = desired[r][c];
        if (ROTATE_DISPLAY) matrix.setPoint(c, r, desired[r][c]);
        else matrix.setPoint(r, c, desired[r][c]);
      }
    }
  }
  matrix.control(MD_MAX72XX::UPDATE, MD_MAX72XX::ON);
}

void clearDisplayBuffer() {
  matrix.clear();
  for (int r = 0; r < HEIGHT; r++)
    for (int c = 0; c < WIDTH; c++)
      displayBuffer[r][c] = false;
}

void setup() {
  Serial.begin(9600);
  matrix.begin();
  clearDisplayBuffer();
  pinMode(BUTTON_PIN, INPUT_PULLUP);
  randomSeed(analogRead(A0));
  resetGame();
  Serial.println("Game initialized.");
}

void resetGame() {
  for (int r = 0; r < HEIGHT; r++)
    for (int c = 0; c < WIDTH; c++)
      stack[r][c] = 0;

  blockPos = 0;
  blockWidth = 4;
  blockRow = 0;
  stackHeight = 0;
  gameOver = false;
  blockMoving = true;
  moveDelay = 200;
  moveRight = true;
  blockHidden = false;
  hiddenCount = 0;
  testMode = false;
  testRow = 0;
  needsRedraw = true;
  clearDisplayBuffer();
  Serial.println("Game reset");
}

void dropBlock() {
  // Count how many LEDs of the block are actually visible right now
  int visibleCount = 0;
  for (int i = 0; i < blockWidth; i++) {
    int c = blockPos + i;
    if (c >= 0 && c < WIDTH) visibleCount++;
  }

  // If fully off screen, treat as miss
  if (visibleCount == 0) {
    gameOver = true;
    Serial.println("Game over: Dropped while fully hidden!");
    needsRedraw = true;
    return;
  }

  blockMoving = false;

  // First row — place whatever is visible, no trimming needed
  if (blockRow == 0) {
    for (int i = 0; i < blockWidth; i++) {
      int col = blockPos + i;
      if (col >= 0 && col < WIDTH)
        stack[blockRow][col] = 1;
    }
    stackHeight = 1;
    blockRow = 1;
    blockWidth = min(blockWidth, visibleCount); // Trim to what was visible
    blockPos = random(0, max(1, WIDTH - blockWidth + 1));
    // Start block just off the edge so it slides in naturally
    if (moveRight) blockPos = -blockWidth;
    else blockPos = WIDTH;
    blockMoving = true;
    blockHidden = false;
    hiddenCount = 0;
    moveRight = true;
    moveDelay = max(50, moveDelay - 10);
    needsRedraw = true;
    return;
  }

  // Find overlap between visible block columns and the row below
  bool blockCols[WIDTH] = {false};
  for (int i = 0; i < blockWidth; i++) {
    int col = blockPos + i;
    if (col >= 0 && col < WIDTH)
      blockCols[col] = true;
  }

  bool overlapCols[WIDTH] = {false};
  int overlapCount = 0;
  for (int c = 0; c < WIDTH; c++) {
    if (blockCols[c] && stack[blockRow - 1][c] == 1) {
      overlapCols[c] = true;
      overlapCount++;
    }
  }

  if (overlapCount == 0) {
    gameOver = true;
    Serial.println("Game over: No overlap!");
    needsRedraw = true;
    return;
  }

  // Place only overlapping columns
  for (int c = 0; c < WIDTH; c++)
    if (overlapCols[c]) stack[blockRow][c] = 1;

  if (blockRow >= stackHeight) stackHeight = blockRow + 1;

  if (stackHeight >= HEIGHT) {
    gameOver = true;
    Serial.println("You win!");
    needsRedraw = true;
    return;
  }

  // New block width = trimmed overlap width
  int newLeft = WIDTH, newRight = -1;
  for (int c = 0; c < WIDTH; c++) {
    if (overlapCols[c]) {
      if (c < newLeft) newLeft = c;
      if (c > newRight) newRight = c;
    }
  }
  blockWidth = newRight - newLeft + 1;
  blockRow++;

  // Start the new block just off the LEFT edge — slides in 1 LED at a time
  blockPos = -blockWidth;  // Fully off-screen left
  moveRight = true;         // Moving right, so it slides in from the left
  blockMoving = true;
  blockHidden = false;
  hiddenCount = 0;
  moveDelay = max(50, moveDelay - 10);
  needsRedraw = true;

  Serial.print("New block width: ");
  Serial.println(blockWidth);
}

void loop() {
  needsRedraw = false;

  if (Serial.available()) {
    char c = Serial.read();
    if (c == 't') {
      testMode = !testMode;
      if (testMode) { testRow = 0; Serial.println("Test mode ON."); }
      else { Serial.println("Test mode OFF."); resetGame(); }
      needsRedraw = true;
    } else if (testMode && c == 'n') {
      testRow = (testRow + 1) % HEIGHT;
      needsRedraw = true;
    }
  }

  if (testMode) {
    if (digitalRead(BUTTON_PIN) == LOW) {
      delay(50);
      if (digitalRead(BUTTON_PIN) == LOW) {
        testRow = (testRow + 1) % HEIGHT;
        needsRedraw = true;
        while (digitalRead(BUTTON_PIN) == LOW);
      }
    }
    if (needsRedraw) updateDisplay();
    return;
  }

  if (gameOver) {
    static unsigned long lastFlash = 0;
    static bool flashOn = false;
    if (millis() - lastFlash >= 500) {
      flashOn = !flashOn;
      lastFlash = millis();
      if (flashOn) updateDisplay();
      else clearDisplayBuffer();
    }
    if (digitalRead(BUTTON_PIN) == LOW) {
      delay(50);
      if (digitalRead(BUTTON_PIN) == LOW) {
        resetGame();
        while (digitalRead(BUTTON_PIN) == LOW);
      }
    }
    return;
  }

  // Move block on timer
  if (blockMoving && millis() - lastMove >= moveDelay) {
    lastMove = millis();

    if (blockHidden) {
      // Fully off screen — count hidden ticks then reappear from same edge
      hiddenCount++;
      if (hiddenCount >= HIDDEN_STEPS) {
        blockHidden = false;
        hiddenCount = 0;
        // Place just off the edge it disappeared from, moving back
        if (!moveRight) {
          // Disappeared off left, reappear sliding in from left
          blockPos = -blockWidth + 1;
          moveRight = true;
        } else {
          // Disappeared off right, reappear sliding in from right
          blockPos = WIDTH - 1;
          moveRight = false;
        }
      }
    } else {
      // Normal movement — just increment, let display clipping handle the rest
      blockPos += moveRight ? 1 : -1;

      // Check if fully off screen
      if (blockPos >= WIDTH) {
        // Fully off right edge
        blockHidden = true;
        hiddenCount = 0;
        moveRight = false; // Will come back from right
      } else if (blockPos + blockWidth - 1 < 0) {
        // Fully off left edge
        blockHidden = true;
        hiddenCount = 0;
        moveRight = true; // Will come back from left
      }
    }

    needsRedraw = true;
  }

  // Button press
  if (digitalRead(BUTTON_PIN) == LOW) {
    delay(50);
    if (digitalRead(BUTTON_PIN) == LOW) {
      dropBlock();
      while (digitalRead(BUTTON_PIN) == LOW);
    }
  }

  if (needsRedraw) updateDisplay();
}