#include <Adafruit_GFX.h>
#include <Adafruit_SSD1306.h>
#include <Wire.h>
#include <avr/pgmspace.h> // Necessary for PROGMEM functions

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

// Button pins
#define PIN_L 2
#define PIN_R 3
#define PIN_A 4
#define PIN_B 5

// Menu - Using PROGMEM syntax to save SRAM for global array
const char *const menuItems[] PROGMEM = {PSTR("Pong"), PSTR("Snake")};
int menuIndex = 0;
const int MENU_COUNT = 2;

// Button helpers
bool btnPressed(uint8_t pin) { return digitalRead(pin) == LOW; }

// =============================
// Debounce helper
bool wasPressed(uint8_t pin) {
  static unsigned long lastTime[14] = {0};
  static bool lastState[14] = {HIGH};
  bool cur = !digitalRead(pin);
  int idx = pin;
  unsigned long now = millis();
  if (cur != lastState[idx] && now - lastTime[idx] > 40) {
    lastState[idx] = cur;
    lastTime[idx] = now;
    if (cur) return true;
  }
  return false;
}

// =============================
// Setup
void setup() {
  pinMode(PIN_L, INPUT_PULLUP);
  pinMode(PIN_R, INPUT_PULLUP);
  pinMode(PIN_A, INPUT_PULLUP);
  pinMode(PIN_B, INPUT_PULLUP);

  if (!display.begin(SSD1306_SWITCHCAPVCC, 0x3C)) while (true);
  display.clearDisplay();
  display.display();
  delay(50);
  showSplash();
}

// Splash
void showSplash() {
  display.clearDisplay();
  display.setTextSize(2);
  display.setTextColor(WHITE);
  display.setCursor(14, 22);
  display.println(F("ikoGames")); // Stored in Flash
  display.display();
  delay(1400);
}

// =============================
// Show menu
int showMenuLoop() {
  while (true) {
    display.clearDisplay();
    display.setTextSize(1);
    display.setCursor(0,0);
    display.println(F("== ikoGames ==")); // Stored in Flash

    for(int i=0;i<MENU_COUNT;i++){
      // Read the string pointer from Flash
      char* menuText = (char*)pgm_read_ptr(&(menuItems[i])); 
      
      if(i==menuIndex) display.print(F("> ")); else display.print(F("  ")); 
      
      // Print the string
      display.println(menuText);
    }
    display.display();

    if(wasPressed(PIN_L)){
      menuIndex--;
      if(menuIndex<0) menuIndex=MENU_COUNT-1;
    }
    if(wasPressed(PIN_R)){
      menuIndex++;
      if(menuIndex>=MENU_COUNT) menuIndex=0;
    }
    if(wasPressed(PIN_A)) return menuIndex;
    delay(80);
  }
}

// =============================
// PONG
void pongGame() {
  int paddleY = 22, ballX=64, ballY=32, dx=2, dy=1;
  const int paddleH=20;
  unsigned long t=millis();

  while(true){
    if(wasPressed(PIN_B)) return;

    if(btnPressed(PIN_L)) paddleY-=2;
    if(btnPressed(PIN_R)) paddleY+=2;
    if(paddleY<0) paddleY=0;
    if(paddleY>SCREEN_HEIGHT-paddleH) paddleY=SCREEN_HEIGHT-paddleH;

    if(millis()-t>20){
      t=millis();
      ballX+=dx; ballY+=dy;
      if(ballY<=0||ballY>=SCREEN_HEIGHT-1) dy=-dy;

      // paddle collision
      if(ballX<=8 && ballY>=paddleY && ballY<=paddleY+paddleH){
        dx=-dx;
        int hit=ballY-(paddleY+paddleH/2);
        dy=constrain(hit/6,-3,3);
        ballX=10;
      }

      if(ballX>=SCREEN_WIDTH-3) dx=-dx;

      if(ballX<-4 || ballX>SCREEN_WIDTH+4){
        display.clearDisplay();
        display.setTextSize(1);
        display.setCursor(18,25);
        display.println(F("Game Over - B to return")); 
        display.display();
        while(!wasPressed(PIN_B)) delay(30);
        return;
      }
    }

    display.clearDisplay();
    display.fillRect(2,paddleY,6,paddleH,WHITE);
    for(int i=0;i<SCREEN_HEIGHT;i+=6) display.drawFastVLine(64,i,3,WHITE);
    display.fillRect(ballX-1,ballY-1,3,3,WHITE);
    display.display();
    delay(5);
  }
}

// =============================
// SNAKE
#define CELL 4
#define COLS (SCREEN_WIDTH/CELL)
#define ROWS (SCREEN_HEIGHT/CELL)
#define MAX_SNAKE_LEN 256 // Maximum snake length is reduced for memory savings

// SRAM FIX: uint8_t and smaller array size saves over 1.5KB of SRAM
uint8_t sX[MAX_SNAKE_LEN], sY[MAX_SNAKE_LEN], sLen, dir;

void spawnFood(int &fx,int &fy){
  bool ok=false;
  while(!ok){
    fx=random(0,COLS);
    fy=random(0,ROWS);
    ok=true;
    for(int i=0;i<sLen;i++) if(sX[i]==fx && sY[i]==fy) ok=false;
  }
}

void snakeGame(){
  sLen=4;
  int startX=COLS/2, startY=ROWS/2;
  for(int i=0;i<sLen;i++){sX[i]=startX-i;sY[i]=startY;}
  dir=1; // right
  int foodX, foodY;
  spawnFood(foodX, foodY);
  unsigned long lastStep=millis();
  int speed=180;

  while(true){
    if(wasPressed(PIN_B)) return;

    // controls L=left, R=right, A=up, B=down
    if(wasPressed(PIN_L) && dir!=1) dir=3;
    if(wasPressed(PIN_R) && dir!=3) dir=1;
    if(wasPressed(PIN_A) && dir!=2) dir=0;
    if(wasPressed(PIN_B) && dir!=0) dir=2;

    if(millis()-lastStep>speed){
      lastStep=millis();
      for(int i=sLen-1;i>0;i--){sX[i]=sX[i-1]; sY[i]=sY[i-1];}
      if(dir==0) sY[0]--; else if(dir==1) sX[0]++; else if(dir==2) sY[0]++; else if(dir==3) sX[0]--;

      // wrap
      if(sX[0]<0) sX[0]=COLS-1; if(sX[0]>=COLS) sX[0]=0;
      if(sY[0]<0) sY[0]=ROWS-1; if(sY[0]>=ROWS) sY[0]=0;

      // self collision
      for(int i=1;i<sLen;i++) if(sX[i]==sX[0] && sY[i]==sY[0]){
        display.clearDisplay();
        display.setTextSize(1);
        display.setCursor(8,25);
        display.println(F("Snake Over - B to return"));
        display.display();
        while(!wasPressed(PIN_B)) delay(40);
        return;
      }

      // eat food
      if(sX[0]==foodX && sY[0]==foodY){
        sLen++;
        if(sLen>MAX_SNAKE_LEN-1) sLen=MAX_SNAKE_LEN-1; 
        sX[sLen-1]=sX[sLen-2]; sY[sLen-1]=sY[sLen-2];
        spawnFood(foodX,foodY);
        if(speed>60) speed-=6;
      }
    }

    // draw
    display.clearDisplay();
    display.fillRect(foodX*CELL,foodY*CELL,CELL,CELL,WHITE);
    for(int i=0;i<sLen;i++) display.fillRect(sX[i]*CELL,sY[i]*CELL,CELL,CELL,WHITE);
    display.display();
    delay(10);
  }
}

// =============================
// Main Loop
void loop(){
  int choice=showMenuLoop();
  if(choice==0) pongGame();
  else if(choice==1) snakeGame();
  delay(150);
}