#include <U8g2lib.h>

U8G2_SSD1306_128X64_NONAME_F_HW_I2C u8g2(U8G2_R0, /* reset=*/ U8X8_PIN_NONE);

#define UP_BUTTON_PIN 2
#define DOWN_BUTTON_PIN 4
#define SELECT_BUTTON_PIN 7
#define WIDTH 64 // OLED display width, in pixels
#define HEIGHT 128 // OLED display height, in pixels
const uint8_t MARGIN_TOP = 19;
const uint8_t MARGIN_LEFT = 3;
const uint8_t SIZE = 5;
const uint8_t TYPES = 6;

const uint8_t pieces_S_l[2][2][4] PROGMEM = {{
    {0, 0, 1, 1}, {0, 1, 1, 2}
  },
  {
    {0, 1, 1, 2}, {1, 1, 0, 0}
  }};
const uint8_t pieces_S_r[2][2][4] PROGMEM = {{
    {1, 1, 0, 0}, {0, 1, 1, 2}
  },
  {
    {0, 1, 1, 2}, {0, 0, 1, 1}
  }};
const uint8_t pieces_L_l[4][2][4] PROGMEM = {{
    {0, 0, 0, 1}, {0, 1, 2, 2}
  },
  {
    {0, 1, 2, 2}, {1, 1, 1, 0}
  },
  {
    {0, 1, 1, 1}, {0, 0, 1, 2}
  },
  {
    {0, 0, 1, 2}, {1, 0, 0, 0}
  }};
const uint8_t pieces_Sq[1][2][4] PROGMEM = {{
    {0, 1, 0, 1}, {0, 0, 1, 1}
  }};
const uint8_t pieces_T[4][2][4] PROGMEM = {{
    {0, 0, 1, 0}, {0, 1, 1, 2}
  },
  {
    {0, 1, 1, 2}, {1, 0, 1, 1}
  },
  {
    {1, 0, 1, 1}, {0, 1, 1, 2}
  },
  {
    {0, 1, 1, 2}, {0, 0, 1, 0}
  }};
const uint8_t pieces_l[2][2][4] PROGMEM = {{
    {0, 1, 2, 3}, {0, 0, 0, 0}
  },
  {
    {0, 0, 0, 0}, {0, 1, 2, 3}
  }};

// Use uint16_t for larger numbers
uint16_t interval = 500;
uint16_t score;

uint8_t currentType, nextType, rotation;
uint8_t pieceX, pieceY;
uint8_t piece[2][4];
long timer, delayer;

boolean grid[10][18];
boolean b1, b2, b3;

void tetrisSetup() {
  pinMode(2, INPUT_PULLUP);
  pinMode(3, INPUT_PULLUP);
  pinMode(4, INPUT_PULLUP);
  pinMode(5, INPUT_PULLUP);

  Serial.begin(9600);

  u8g2.begin();

  randomSeed(analogRead(0));
  nextType = random(TYPES);
  generate();

  timer = millis();
}

void tetrisLoop() {
  if (millis() - timer > interval) {
    checkLines();
    refresh();

    if (nextCollision()) {
      for (short i = 0; i < 4; i++)
        grid[pieceX + piece[0][i]][pieceY + piece[1][i]] = 1;
      generate();
    } else
      pieceY++;

    timer = millis();
  }

  if (!digitalRead(2)) {
    if (b1) {
      if (!nextHorizontalCollision(piece, -1)) {
        pieceX--;
        refresh();
      }
      b1 = false;
    }
  } else { 
    b1 = true;
  }

  if (!digitalRead(4)) {
    if (b2) {
      if (!nextHorizontalCollision(piece, 1)) {
        pieceX++;
        refresh();
      }
      b2 = false;
    }
  } else { 
    b2 = true;
  }

  if (!digitalRead(5)) {
    interval = 20;
  } else {
    interval = 400;
  }

  if (!digitalRead(8)) {
    if (b3) {
      if (rotation == getMaxRotation(currentType) - 1 && canRotate(0)) {
        rotation = 0;
      } else if (canRotate(rotation + 1)) {
        rotation++;
      }

      copyPiece(piece, currentType, rotation);
      refresh();

      b3 = false;
      delayer = millis();
    }
  } else if (millis() - delayer > 50) { 
    b3 = true;
  }
}

void checkLines() {
  boolean full;
  for (short y = 17; y >= 0; y--) {
    full = true;
    for (short x = 0; x < 10; x++) {
      full = full && grid[x][y];
    }
    if (full) {
      breakLine(y);
      y++;
    }
  }
}

void breakLine(short line) {
  for (short y = line; y >= 0; y--) {
    for (short x = 0; x < 10; x++) {
      grid[x][y] = grid[x][y-1];
    }
  }
  for (short x = 0; x < 10; x++) {
    grid[x][0] = 0;
  }
}

void refresh() {
  u8g2.firstPage();
  do {
    drawLayout();
    drawGrid();
    drawPiece(currentType, 0, pieceX, pieceY);
  } while (u8g2.nextPage());
}

void drawGrid() {
  for (short x = 0; x < 10; x++)
    for (short y = 0; y < 18; y++)
      if (grid[x][y])
        u8g2.drawBox(MARGIN_LEFT + (SIZE + 1) * x, MARGIN_TOP + (SIZE + 1) * y, SIZE, SIZE);
}

boolean nextHorizontalCollision(uint8_t piece[2][4], int amount) {
  for (uint8_t i = 0; i < 4; i++) {
    int newX = pieceX + piece[0][i] + amount;
    if (newX > 9 || newX < 0 || grid[newX][pieceY + piece[1][i]])
      return true;
  }
  return false;
}


boolean nextCollision() {
  for (short i = 0; i < 4; i++) {
    short y = pieceY + piece[1][i] + 1;
    short x = pieceX + piece[0][i];
    if (y > 17 || grid[x][y])
      return true;
  }
  return false;
}

void generate() {
  currentType = nextType;
  nextType = random(TYPES);
  if (currentType != 5)
    pieceX = random(9);
  else
    pieceX = random(7);
  pieceY = 0;
  rotation = 0;
  copyPiece(piece, currentType, rotation);
}

void drawPiece(short type, short rotation, short x, short y) {
  for (short i = 0; i < 4; i++)
    u8g2.drawBox(MARGIN_LEFT + (SIZE + 1) * (x + piece[0][i]), MARGIN_TOP + (SIZE + 1) * (y + piece[1][i]), SIZE, SIZE);
}

void drawNextPiece() {
  short nPiece[2][4];
  copyPiece(nPiece, nextType, 0);
  for (short i = 0; i < 4; i++)
    u8g2.drawBox(50 + 3 * nPiece[0][i], 4 + 3 * nPiece[1][i], 2, 2);
}

void copyPiece(uint8_t piece[2][4], uint8_t type, uint8_t rotation) {
  switch(type) {
    case 0: //L_l
      for (uint8_t i = 0; i < 4; i++) {
        piece[0][i] = pieces_L_l[rotation][0][i];
        piece[1][i] = pieces_L_l[rotation][1][i];
      }
      break;
    case 1: //S_l
      for (uint8_t i = 0; i < 4; i++) {
        piece[0][i] = pieces_S_l[rotation][0][i];
        piece[1][i] = pieces_S_l[rotation][1][i];
      }
      break;
    case 2: //S_r
      for (uint8_t i = 0; i < 4; i++) {
        piece[0][i] = pieces_S_r[rotation][0][i];
        piece[1][i] = pieces_S_r[rotation][1][i];
      }
      break;
    case 3: //Sq
      for (uint8_t i = 0; i < 4; i++) {
        piece[0][i] = pieces_Sq[0][0][i];
        piece[1][i] = pieces_Sq[0][1][i];
      }
      break;
    case 4: //T
      for (uint8_t i = 0; i < 4; i++) {
        piece[0][i] = pieces_T[rotation][0][i];
        piece[1][i] = pieces_T[rotation][1][i];
      }
      break;
    case 5: //l
      for (uint8_t i = 0; i < 4; i++) {
        piece[0][i] = pieces_l[rotation][0][i];
        piece[1][i] = pieces_l[rotation][1][i];
      }
      break;
  }
}

boolean canRotate(short rotation) {
  short piece[2][4];
  copyPiece(piece, currentType, rotation);
  return !nextHorizontalCollision(piece, 0);
}

void drawLayout() {
  u8g2.drawHLine(0, 15, WIDTH);
  u8g2.drawFrame(0, 0, WIDTH, HEIGHT);
  drawNextPiece();
  char text[6];
  itoa(score, text, 10);
  drawText(text, getNumberLength(score), 7, 4);
}

short getNumberLength(int n) {
  short counter = 1;
  while (n >= 10) {
    n /= 10;
    counter++;
  }
  return counter;
}

void drawText(char text[], short length, int x, int y) {
  u8g2.setFont(u8g2_font_5x7_tr);
  u8g2.setCursor(x, y);

  for (short i = 0; i < length; i++)
    u8g2.write(text[i]);
}

int mainMenuIndex = 0;
int subMenuIndex = 0;
String mainMenuItems[] = {"Games", "Songs", "Camera"};
String subMenuItemsSongs[] = {"Songs", "Visualizer", "Back"};
String subMenuItemsGames[] = {"Tetris", "Pong", "Back"};
bool inSubMenu = false;
bool inSongs = false;
bool inGames = false;
bool inVisualizer = false;
int progress = 10;
int y = 50;
const unsigned long PADDLE_RATE = 33;
const unsigned long BALL_RATE = 16;
const uint8_t PADDLE_HEIGHT = 10;
#define SCREEN_WIDTH 128
#define SCREEN_HEIGHT 64

void drawCourt() {
    u8g2.drawFrame(0, 0, SCREEN_WIDTH, SCREEN_HEIGHT);
}

void pong() {
    uint8_t ball_x = 64, ball_y = 32;
    uint8_t ball_dir_x = 1, ball_dir_y = 1;
    unsigned long ball_update;

    unsigned long paddle_update;
    const uint8_t CPU_X = 12;
    uint8_t cpu_y = 16;

    const uint8_t PLAYER_X = 115;
    uint8_t player_y = 16;

    u8g2.begin();
    unsigned long start = millis();

    pinMode(UP_BUTTON_PIN, INPUT);
    pinMode(DOWN_BUTTON_PIN, INPUT);
    digitalWrite(UP_BUTTON_PIN, 1);
    digitalWrite(DOWN_BUTTON_PIN, 1);
    u8g2.clearBuffer();
    drawCourt();

    while (millis() - start < 2000);

    u8g2.sendBuffer();  // Send buffer to update the display
    ball_update = millis();
    paddle_update = ball_update;

    while (true) {
        bool update = false;
        unsigned long time = millis();

        static bool up_state = false;
        static bool down_state = false;

        up_state |= (digitalRead(UP_BUTTON_PIN) == LOW);
        down_state |= (digitalRead(DOWN_BUTTON_PIN) == LOW);

        if (time > ball_update) {
            uint8_t new_x = ball_x + ball_dir_x;
            uint8_t new_y = ball_y + ball_dir_y;

            if (new_x == 0 || new_x == 127) {
                ball_x = 64;
                ball_y = 32;
                new_x = 64;
                new_y = 32;
                cpu_y = 16;
                player_y = 16;
                ball_dir_x = 1;
                ball_dir_y = 1;
            }

            if (new_y == 0 || new_y == 63) {
                ball_dir_y = -ball_dir_y;
                new_y += ball_dir_y + ball_dir_y;
            }

            if (new_x == CPU_X && new_y >= cpu_y && new_y <= cpu_y + PADDLE_HEIGHT) {
                ball_dir_x = -ball_dir_x;
                new_x += ball_dir_x + ball_dir_x;
            }

            if (new_x == PLAYER_X && new_y >= player_y && new_y <= player_y + PADDLE_HEIGHT) {
                ball_dir_x = -ball_dir_x;
                new_x += ball_dir_x + ball_dir_x;
            }

            // Draw the ball at the new position
            u8g2.drawPixel(ball_x, ball_y);

            // Update the ball position
            ball_x = new_x;
            ball_y = new_y;

            ball_update += BALL_RATE;

            update = true;
        }

        if (time > paddle_update) {
            // Update the paddles
            const uint8_t half_paddle = PADDLE_HEIGHT >> 1;
            if (cpu_y + half_paddle > ball_y) {
                cpu_y -= 1;
            }
            if (cpu_y + half_paddle < ball_y) {
                cpu_y += 1;
            }
            if (cpu_y < 1) cpu_y = 1;
            if (cpu_y + PADDLE_HEIGHT > 63) cpu_y = 63 - PADDLE_HEIGHT;

            if (up_state) {
                player_y -= 1;
            }
            if (down_state) {
                player_y += 1;
            }
            up_state = down_state = false;
            if (player_y < 1) player_y = 1;
            if (player_y + PADDLE_HEIGHT > 63) player_y = 63 - PADDLE_HEIGHT;

            // Draw the paddles at the new positions
            for (int i = 0; i < PADDLE_HEIGHT; i++) {
                u8g2.drawPixel(CPU_X, cpu_y + i);
                u8g2.drawPixel(PLAYER_X, player_y + i);
            }

            update = true;
            paddle_update += PADDLE_RATE;
        }

        if (update) {
            // Send buffer to update the display
            u8g2.sendBuffer();
            drawCourt();
            delayMicroseconds(200);
            u8g2.clearBuffer();
            drawCourt();
        }
        if (digitalRead(SELECT_BUTTON_PIN) == LOW) {
          inSubMenu = true;
          inGames = true;
          break;
        }
    }
}

void setup() {
    pinMode(UP_BUTTON_PIN, INPUT_PULLUP);
    pinMode(DOWN_BUTTON_PIN, INPUT_PULLUP);
    pinMode(SELECT_BUTTON_PIN, INPUT_PULLUP);

    u8g2.begin();
}

void loop() {
   if (digitalRead(UP_BUTTON_PIN) == LOW) {
    if (inSubMenu && !inVisualizer) {
      subMenuIndex--;
      if (subMenuIndex < 0) {
        subMenuIndex = sizeof(subMenuItemsSongs)/sizeof(subMenuItemsSongs[0]) - 1;
      }
    } else if (!inSubMenu && !inVisualizer) {
      mainMenuIndex--;
      if (mainMenuIndex < 0) {
        mainMenuIndex = sizeof(mainMenuItems)/sizeof(mainMenuItems[0]) - 1;
      }
    }
    delay(200);
  }

  if (digitalRead(DOWN_BUTTON_PIN) == LOW) {
    if (inSubMenu && !inVisualizer) {
      subMenuIndex++;
      if (subMenuIndex >= sizeof(subMenuItemsSongs)/sizeof(subMenuItemsSongs[0])) {
        subMenuIndex = 0;
      }
    } else if (!inSubMenu && !inVisualizer) {
      mainMenuIndex++;
      if (mainMenuIndex >= sizeof(mainMenuItems)/sizeof(mainMenuItems[0])) {
        mainMenuIndex = 0;
      }
    }
    delay(200);
  }

  if (digitalRead(SELECT_BUTTON_PIN) == LOW) {
    if (!inSubMenu && !inVisualizer) {
      if (mainMenuItems[mainMenuIndex] == "Songs") {
        inSubMenu = true;
        inSongs = true;
      } else if (mainMenuItems[mainMenuIndex] == "Games") {
        inSubMenu = true;
        inGames = true;
      } else {
        // Add your camera code here
      }
    } else if (inGames && subMenuItemsGames[subMenuIndex] == "Pong") {
      pong();
    } else if (inGames && subMenuItemsGames[subMenuIndex] == "Tetris"){
      tetrisSetup();
      tetrisLoop();
    } else if (inSongs && subMenuItemsSongs[subMenuIndex] == "Visualizer" && !inVisualizer) {
      inVisualizer = true;
    } else if (inSongs && subMenuItemsSongs[subMenuIndex] == "Back") {
      inSubMenu = false;
      inSongs = false;
    } else if (inGames && subMenuItemsGames[subMenuIndex] == "Back") {
      inSubMenu = false;
      inGames = false;
    } else if (inVisualizer) {
      inVisualizer = false;
    }
    delay(200);
  }


    u8g2.clearBuffer();
    u8g2.setFont(u8g2_font_ncenB08_tr);

    if (!inSubMenu && !inVisualizer) {
        for (int i = 0; i < sizeof(mainMenuItems) / sizeof(mainMenuItems[0]); i++) {
            if (i == mainMenuIndex) {
                u8g2.drawStr(0, i * 10 + 10, (">" + mainMenuItems[i]).c_str());
            } else {
                u8g2.drawStr(10, i * 10 + 10, mainMenuItems[i].c_str());
            }
        }
    } else if (inSubMenu && !inVisualizer) {
        if (inSongs) {
            for (int i = 0; i < sizeof(subMenuItemsSongs) / sizeof(subMenuItemsSongs[0]); i++) {
                if (i == subMenuIndex) {
                    u8g2.drawStr(0, i * 10 + 10, (">" + subMenuItemsSongs[i]).c_str());
                } else {
                    u8g2.drawStr(10, i * 10 + 10, subMenuItemsSongs[i].c_str());
                }
            }
        } else if (inGames) {
            for (int i = 0; i < sizeof(subMenuItemsGames) / sizeof(subMenuItemsGames[0]); i++) {
                if (i == subMenuIndex) {
                    u8g2.drawStr(0, i * 10 + 10, (">" + subMenuItemsGames[i]).c_str());
                } else {
                    u8g2.drawStr(10, i * 10 + 10, subMenuItemsGames[i].c_str());
                }
            }
        }
    } else if (inVisualizer) {
        // Your visualizer code here
        int readup = digitalRead(UP_BUTTON_PIN);
        int readdown = digitalRead(DOWN_BUTTON_PIN);
        if (readup == LOW) {
            progress = progress + 1;
            y = y - 1;
        }
        if (readdown == LOW) {
            progress = progress - 1;
            y = y + 1;
        }
        u8g2.setBitmapMode(1);
        u8g2.drawFrame(55, 5, 20, 57);
        u8g2.drawBox(57, y, 16, progress);
        if (progress > 52) {
            progress = 53;
            y = 7;
        }
        if (progress < 2) {
            progress = 1;
            y = 57;
        }
    }
    u8g2.sendBuffer();
}