#include <Arduino.h>
#include <Wire.h>
#include <U8g2lib.h>

U8G2_SSD1306_128X64_NONAME_1_HW_I2C u8g2(U8G2_R0, U8X8_PIN_NONE);

// -------------------- Buttons --------------------
const int BTN_RIGHT = 2;
const int BTN_LEFT  = 3;
const int BTN_FIRE  = 4;

// -------------------- Screen ---------------------
const int SCREEN_W = 128;
const int SCREEN_H = 64;

// -------------------- Player ---------------------
int playerX = 60;
const int playerY = 58;
const int playerW = 10;
const int playerH = 4;

// -------------------- Bullet ---------------------
bool bulletActive = false;
int bulletX = 0;
int bulletY = 0;

// fire button edge detect
bool lastFireState = HIGH;

// -------------------- Enemy grid -----------------
const int ROWS = 3;
const int COLS = 5;
bool enemyAlive[ROWS][COLS];

int enemyOffsetX = 14;
int enemyOffsetY = 10;
int enemyDir = 1;
int enemyStepDown = 4;
int enemyAnim = 0;

// -------------------- Game state -----------------
int level = 1;
int score = 0;
bool gameStarted = false;
bool gameOver = false;
bool levelCleared = false;
bool gameWon = false;

// -------------------- Timing ---------------------
unsigned long lastEnemyMove = 0;
unsigned long lastAnimToggle = 0;
unsigned long lastBulletMove = 0;

// speed per level
int enemyMoveInterval = 500;

// -------------------- Helpers --------------------
void initEnemies() {
  for (int r = 0; r < ROWS; r++) {
    for (int c = 0; c < COLS; c++) {
      enemyAlive[r][c] = true;
    }
  }
  enemyOffsetX = 14;
  enemyOffsetY = 10;
  enemyDir = 1;
}

void setupLevel(int newLevel) {
  level = newLevel;
  playerX = 60;
  bulletActive = false;
  levelCleared = false;
  gameOver = false;
  gameWon = false;

  initEnemies();

  enemyMoveInterval = 500 - (level - 1) * 70;
  if (enemyMoveInterval < 160) enemyMoveInterval = 160;

  lastEnemyMove = millis();
  lastAnimToggle = millis();
  lastBulletMove = millis();
}

void resetGame() {
  score = 0;
  gameStarted = true;
  setupLevel(1);
}

int livingEnemies() {
  int count = 0;
  for (int r = 0; r < ROWS; r++) {
    for (int c = 0; c < COLS; c++) {
      if (enemyAlive[r][c]) count++;
    }
  }
  return count;
}

void fireBullet() {
  if (!bulletActive) {
    bulletActive = true;
    bulletX = playerX + playerW / 2;
    bulletY = playerY - 3;
  }
}

void drawPlayer() {
  u8g2.drawBox(playerX + 2, playerY, 6, 3);
  u8g2.drawPixel(playerX + 1, playerY + 1);
  u8g2.drawPixel(playerX + 8, playerY + 1);
  u8g2.drawLine(playerX + 4, playerY - 2, playerX + 5, playerY - 2);
  u8g2.drawLine(playerX + 3, playerY - 1, playerX + 6, playerY - 1);
}

void drawEnemy(int x, int y, int frame) {
  if (frame == 0) {
    u8g2.drawPixel(x + 1, y);
    u8g2.drawPixel(x + 6, y);
    u8g2.drawLine(x + 2, y + 1, x + 5, y + 1);
    u8g2.drawBox(x, y + 2, 8, 2);
    u8g2.drawPixel(x + 1, y + 4);
    u8g2.drawPixel(x + 6, y + 4);
    u8g2.drawPixel(x, y + 5);
    u8g2.drawPixel(x + 2, y + 5);
    u8g2.drawPixel(x + 5, y + 5);
    u8g2.drawPixel(x + 7, y + 5);
  } else {
    u8g2.drawPixel(x + 1, y);
    u8g2.drawPixel(x + 6, y);
    u8g2.drawLine(x + 2, y + 1, x + 5, y + 1);
    u8g2.drawBox(x, y + 2, 8, 2);
    u8g2.drawPixel(x + 2, y + 4);
    u8g2.drawPixel(x + 5, y + 4);
    u8g2.drawPixel(x + 1, y + 5);
    u8g2.drawPixel(x + 6, y + 5);
  }
}

void drawEnemies() {
  for (int r = 0; r < ROWS; r++) {
    for (int c = 0; c < COLS; c++) {
      if (enemyAlive[r][c]) {
        int ex = enemyOffsetX + c * 20;
        int ey = enemyOffsetY + r * 12;
        drawEnemy(ex, ey, enemyAnim);
      }
    }
  }
}

void moveEnemies() {
  bool hitEdge = false;

  for (int r = 0; r < ROWS; r++) {
    for (int c = 0; c < COLS; c++) {
      if (!enemyAlive[r][c]) continue;
      int ex = enemyOffsetX + c * 20;
      if ((enemyDir == 1 && ex + 8 >= SCREEN_W - 2) ||
          (enemyDir == -1 && ex <= 2)) {
        hitEdge = true;
      }
    }
  }

  if (hitEdge) {
    enemyDir = -enemyDir;
    enemyOffsetY += enemyStepDown;
  } else {
    enemyOffsetX += enemyDir * 3;
  }
}

void updateBullet() {
  if (!bulletActive) return;

  bulletY -= 4;

  // turn off before it reaches top edge
  if (bulletY < 2) {
    bulletActive = false;
    return;
  }

  for (int r = 0; r < ROWS; r++) {
    for (int c = 0; c < COLS; c++) {
      if (!enemyAlive[r][c]) continue;

      int ex = enemyOffsetX + c * 20;
      int ey = enemyOffsetY + r * 12;

      if (bulletX >= ex && bulletX <= ex + 8 &&
          bulletY >= ey && bulletY <= ey + 6) {
        enemyAlive[r][c] = false;
        bulletActive = false;
        score += 10;
        return;
      }
    }
  }
}

bool enemiesReachedPlayer() {
  for (int r = 0; r < ROWS; r++) {
    for (int c = 0; c < COLS; c++) {
      if (!enemyAlive[r][c]) continue;
      int ey = enemyOffsetY + r * 12;
      if (ey + 6 >= playerY - 2) return true;
    }
  }
  return false;
}

void drawHUD() {
  char buf[24];
  u8g2.setFont(u8g2_font_5x7_tr);
  sprintf(buf, "L:%d S:%d", level, score);
  u8g2.drawStr(2, 7, buf);
}

void drawBullet() {
  if (bulletActive && bulletY >= 2 && bulletY < 64) {
    u8g2.drawBox(bulletX, bulletY, 1, 3);
  }
}

void drawStars() {
  for (int i = 0; i < 8; i++) {
    int sx = (i * 17 + millis() / 20) % 128;
    int sy = (i * 11 + level * 7) % 64;
    u8g2.drawPixel(sx, sy);
  }
}

void drawStartScreen() {
  u8g2.firstPage();
  do {
    drawStars();
    u8g2.setFont(u8g2_font_6x10_tr);
    u8g2.drawStr(20, 16, "SPACE INVADERS");
    u8g2.drawStr(38, 30, "5 LEVELS");
    u8g2.setFont(u8g2_font_5x7_tr);
    u8g2.drawStr(22, 42, "D3 LEFT  D2 RIGHT");
    u8g2.drawStr(24, 50, "D4 SHOOT / START");
    u8g2.drawStr(28, 61, "PRESS D4 TO PLAY");
  } while (u8g2.nextPage());
}

void drawLevelClearScreen() {
  u8g2.firstPage();
  do {
    drawStars();
    u8g2.setFont(u8g2_font_6x10_tr);
    u8g2.drawStr(24, 22, "LEVEL CLEAR!");
    char buf[20];
    sprintf(buf, "LEVEL %d DONE", level);
    u8g2.drawStr(30, 36, buf);
    u8g2.setFont(u8g2_font_5x7_tr);
    u8g2.drawStr(26, 54, "PRESS D4 FOR NEXT");
  } while (u8g2.nextPage());
}

void drawGameOverScreen() {
  u8g2.firstPage();
  do {
    u8g2.setFont(u8g2_font_6x10_tr);
    u8g2.drawStr(30, 20, "GAME OVER");
    char buf1[20];
    char buf2[20];
    sprintf(buf1, "LEVEL: %d", level);
    sprintf(buf2, "SCORE: %d", score);
    u8g2.drawStr(36, 34, buf1);
    u8g2.drawStr(36, 46, buf2);
    u8g2.setFont(u8g2_font_5x7_tr);
    u8g2.drawStr(22, 60, "PRESS D4 TO RESTART");
  } while (u8g2.nextPage());
}

void drawWinScreen() {
  u8g2.firstPage();
  do {
    drawStars();
    u8g2.setFont(u8g2_font_6x10_tr);
    u8g2.drawStr(22, 18, "YOU WIN!");
    u8g2.drawStr(18, 32, "ALL 5 LEVELS");
    char buf[20];
    sprintf(buf, "SCORE: %d", score);
    u8g2.drawStr(30, 46, buf);
    u8g2.setFont(u8g2_font_5x7_tr);
    u8g2.drawStr(22, 60, "PRESS D4 TO RESTART");
  } while (u8g2.nextPage());
}

void drawGame() {
  u8g2.firstPage();
  do {
    drawStars();
    drawHUD();
    drawEnemies();
    drawBullet();
    drawPlayer();
    u8g2.drawFrame(0, 0, 128, 64);
  } while (u8g2.nextPage());
}

void setup() {
  pinMode(BTN_RIGHT, INPUT_PULLUP);
  pinMode(BTN_LEFT, INPUT_PULLUP);
  pinMode(BTN_FIRE, INPUT_PULLUP);

  u8g2.begin();
  randomSeed(analogRead(A0));
}

void loop() {
  bool fireNow = digitalRead(BTN_FIRE);

  if (!gameStarted) {
    drawStartScreen();
    if (fireNow == LOW && lastFireState == HIGH) {
      resetGame();
    }
    lastFireState = fireNow;
    return;
  }

  if (gameOver) {
    drawGameOverScreen();
    if (fireNow == LOW && lastFireState == HIGH) {
      resetGame();
    }
    lastFireState = fireNow;
    return;
  }

  if (gameWon) {
    drawWinScreen();
    if (fireNow == LOW && lastFireState == HIGH) {
      resetGame();
    }
    lastFireState = fireNow;
    return;
  }

  if (levelCleared) {
    drawLevelClearScreen();
    if (fireNow == LOW && lastFireState == HIGH) {
      if (level >= 5) {
        gameWon = true;
        levelCleared = false;
      } else {
        setupLevel(level + 1);
      }
    }
    lastFireState = fireNow;
    return;
  }

  // ---------- Controls ----------
  if (digitalRead(BTN_LEFT) == LOW) {
    playerX -= 3;
  }
  if (digitalRead(BTN_RIGHT) == LOW) {
    playerX += 3;
  }

  if (playerX < 1) playerX = 1;
  if (playerX > SCREEN_W - playerW - 1) playerX = SCREEN_W - playerW - 1;

  // fire only once per press
  if (fireNow == LOW && lastFireState == HIGH && !bulletActive) {
    fireBullet();
  }
  lastFireState = fireNow;

  // ---------- Timed updates ----------
  unsigned long now = millis();

  if (now - lastAnimToggle > 180) {
    enemyAnim = !enemyAnim;
    lastAnimToggle = now;
  }

  if (now - lastEnemyMove > (unsigned long)enemyMoveInterval) {
    moveEnemies();
    lastEnemyMove = now;
  }

  if (now - lastBulletMove > 35) {
    updateBullet();
    lastBulletMove = now;
  }

  // ---------- State checks ----------
  if (livingEnemies() == 0) {
    levelCleared = true;
  }

  if (enemiesReachedPlayer()) {
    gameOver = true;
  }

  drawGame();
}