#include <Wire.h>
#include <Adafruit_SSD1306.h>

Adafruit_SSD1306 display(128, 64, &Wire, -1);

// Buttons
#define BTN_UP     27
#define BTN_DOWN   25
#define BTN_LEFT   32
#define BTN_RIGHT  14
#define BTN_SELECT 26

// Lanes
int laneX[4] = {10, 40, 70, 100};
int silhouetteY = 52;

// Arrow speed
float arrowSpeed = 0.05;

// --------- NOTE STRUCT ----------
struct Note {
  unsigned long time;
  int lane;
  bool active;
};

Note notes[128];

// --------- BPM / BEAT / DIFFICULTY ----------
int bpm = 120;
int beatDiv = 1;
unsigned long spawnInterval = 500;

unsigned long lastSpawn = 0;

int difficultyLevel = 1;
unsigned long nextDifficultyTime = 0;

// --------- HEALTH / SURVIVAL ----------
int health = 100;
int maxHealth = 100;
bool isDead = false;

unsigned long survivalStart = 0;
unsigned long survivalTime = 0;

// --------- HIT ZONE / JUDGMENT ----------
int hitTop = 40;
int hitBottom = 64;

String lastJudgment = "";
unsigned long judgmentTimer = 0;

// --------- INPUT ----------
bool lanePressed(int lane) {
  if (lane == 0) return digitalRead(BTN_LEFT)  == LOW;
  if (lane == 1) return digitalRead(BTN_DOWN)  == LOW;
  if (lane == 2) return digitalRead(BTN_UP)    == LOW;
  if (lane == 3) return digitalRead(BTN_RIGHT) == LOW;
  return false;
}

// --------- DRAW ----------
void drawSilhouettes() {
  for (int i = 0; i < 4; i++) {
    display.drawRect(laneX[i], silhouetteY, 12, 12, SSD1306_WHITE);
  }
}

void drawArrow(int lane, int y) {
  int x = laneX[lane];
  display.fillTriangle(x+6, y, x, y+12, x+12, y+12, SSD1306_WHITE);
}

void drawArrows() {
  unsigned long now = millis() - survivalStart;

  for (int i = 0; i < 128; i++) {
    if (!notes[i].active) continue;

    int y = (int)((long)(now - notes[i].time) * arrowSpeed);

    if (y >= 0 && y < 64) {
      drawArrow(notes[i].lane, y);
    }
  }
}

void drawJudgment() {
  if (millis() - judgmentTimer < 400) {
    display.setCursor(40, 20);
    display.setTextSize(1);
    display.setTextColor(SSD1306_WHITE);
    display.println(lastJudgment);
  }
}

void drawHealth() {
  display.drawRect(0, 0, 128, 6, SSD1306_WHITE);
  int w = map(health, 0, maxHealth, 0, 126);
  display.fillRect(1, 1, w, 4, SSD1306_WHITE);
}

void drawSurvivalTime() {
  display.setCursor(0, 10);
  display.setTextSize(1);
  display.print("Time: ");
  display.print(survivalTime / 1000);
  display.print("s");
}

void drawDifficulty() {
  display.setCursor(0, 20);
  display.setTextSize(1);
  display.print("Diff: ");
  display.print(difficultyLevel);
}

// --------- RANDOM ARROW SPAWN ----------
void spawnArrow() {
  int lane = random(0, 4);

  for (int i = 0; i < 128; i++) {
    if (!notes[i].active) {
      notes[i].active = true;
      notes[i].lane = lane;
      notes[i].time = millis() - survivalStart + 800;
      return;
    }
  }
}

// --------- HIT DETECTION ----------
void checkHits() {
  unsigned long now = millis() - survivalStart;

  for (int i = 0; i < 128; i++) {
    if (!notes[i].active) continue;

    int y = (int)((long)(now - notes[i].time) * arrowSpeed);

    if (y < 0) continue;

    // HIT
    if (y >= hitTop && y <= hitBottom) {
      if (lanePressed(notes[i].lane)) {
        lastJudgment = "SICK!";
        judgmentTimer = millis();

        health += 3;
        if (health > maxHealth) health = maxHealth;

        notes[i].active = false;
        return;
      }
    }

    // MISS
    if (y > hitBottom + 5 && notes[i].active) {
      lastJudgment = "MISS!";
      judgmentTimer = millis();

      health -= 10;
      if (health <= 0) isDead = true;

      notes[i].active = false;
    }
  }
}

// --------- BPM SELECTION MENU ----------
void bpmMenu() {
  bpm = 120;
  beatDiv = 1;

  while (true) {
    display.clearDisplay();
    display.setTextSize(1);
    display.setCursor(0, 0);
    display.println("Set BPM (UP/DOWN):");
    display.print("> ");
    display.println(bpm);

    display.println("");
    display.println("Beat division (LEFT/RIGHT):");
    display.print("> 1/");
    display.println(beatDiv);

    display.println("");
    display.println("Press SELECT to start");

    display.display();

    if (digitalRead(BTN_UP) == LOW) {
      bpm += 5;
      delay(150);
    }
    if (digitalRead(BTN_DOWN) == LOW) {
      bpm -= 5;
      if (bpm < 30) bpm = 30;
      delay(150);
    }
    if (digitalRead(BTN_LEFT) == LOW) {
      beatDiv++;
      if (beatDiv > 8) beatDiv = 8;
      delay(150);
    }
    if (digitalRead(BTN_RIGHT) == LOW) {
      beatDiv--;
      if (beatDiv < 1) beatDiv = 1;
      delay(150);
    }
    if (digitalRead(BTN_SELECT) == LOW) {
      delay(300);
      break;
    }
  }

  spawnInterval = (60000 / bpm) / beatDiv;
}

// --------- GAMEPLAY ----------
void startSurvival() {
  survivalStart = millis();
  health = maxHealth;
  isDead = false;
  survivalTime = 0;

  difficultyLevel = 1;
  nextDifficultyTime = millis() + 10000;

  for (int i = 0; i < 128; i++) notes[i].active = false;

  lastSpawn = millis();

  while (true) {
    unsigned long now = millis();
    survivalTime = now - survivalStart;

    // difficulty scaling
    if (now > nextDifficultyTime) {
      difficultyLevel++;
      bpm += 5;
      spawnInterval = (60000 / bpm) / beatDiv;
      nextDifficultyTime = now + 10000;
    }

    display.clearDisplay();

    drawSilhouettes();
    drawArrows();
    checkHits();
    drawJudgment();
    drawHealth();
    drawSurvivalTime();
    drawDifficulty();

    if (now - lastSpawn > spawnInterval) {
      spawnArrow();
      lastSpawn = now;
    }

    if (isDead) {
      display.clearDisplay();
      display.setCursor(10, 20);
      display.setTextSize(2);
      display.println("GAME OVER");
      display.display();
      delay(2000);
      return;
    }

    display.display();
  }
}

// --------- SETUP ----------
void setup() {
  pinMode(BTN_UP, INPUT_PULLUP);
  pinMode(BTN_DOWN, INPUT_PULLUP);
  pinMode(BTN_LEFT, INPUT_PULLUP);
  pinMode(BTN_RIGHT, INPUT_PULLUP);
  pinMode(BTN_SELECT, INPUT_PULLUP);

  Wire.begin(21, 22);
  display.begin(SSD1306_SWITCHCAPVCC, 0x3C);

  randomSeed(analogRead(0));
}

// --------- LOOP ----------
void loop() {
  bpmMenu();       // NOW the menu appears first
  startSurvival(); // THEN the game starts
}