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

#define SCREEN_WIDTH 128
#define SCREEN_HEIGHT 64
#define GRID 4

Adafruit_SSD1306 display(SCREEN_WIDTH, SCREEN_HEIGHT, &Wire, -1);

// ---------- BUTTONS ----------
#define BTN_UP     32
#define BTN_DOWN   33
#define BTN_LEFT   25
#define BTN_RIGHT  26
#define BTN_OK     27
#define BTN_BACK   14
#define PIEZO      13   // Piezo buzzer output

#define PRESSED(pin) (digitalRead(pin) == LOW)

enum AppState { APP_BOOT, APP_HOME, APP_SNAKE, APP_CLOCK, APP_CALC, APP_TETRIS, APP_NOTEPAD, APP_DICE, APP_COUNTER, APP_SETTINGS, APP_ABOUT, APP_MUSIC };
AppState currentApp = APP_BOOT;

// ---------- HOME MENU ----------
struct App {
  const char* name;
  const uint8_t icon[8];
  AppState state;
};

App apps[] = {
  {"Snake",      {0x3C,0x42,0x81,0x81,0x81,0x42,0x3C,0x00}, APP_SNAKE},
  {"Clock",      {0x3C,0x42,0x81,0x81,0x81,0x42,0x3C,0x00}, APP_CLOCK},
  {"Calculator", {0x3C,0x24,0x24,0x3C,0x24,0x24,0x3C,0x00}, APP_CALC},
  {"Tetris",     {0x3C,0x24,0x7E,0x7E,0x24,0x3C,0x00,0x00}, APP_TETRIS},
  {"Notepad",    {0x3C,0x24,0x3C,0x24,0x3C,0x24,0x3C,0x00}, APP_NOTEPAD},
  {"Dice",       {0x3C,0x42,0x42,0x3C,0x42,0x42,0x3C,0x00}, APP_DICE},
  {"Counter",    {0x3C,0x42,0x3C,0x42,0x3C,0x42,0x3C,0x00}, APP_COUNTER},
  {"Settings",   {0x3C,0x24,0x3C,0x42,0x3C,0x24,0x3C,0x00}, APP_SETTINGS},
  {"About",      {0x3C,0x42,0xA5,0x81,0xA5,0x42,0x3C,0x00}, APP_ABOUT},
  {"Music",      {0x3C,0x42,0x81,0xBD,0x81,0x42,0x3C,0x00}, APP_MUSIC}
};
const int appCount = sizeof(apps)/sizeof(App);
int selectedApp = 0;
int menuOffset = 0;

// ---------- SNAKE ----------
#define MAX_LEN 120
int snakeX[MAX_LEN], snakeY[MAX_LEN];
int length, dirX, dirY;
int foodX, foodY;
unsigned long lastMove;
int speed;
bool snakeOver;

// ---------- CALCULATOR ----------
long calcValue1 = 0;
long calcValue2 = 0;
char operation = 0;
bool enteringSecond = false;
unsigned long lastCalcPress = 0;

// ---------- TETRIS ----------
#define T_WIDTH 10
#define T_HEIGHT 16
int tetrisBoard[T_WIDTH][T_HEIGHT];
int tetX, tetY;
int tetShape[4][4];
int tetrisScore;
bool tetrisGameOver;
const int shapes[7][4][4] = {
  { {1,1,1,1},{0,0,0,0},{0,0,0,0},{0,0,0,0} }, // I
  { {1,1},{1,1},{0,0},{0,0} }, // O
  { {0,1,0},{1,1,1},{0,0,0},{0,0,0} }, // T
  { {1,1,0},{0,1,1},{0,0,0},{0,0,0} }, // S
  { {0,1,1},{1,1,0},{0,0,0},{0,0,0} }, // Z
  { {1,0,0},{1,1,1},{0,0,0},{0,0,0} }, // L
  { {0,0,1},{1,1,1},{0,0,0},{0,0,0} }  // J
};

// ---------- NOTEPAD ----------
char notepadText[32];
int cursorPos = 0;

// ---------- DICE ----------
int diceValue = 1;

// ---------- COUNTER ----------
int counterValue = 0;

// ---------- MUSIC ----------
const int musicNotes[] = { 659, 659, 0, 659, 0, 523, 659, 0, 784 }; // Bummer/Mobilnik melody simplified
const int musicDurations[] = {200,200,100,200,100,200,200,100,400};
const int musicLength = sizeof(musicNotes)/sizeof(int);
int musicIndex = 0;
unsigned long lastNote = 0;

// ---------- UTIL ----------
void waitBtn() { delay(150); }
void spawnFood() { foodX = random(0, SCREEN_WIDTH / GRID); foodY = random(0, SCREEN_HEIGHT / GRID); }
void resetSnake() {
  length = 3; dirX = 1; dirY = 0; speed = 200; snakeOver = false; lastMove = millis();
  for (int i = 0; i < length; i++) { snakeX[i] = 10 - i; snakeY[i] = 10; }
  spawnFood();
}
void playTone(int freq,int dur){ tone(PIEZO,freq,dur); }

// ---------- BOOT ----------
void bootScreen() {
  display.clearDisplay();
  display.setTextColor(SSD1306_WHITE);
  display.fillCircle(40, 32, 10, SSD1306_WHITE);
  display.fillTriangle(50,32,65,22,65,42,SSD1306_WHITE);
  display.fillCircle(36,29,2,SSD1306_BLACK);
  display.setTextSize(1);
  display.setCursor(75, SCREEN_HEIGHT/2 - 4);
  display.print("HELLO");
  display.display();
  delay(1500);
  currentApp = APP_HOME;
}

// ---------- DRAW ICON ----------
void drawIcon(uint8_t x, uint8_t y, const uint8_t icon[8]) {
  for (int i=0;i<8;i++){
    for(int j=0;j<8;j++){
      if(icon[i] & (1<<(7-j))) display.drawPixel(x+j,y+i,SSD1306_WHITE);
    }
  }
}

// ---------- HOME LOOP ----------
void homeLoop() {
  if (PRESSED(BTN_UP))   { selectedApp = (selectedApp - 1 + appCount) % appCount; waitBtn(); playTone(1000,50);}
  if (PRESSED(BTN_DOWN)) { selectedApp = (selectedApp + 1) % appCount; waitBtn(); playTone(1000,50);}

  if (selectedApp < menuOffset) menuOffset = selectedApp;
  if (selectedApp >= menuOffset + 4) menuOffset = selectedApp - 3;

  if (PRESSED(BTN_OK)) {
    currentApp = apps[selectedApp].state;
    waitBtn();
    playTone(1200,80);
    if(currentApp == APP_SNAKE) resetSnake();
    if(currentApp == APP_TETRIS){
      memset(tetrisBoard,0,sizeof(tetrisBoard));
      tetrisScore=0;
      tetrisGameOver=false;
      tetX=3; tetY=0;
      int s=random(0,7);
      memcpy(tetShape,shapes[s],sizeof(tetShape));
    }
  }
  if(PRESSED(BTN_BACK)){ currentApp = APP_HOME; waitBtn(); }

  display.clearDisplay();
  display.setTextSize(1);
  display.setTextColor(SSD1306_WHITE);
  display.setCursor(0,0);
  display.println("Fishy OS");
  display.println("----------------");
  for(int i=0;i<4 && (i+menuOffset)<appCount;i++){
    int idx = i+menuOffset;
    display.print(idx==selectedApp ? "> " : "  ");
    display.println(apps[idx].name);
    drawIcon(100,i*10+10,apps[idx].icon);
  }
  display.display();
}

// ---------- SNAKE LOOP ----------
void snakeLoop() {
  if (PRESSED(BTN_LEFT)  && dirX != 1)  { dirX = -1; dirY = 0; playTone(100,50);}
  if (PRESSED(BTN_RIGHT) && dirX != -1) { dirX = 1;  dirY = 0; playTone(100,50);}
  if (PRESSED(BTN_UP)    && dirY != 1)  { dirX = 0;  dirY = -1; playTone(100,50);}
  if (PRESSED(BTN_DOWN)  && dirY != -1) { dirX = 0;  dirY = 1; playTone(100,50);}
  if (PRESSED(BTN_OK)) { currentApp = APP_HOME; waitBtn(); playTone(800,80);}

  if (millis() - lastMove > speed && !snakeOver) {
    lastMove = millis();
    for(int i=length;i>0;i--){ snakeX[i]=snakeX[i-1]; snakeY[i]=snakeY[i-1]; }
    snakeX[0]+=dirX; snakeY[0]+=dirY;
    if(snakeX[0]<0) snakeX[0]=SCREEN_WIDTH/GRID-1;
    if(snakeX[0]>=SCREEN_WIDTH/GRID) snakeX[0]=0;
    if(snakeY[0]<0) snakeY[0]=SCREEN_HEIGHT/GRID-1;
    if(snakeY[0]>=SCREEN_HEIGHT/GRID) snakeY[0]=0;
    for(int i=1;i<length;i++) if(snakeX[0]==snakeX[i] && snakeY[0]==snakeY[i]) snakeOver=true;
    if(snakeX[0]==foodX && snakeY[0]==foodY){ length++; speed=max(60,speed-10); spawnFood(); playTone(1500,80);}
  }

  display.clearDisplay();
  display.setTextColor(SSD1306_WHITE);
  if(snakeOver){
    display.setCursor(30,30); display.print("GAME OVER");
  }else{
    display.fillRect(foodX*GRID,foodY*GRID,GRID,GRID,SSD1306_WHITE);
    for(int i=0;i<length;i++) display.fillRect(snakeX[i]*GRID,snakeY[i]*GRID,GRID,GRID,SSD1306_WHITE);
  }
  display.display();
}

// ---------- CLOCK LOOP ----------
void clockLoop() {
  if (PRESSED(BTN_OK) || PRESSED(BTN_BACK)) { currentApp = APP_HOME; waitBtn(); playTone(800,80);}
  display.clearDisplay();
  display.setTextSize(2);
  display.setCursor(10,25);
  unsigned long t = millis()/1000;
  int h = (t/3600)%24;
  int m = (t/60)%60;
  int s = t%60;
  char buf[9];
  sprintf(buf,"%02d:%02d:%02d",h,m,s);
  display.print(buf);
  display.display();
}

// ---------- CALCULATOR LOOP ----------
void calcLoop() {
  unsigned long now = millis();
  if (PRESSED(BTN_UP) && now - lastCalcPress > 150) { if(!enteringSecond) calcValue1++; else calcValue2++; lastCalcPress=now; playTone(600,50);}
  if (PRESSED(BTN_DOWN) && now - lastCalcPress > 150) { if(!enteringSecond) calcValue1--; else calcValue2--; lastCalcPress=now; playTone(600,50);}
  if ((PRESSED(BTN_LEFT) || PRESSED(BTN_RIGHT)) && now - lastCalcPress > 150) { if(!enteringSecond) operation = (operation=='+'?'-':'+'); lastCalcPress=now; playTone(700,50);}
  if (PRESSED(BTN_OK) && now - lastCalcPress > 150) {
    if(!enteringSecond) enteringSecond=true;
    else {
      long res=0;
      if(operation=='+') res=calcValue1+calcValue2;
      else if(operation=='-') res=calcValue1-calcValue2;
      else if(operation=='*') res=calcValue1*calcValue2;
      else if(operation=='/') res=calcValue2!=0 ? calcValue1/calcValue2 : 0;
      calcValue1=res; calcValue2=0; operation=0; enteringSecond=false;
      playTone(1000,80);
    }
    lastCalcPress=now;
  }
  if (PRESSED(BTN_BACK)) { currentApp=APP_HOME; enteringSecond=false; waitBtn(); playTone(800,80);}

  display.clearDisplay();
  display.setTextSize(1);
  display.setCursor(0,0);
  display.println("Calculator:");
  display.setCursor(0,12);
  display.print(calcValue1);
  if(operation) display.print(operation);
  if(enteringSecond) display.print(calcValue2);
  display.display();
}

// ---------- NOTEPAD LOOP ----------
void notepadLoop() {
  unsigned long now = millis();
  if (PRESSED(BTN_BACK)) { currentApp=APP_HOME; waitBtn(); cursorPos=0; playTone(800,80);}
  if (PRESSED(BTN_UP) && now - lastCalcPress>150) { if(cursorPos<31){ notepadText[cursorPos]++; if(notepadText[cursorPos]>126) notepadText[cursorPos]=32; lastCalcPress=now; playTone(600,50);} }
  if (PRESSED(BTN_DOWN) && now - lastCalcPress>150) { if(cursorPos<31){ notepadText[cursorPos]--; if(notepadText[cursorPos]<32) notepadText[cursorPos]=126; lastCalcPress=now; playTone(600,50);} }
  if (PRESSED(BTN_RIGHT) && now - lastCalcPress>150) { if(cursorPos<31) cursorPos++; lastCalcPress=now; playTone(700,50);}
  if (PRESSED(BTN_LEFT) && now - lastCalcPress>150) { if(cursorPos>0) cursorPos--; lastCalcPress=now; playTone(700,50);}

  display.clearDisplay();
  display.setTextSize(1);
  display.setCursor(0,0);
  display.print("Notepad:");
  display.setCursor(0,12);
  display.print(notepadText);
  display.display();
}

// ---------- MUSIC LOOP ----------
void musicLoop() {
  if(PRESSED(BTN_BACK)){ currentApp=APP_HOME; waitBtn(); playTone(800,80); musicIndex=0;}
  display.clearDisplay();
  display.setTextSize(1);
  display.setCursor(0,0);
  display.println("Music: Bummer/Mobilnik");

  unsigned long now = millis();
  if(now - lastNote > musicDurations[musicIndex]){
    lastNote = now;
    if(musicNotes[musicIndex]>0) playTone(musicNotes[musicIndex],musicDurations[musicIndex]);
    musicIndex = (musicIndex+1)%musicLength;
  }
  display.display();
}

// ---------- DICE LOOP ----------
void diceLoop() {
  if (PRESSED(BTN_OK) || PRESSED(BTN_BACK)) { currentApp=APP_HOME; waitBtn(); playTone(800,80);}
  if (PRESSED(BTN_UP)) diceValue = random(1,7);
  display.clearDisplay();
  display.setTextSize(2);
  display.setCursor(40,25);
  display.print(diceValue);
  display.display();
}

// ---------- COUNTER LOOP ----------
void counterLoop() {
  if (PRESSED(BTN_UP)) counterValue++;
  if (PRESSED(BTN_DOWN)) counterValue--;
  if (PRESSED(BTN_BACK)) counterValue=0;
  if (PRESSED(BTN_OK)) { currentApp=APP_HOME; waitBtn(); playTone(800,80);}

  display.clearDisplay();
  display.setTextSize(2);
  display.setCursor(40,25);
  display.print(counterValue);
  display.display();
}

// ---------- SETTINGS LOOP ----------
void settingsLoop() {
  if (PRESSED(BTN_BACK) || PRESSED(BTN_OK)) { currentApp=APP_HOME; waitBtn(); playTone(800,80);}
  display.clearDisplay();
  display.setTextSize(1);
  display.setCursor(0,0);
  display.println("Settings:");
  display.println("Speed / buttons / sound");
  display.display();
}

// ---------- ABOUT LOOP ----------
void aboutLoop() {
  if(PRESSED(BTN_OK) || PRESSED(BTN_BACK)){ currentApp=APP_HOME; waitBtn(); playTone(800,80);}
  display.clearDisplay();
  display.setTextSize(1);
  display.setCursor(10,10);
  display.println("Fishy OS v1");
  display.println("Mini OS ESP32");
  display.println("by Aroniukas");
  display.display();
}

// ---------- TETRIS HELPERS ----------
void drawTetromino(int x,int y,int shape[4][4],bool color){
  for(int i=0;i<4;i++){
    for(int j=0;j<4;j++){
      if(shape[i][j]){
        display.fillRect((x+j)*GRID,(y+i)*GRID,GRID,GRID,color?SSD1306_WHITE:SSD1306_BLACK);
      }
    }
  }
}
bool checkCollision(int x,int y,int shape[4][4]){
  for(int i=0;i<4;i++){
    for(int j=0;j<4;j++){
      if(shape[i][j]){
        int nx=x+j;
        int ny=y+i;
        if(nx<0 || nx>=T_WIDTH || ny>=T_HEIGHT) return true;
        if(ny>=0 && tetrisBoard[nx][ny]) return true;
      }
    }
  }
  return false;
}
void lockTetromino(){
  for(int i=0;i<4;i++){
    for(int j=0;j<4;j++){
      if(tetShape[i][j]){
        int nx=tetX+j;
        int ny=tetY+i;
        if(ny>=0) tetrisBoard[nx][ny]=1;
      }
    }
  }
  for(int y=T_HEIGHT-1;y>=0;y--){
    bool full=true;
    for(int x=0;x<T_WIDTH;x++) if(!tetrisBoard[x][y]) full=false;
    if(full){
      for(int yy=y;yy>0;yy--) for(int x=0;x<T_WIDTH;x++) tetrisBoard[x][yy]=tetrisBoard[x][yy-1];
      for(int x=0;x<T_WIDTH;x++) tetrisBoard[x][0]=0;
      tetrisScore+=100;
      playTone(1500,80);
      y++;
    }
  }
}
void rotateTetromino(){
  int tmp[4][4];
  for(int i=0;i<4;i++) for(int j=0;j<4;j++) tmp[i][j]=tetShape[3-j][i];
  if(!checkCollision(tetX,tetY,tmp)) memcpy(tetShape,tmp,sizeof(tetShape));
}

// ---------- TETRIS LOOP ----------
void tetrisLoop(){
  static unsigned long lastFall=0;
  if(tetrisGameOver){
    display.clearDisplay();
    display.setTextSize(1);
    display.setCursor(20,25);
    display.println("GAME OVER");
    display.setCursor(20,35);
    display.print("Score: "); display.print(tetrisScore);
    display.display();
    if(PRESSED(BTN_OK) || PRESSED(BTN_BACK)){ currentApp=APP_HOME; waitBtn(); playTone(800,80);}
    return;
  }

  if (PRESSED(BTN_LEFT)){ if(!checkCollision(tetX-1,tetY,tetShape)) {tetX--; playTone(500,50);} waitBtn();}
  if (PRESSED(BTN_RIGHT)){ if(!checkCollision(tetX+1,tetY,tetShape)) {tetX++; playTone(500,50);} waitBtn();}
  if (PRESSED(BTN_UP)){ rotateTetromino(); playTone(700,50); waitBtn();}
  if (PRESSED(BTN_DOWN)){ if(!checkCollision(tetX,tetY+1,tetShape)) {tetY++; playTone(400,50);} waitBtn();}
  if (PRESSED(BTN_BACK)){ currentApp=APP_HOME; waitBtn(); playTone(800,80);}

  if(millis()-lastFall>500){
    lastFall=millis();
    if(!checkCollision(tetX,tetY+1,tetShape)){ tetY++; }
    else{
      lockTetromino();
      int s=random(0,7);
      memcpy(tetShape,shapes[s],sizeof(tetShape));
      tetX=3; tetY=0;
      if(checkCollision(tetX,tetY,tetShape)) tetrisGameOver=true;
    }
  }

  display.clearDisplay();
  display.setTextSize(1);
  for(int x=0;x<T_WIDTH;x++){
    for(int y=0;y<T_HEIGHT;y++){
      if(tetrisBoard[x][y]) display.fillRect(x*GRID,y*GRID,GRID,GRID,SSD1306_WHITE);
    }
  }
  drawTetromino(tetX,tetY,tetShape,true);
  display.setCursor(0,0);
  display.print("Score:"); display.print(tetrisScore);
  display.display();
}

// ---------- SETUP ----------
void setup() {
  Wire.begin();
  display.begin(SSD1306_SWITCHCAPVCC, 0x3C);
  display.clearDisplay();
  pinMode(BTN_UP,INPUT_PULLUP);
  pinMode(BTN_DOWN,INPUT_PULLUP);
  pinMode(BTN_LEFT,INPUT_PULLUP);
  pinMode(BTN_RIGHT,INPUT_PULLUP);
  pinMode(BTN_OK,INPUT_PULLUP);
  pinMode(BTN_BACK,INPUT_PULLUP);
  pinMode(PIEZO,OUTPUT);
  randomSeed(analogRead(0));
  bootScreen();
}

// ---------- LOOP ----------
void loop() {
  switch(currentApp){
    case APP_HOME: homeLoop(); break;
    case APP_SNAKE: snakeLoop(); break;
    case APP_CLOCK: clockLoop(); break;
    case APP_CALC: calcLoop(); break;
    case APP_TETRIS: tetrisLoop(); break;
    case APP_NOTEPAD: notepadLoop(); break;
    case APP_DICE: diceLoop(); break;
    case APP_COUNTER: counterLoop(); break;
    case APP_SETTINGS: settingsLoop(); break;
    case APP_ABOUT: aboutLoop(); break;
    case APP_MUSIC: musicLoop(); break;
    default: homeLoop(); break;
  }
}