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

#define SCREEN_WIDTH 128
#define SCREEN_HEIGHT 64
Adafruit_SSD1306 display(SCREEN_WIDTH, SCREEN_HEIGHT, &Wire, -1);

// Pin mapping untuk XIAO ESP32-C3
const int PIN_LEFT  = D0;
const int PIN_UP    = D1;
const int PIN_RIGHT = D2;
const int PIN_DOWN  = D3;
const int PIN_A     = D10;  // hold -> dodge right
const int PIN_B     = D9;   // hold -> dodge left
const int PIN_BUZZ  = D7;

// Game state
enum GameState { MENU, RUNNING, GAMEOVER };
GameState state = MENU;

// Grid config
const int COLS = 9;
const int ROWS = 5;
int cellW, cellH;
int gridOffsetX, gridOffsetY;

// Player
int playerGroup = 1;   // 0=A,1=B,2=C
int playerSub   = 1;   // 0,1,2 (default 1)

// Buttons (debounce utk D-pad single click)
bool lastLeft = HIGH, lastRight = HIGH;
unsigned long lastChangeLeft=0, lastChangeRight=0;
const unsigned long DEBOUNCE_MS = 40;

// Obstacles
struct Ob {
  bool active;
  uint8_t group;
  uint8_t row;
  uint8_t mask;
};
const int MAX_OBS = 8;
Ob obs[MAX_OBS];

// Coins
struct Coin {
  bool active;
  uint8_t group;
  uint8_t row;
  uint8_t sub;
};
const int MAX_COINS = 6;
Coin coins[MAX_COINS];

// Timing
unsigned long lastStep = 0;
const unsigned long STEP_MS = 300;
unsigned long lastSpawnObs = 0, lastSpawnCoin = 0;
const unsigned long MIN_SPAWN_MS = 700;
const unsigned long MAX_SPAWN_MS = 1400;
const unsigned long MIN_COIN_MS = 1200;
const unsigned long MAX_COIN_MS = 3000;
unsigned long lastFrame = 0;
const unsigned long FRAME_MS = 50;

int score = 0;

void setup() {
  Serial.begin(115200);

  // I2C OLED pakai pin D4 (SDA) dan D5 (SCL)
  Wire.begin(D4, D5);
  if (!display.begin(SSD1306_SWITCHCAPVCC, 0x3C)) {
    Serial.println("SSD1306 fail");
    for(;;);
  }
  display.clearDisplay(); display.display();

  pinMode(PIN_LEFT, INPUT_PULLUP);
  pinMode(PIN_UP, INPUT_PULLUP);
  pinMode(PIN_RIGHT, INPUT_PULLUP);
  pinMode(PIN_DOWN, INPUT_PULLUP);
  pinMode(PIN_A, INPUT_PULLUP);
  pinMode(PIN_B, INPUT_PULLUP);
  pinMode(PIN_BUZZ, OUTPUT);
  digitalWrite(PIN_BUZZ, LOW);

  randomSeed(analogRead(A0)); // pakai ADC buat random seed

  cellW = SCREEN_WIDTH / COLS;
  cellH = SCREEN_HEIGHT / ROWS;
  gridOffsetX = (SCREEN_WIDTH - (cellW * COLS)) / 2;
  gridOffsetY = (SCREEN_HEIGHT - (cellH * ROWS)) / 2;

  for (int i=0;i<MAX_OBS;i++) obs[i].active = false;
  for (int i=0;i<MAX_COINS;i++) coins[i].active = false;
}

void loop() {
  unsigned long now = millis();

  // Baca tombol
  bool leftRaw  = digitalRead(PIN_LEFT) == LOW;
  bool rightRaw = digitalRead(PIN_RIGHT) == LOW;
  bool upRaw    = digitalRead(PIN_UP) == LOW;
  bool downRaw  = digitalRead(PIN_DOWN) == LOW;
  bool aRaw     = digitalRead(PIN_A) == LOW;
  bool bRaw     = digitalRead(PIN_B) == LOW;

  // D-Pad (click only, edge detection)
  if (leftRaw != lastLeft && now - lastChangeLeft > DEBOUNCE_MS) {
    lastChangeLeft = now;
    if (!lastLeft && leftRaw) {
      if (playerGroup > 0) playerGroup--;
      playerSub = 1;
      tone(PIN_BUZZ, 700, 50);
    }
    lastLeft = leftRaw;
  }
  if (rightRaw != lastRight && now - lastChangeRight > DEBOUNCE_MS) {
    lastChangeRight = now;
    if (!lastRight && rightRaw) {
      if (playerGroup < 2) playerGroup++;
      playerSub = 1;
      tone(PIN_BUZZ, 700, 50);
    }
    lastRight = rightRaw;
  }

  // A/B (hold → continuous dodge)
  if (aRaw) playerSub = 2;       // hold kanan
  else if (bRaw) playerSub = 0;  // hold kiri
  else playerSub = 1;            // release → tengah

  // Menu handling
  if (state == MENU) {
    drawMenu();
    if (aRaw) startGame();
    return;
  }

  if (state == GAMEOVER) {
    drawGameOver();
    if (aRaw) state = MENU;
    return;
  }

  // Running
  if (now - lastStep >= STEP_MS) {
    lastStep = now;
    updateGame();
  }

  // Spawns
  if (now - lastSpawnObs >= (unsigned long)random(MIN_SPAWN_MS, MAX_SPAWN_MS)) {
    lastSpawnObs = now; spawnObstacle();
  }
  if (now - lastSpawnCoin >= (unsigned long)random(MIN_COIN_MS, MAX_COIN_MS)) {
    lastSpawnCoin = now; spawnCoin();
  }

  // Render
  if (now - lastFrame >= FRAME_MS) {
    lastFrame = now;
    render();
  }
}

// ---------------- Game logic ----------------
uint8_t randomPattern() {
  int r = random(0,100);
  if (r < 30) return 0b111;
  int choose = random(0,3);
  if (choose == 0) return 0b011;
  if (choose == 1) return 0b110;
  return 0b101;
}
int activeGroupCount() {
  bool used[3] = {false,false,false};
  for (int i=0;i<MAX_OBS;i++) if (obs[i].active) used[obs[i].group] = true;
  int cnt=0; for (int g=0;g<3;g++) if (used[g]) cnt++;
  return cnt;
}
bool groupActive(uint8_t g) {
  for (int i=0;i<MAX_OBS;i++) if (obs[i].active && obs[i].group==g) return true;
  return false;
}
int findFreeObs() { for (int i=0;i<MAX_OBS;i++) if (!obs[i].active) return i; return -1; }
int findFreeCoin() { for (int i=0;i<MAX_COINS;i++) if (!coins[i].active) return i; return -1; }

void spawnObstacle() {
  int candidates[3]; int c=0;
  for (int g=0; g<3; g++){
    if (activeGroupCount() < 2 || groupActive(g)) candidates[c++] = g;
  }
  if (c==0) return;
  int pick = candidates[random(0,c)];
  int idx = findFreeObs();
  if (idx < 0) return;
  obs[idx].active = true;
  obs[idx].group = pick;
  obs[idx].row = 0;
  obs[idx].mask = randomPattern();
}
void spawnCoin() {
  int idx = findFreeCoin();
  if (idx < 0) return;
  coins[idx].active = true;
  coins[idx].group = random(0,3);
  coins[idx].row = 0;
  coins[idx].sub = random(0,3);
}
void updateGame() {
  // Move obs
  for (int i=0;i<MAX_OBS;i++){
    if (!obs[i].active) continue;
    obs[i].row++;
    if (obs[i].row >= ROWS) obs[i].active = false;
  }
  // Move coins
  for (int i=0;i<MAX_COINS;i++){
    if (!coins[i].active) continue;
    coins[i].row++;
    if (coins[i].row >= ROWS) coins[i].active = false;
  }
  // Collision check
  for (int i=0;i<MAX_OBS;i++){
    if (!obs[i].active) continue;
    if (obs[i].row == ROWS-1) {
      int baseCol = obs[i].group * 3;
      for (int s=0;s<3;s++){
        if (obs[i].mask & (1<<s)) {
          int colIndex = baseCol + s;
          int playerCol = playerGroup*3 + playerSub;
          if (colIndex == playerCol) {
            state = GAMEOVER;
            tone(PIN_BUZZ, 300, 400);
            return;
          }
        }
      }
    }
  }
  // Coin check
  for (int i=0;i<MAX_COINS;i++){
    if (!coins[i].active) continue;
    if (coins[i].row == ROWS-1) {
      int colIndex = coins[i].group*3 + coins[i].sub;
      int playerCol = playerGroup*3 + playerSub;
      bool blocked=false;
      for (int j=0;j<MAX_OBS;j++){
        if (!obs[j].active) continue;
        if (obs[j].row==ROWS-1 && obs[j].group==coins[i].group && (obs[j].mask & (1<<coins[i].sub))) blocked=true;
      }
      if (!blocked && colIndex==playerCol) {
        score++;
        coins[i].active=false;
        tone(PIN_BUZZ, 1200, 60);
      } else {
        coins[i].active=false;
      }
    }
  }
}

// ---------------- UI ----------------
void drawMenu() {
  display.clearDisplay();
  display.setTextSize(2); display.setTextColor(SSD1306_WHITE);
  display.setCursor(18, 10); display.println("GridRun");
  display.setTextSize(1); display.setCursor(6, 38);
  display.println("Hold A/B to dodge");
  display.setCursor(6, 50);
  display.println("Press A to start");
  display.display();
}
void startGame() {
  for (int i=0;i<MAX_OBS;i++) obs[i].active=false;
  for (int i=0;i<MAX_COINS;i++) coins[i].active=false;
  playerGroup=1; playerSub=1; score=0;
  lastStep=millis(); lastSpawnObs=millis(); lastSpawnCoin=millis();
  state=RUNNING;
}
void drawGameOver() {
  display.clearDisplay();
  display.setTextSize(2); display.setTextColor(SSD1306_WHITE);
  display.setCursor(12, 12); display.println("GAME OVER");
  display.setTextSize(1); display.setCursor(8, 50);
  display.printf("Score: %d  A->Menu", score);
  display.display();
}

// ---------------- Render ----------------
void render() {
  display.clearDisplay();
  // Obstacles
  for (int i=0;i<MAX_OBS;i++){
    if (!obs[i].active) continue;
    int baseCol=obs[i].group*3;
    for (int s=0;s<3;s++){
      if (obs[i].mask&(1<<s)){
        int col=baseCol+s;
        int x=gridOffsetX+col*cellW+2;
        int y=gridOffsetY+obs[i].row*cellH+2;
        display.fillRect(x,y,cellW-4,cellH-4,SSD1306_WHITE);
      }
    }
  }
  // Coins
  for (int i=0;i<MAX_COINS;i++){
    if (!coins[i].active) continue;
    int col=coins[i].group*3+coins[i].sub;
    int x=gridOffsetX+col*cellW+4;
    int y=gridOffsetY+coins[i].row*cellH+4;
    display.drawRect(x,y,cellW-8,cellH-8,SSD1306_WHITE);
    display.fillRect(x+(cellW-8)/2-1,y+(cellH-8)/2-1,2,2,SSD1306_WHITE);
  }
  // Player
  int pCol=playerGroup*3+playerSub;
  int px=gridOffsetX+pCol*cellW+3;
  int py=gridOffsetY+(ROWS-1)*cellH+3;
  display.fillRect(px,py,cellW-6,cellH-6,SSD1306_WHITE);
  // HUD
  display.setTextSize(1);
  display.setCursor(2,2); display.printf("Score:%d",score);
  display.display();
}