#include <Adafruit_GFX.h>
#include <Adafruit_SSD1306.h>
#include <string.h>
#include <Wire.h>
#include <map>
#include <utility>

#define WIDTH 128 // OLED display width, in pixels
#define HEIGHT 64 // OLED display height, in pixels
#define WHITE SSD1306_WHITE

// --- touch input (reuse your working pins/threshold) ---
#define N_TOUCH 6
#define THRESHOLD 30
int step_delay = 200;
int step_num = 0;

int  touch_pins[N_TOUCH]      = {3, 4, 2, 27, 1, 26};
int  touch_values[N_TOUCH]    = {0};
bool pin_touched_now[N_TOUCH] = {false};
bool pin_touched_past[N_TOUCH]= {false};

// map a few “buttons” to indices you’ll actually touch
#define BTN_SELECT 0   // press to pause/resume
#define BTN_UP     5   // (we'll use later)
#define BTN_DOWN   2   // (we'll use later)

enum UiMode { MODE_GAME, MODE_MENU, MODE_PROMPT };
enum PromptKind { PROMPT_NONE, PROMPT_SPEED, PROMPT_FILL };

// --- UI state ---
UiMode     uiMode      = MODE_GAME;
PromptKind promptKind  = PROMPT_NONE;
bool       paused      = false;


// --- timing/speed state ---
int  fps = 5;   // keep your current speed
float fill = .25;

// --- serial prompt buffers ---
String inputStr   = "";
bool   inputDone  = false;

// (optional) forward declarations so the compiler knows these exist even if defined later
void startPrompt(PromptKind kind);
void pollPromptSerial();
void applyPrompt();
void drawPrompt();




// Reset pin not used with the XIAO RP2040, set to -1
Adafruit_SSD1306 display(WIDTH, HEIGHT, &Wire, -1, 1700000UL, 1700000UL);

constexpr int CELL_SIZE = 2;
constexpr int GRID_W = WIDTH/CELL_SIZE;
constexpr int GRID_H = HEIGHT/CELL_SIZE;

uint8_t grid[GRID_W * GRID_H];
uint8_t nextGrid[GRID_W * GRID_H];

// index helper function
inline int idx(int x, int y) {return y*GRID_W + x; }

// set/get helpers
inline void setCell(int x, int y, uint8_t alive) {
  if ((unsigned)x < (unsigned)GRID_W && (unsigned)y < (unsigned)GRID_H)
    grid[idx(x, y)] = alive ? 1 : 0;
}

inline uint8_t getCell(int x, int y) {
  if ((unsigned)x < (unsigned)GRID_W && (unsigned)y < (unsigned)GRID_H) {
    return grid[idx(x, y)];
  }
  return 0;
}

// draw a cell

inline void drawCellBlock(int cx, int cy, uint8_t alive) {
  int px = cx*CELL_SIZE;
  int py = cy*CELL_SIZE;
  if (alive) {
    for (int dy = 0; dy < CELL_SIZE; ++dy) 
      for (int dx = 0; dx < CELL_SIZE; ++dx)
        display.drawPixel(px + dx, py+dy, WHITE);
  } 
}

void drawGrid() {
  display.clearDisplay();

  for (int y = 0; y < GRID_H; ++y) {
    for (int x = 0; x < GRID_W; ++x) {
      if (grid[idx(x, y)]) drawCellBlock(x, y, 1);
    }
  }

  // --- minimal status banner (top-left) ---
  // display.setTextSize(1);
  // display.setTextColor(WHITE);
  // display.setCursor(0, 0);
  // display.print(paused ? "PAUSED " : "RUN    ");
  // display.print(fps);
  // display.print("FPS");
  int boxX = 90;
  int boxY = 54;
  int boxW = 38;   // width of the box in pixels
  int boxH = 10;   // height of the box in pixels

  // draw a filled black rectangle to "erase" that region
  display.fillRect(boxX, boxY, boxW, boxH, SSD1306_BLACK);

  // draw a white outline (optional)
  // display.drawRect(boxX, boxY, boxW, boxH, WHITE);

  // print frame number inside
  display.setTextSize(1);
  display.setTextColor(WHITE, SSD1306_BLACK);  // white text on black background
  display.setCursor(boxX + 2, boxY + 1);
  display.print("F:");
  display.print(step_num);

  display.display();
}

void seedTest() {
  for (int y = 0; y < GRID_H; ++y) {
    for (int x = 0; x < GRID_W; ++x) {
      setCell(x, y, ((x^y)&1));
    }
  }

  for (int y = 2; y < 6; ++y)
    for (int x = 2; x < 6; ++x)
      setCell(x, y, 1);
}

int numNeighbors(int x, int y) {
  int sum = 0;
  for (int cur_row = y-1; cur_row <= y + 1; ++cur_row){
    for (int cur_col = x-1; cur_col <= x + 1; ++cur_col){
      if (cur_row == y && cur_col == x) continue;
      int nx = (cur_col + GRID_W)  % GRID_W;       // wrap X
      int ny = (cur_row + GRID_H)  % GRID_H;
      sum += grid[idx(nx, ny)];
    }
  }
  return sum;
}

void setNextCell(int x, int y) {
  int neighbors = numNeighbors(x, y);
  uint8_t cur_val = getCell(x, y);
  if (cur_val == 1) {
    if (neighbors >= 2 && neighbors <= 3) {
      nextGrid[idx(x, y)] = 1;
    }
    else {
      nextGrid[idx(x, y)] = 0;
    }
  } else {
    if (neighbors == 3){
      nextGrid[idx(x, y)] = 1;
    }
    else {
      nextGrid[idx(x, y)] = 0;
    }
  }
}

void step() {
  for (int y = 0; y < GRID_H; ++y) {
    for (int x = 0; x < GRID_W; ++x) {
      setNextCell(x, y);
    }
  }
  memcpy(grid, nextGrid, sizeof(grid));
}


void seedBlinker(int cx, int cy) {
  setCell(cx-1, cy, 1);
  setCell(cx,   cy, 1);
  setCell(cx+1, cy, 1);
}

void randomStartState(float percentFill) {
  int num_pixels = round((GRID_W * GRID_H) * percentFill);
  while (num_pixels > 0) {
    int x = random(GRID_W);
    int y = random(GRID_H);
    setCell(x, y, 1);
    num_pixels -= 1;
  }
}

void update_touch() {
  const int t_max = 200;
  for (int i = 0; i < N_TOUCH; i++) {
    int p = touch_pins[i];

    // discharge
    pinMode(p, OUTPUT);
    digitalWriteFast(p, LOW);
    delayMicroseconds(25);

    noInterrupts();
    pinMode(p, INPUT_PULLUP);

    int t = 0;
    while (!digitalReadFast(p) && t < t_max) t++;
    interrupts();

    touch_values[i]    = t;
    pin_touched_past[i]= pin_touched_now[i];
    pin_touched_now[i] = (t > THRESHOLD);
  }
}

// “edge” (new press this frame)
bool buttonPressed(int btnIndex) {
  return pin_touched_now[btnIndex] && !pin_touched_past[btnIndex];
}


void startPrompt(PromptKind kind) {
  promptKind = kind;
  inputStr = "";
  inputDone = false;
  uiMode = MODE_PROMPT;

  // print a prompt to the terminal
  if (promptKind == PROMPT_SPEED) {
    Serial.println("Enter new frame rate (Frames Per Second), then press Enter:");
  } else if (promptKind == PROMPT_FILL) {
    Serial.println("Enter fill percentage (e.g., 0.25 is 25%), then press Enter:");
  }
}

void pollPromptSerial() {
  while (Serial.available() > 0) {
    char c = Serial.read();
    if (c == '\r') continue;        // ignore CR (Windows)
    if (c == '\n') { inputDone = true; break; }
    inputStr += c;
  }
}

void applyPrompt() {
  if (promptKind == PROMPT_SPEED) {
    long v = inputStr.toInt();         // parse integer
    if (v <= 1)   v = 1;              // clamp sensible bounds
    if (v > 30) v = 30;
    fps = (int)v;
    Serial.print("Frame rate set to ");
    Serial.print(fps);
    Serial.println(" FPS");

  } else if (promptKind == PROMPT_FILL) {
  float f = inputStr.toFloat();
  if (f < 0.0f) f = 0.0f;
  if (f > 1.0f) f = 1.0f;
  fill = f;                         // <-- keep global in sync
  memset(grid, 0, sizeof(grid));
  randomStartState(fill);
  Serial.print("Random fill set to ");
  Serial.println(fill, 3);
}


  // return to menu after applying
  promptKind = PROMPT_NONE;
  uiMode = MODE_MENU;
}

void drawPrompt() {
  display.clearDisplay();
  display.setTextSize(1);
  display.setTextColor(WHITE);
  display.setCursor(0, 0);

  if (promptKind == PROMPT_SPEED) {
    display.println("Set FPS:");
    display.println("(Use Serial Monitor)");
  } else if (promptKind == PROMPT_FILL) {
    display.println("Set Fill (0-1):");
    display.println("(Use Serial Monitor)");
  } else {
    display.println("Input:");
  }

  display.setCursor(0, 12);
  display.println(inputStr);   // echo what user typed
  display.display();
  delay(1000);
}






int menuIndex = 0;    // which item is highlighted
const int MENU_ITEMS = 4;
const char* menuLabels[MENU_ITEMS] = {
  "Resume",
  "Random Seed",
  "Set FPS", 
  "Set Fill"
};

void drawMenu() {
  display.clearDisplay();
  display.setTextSize(1);
  display.setTextColor(WHITE);
  display.setCursor(0,0);
  display.print("FPS: ");
  display.print(fps);
  display.print(" Fill:");
  display.print(fill);
  display.setCursor(0, 10);
  display.print("Frame: ");
  display.print(step_num);


  for (int i = 0; i < MENU_ITEMS; i++) {
    display.setCursor(0, (i + 2)*10);  // each line spaced ~10px
    if (i == menuIndex) {
      display.print("> ");
    } else {
      display.print("  ");
    }
    display.println(menuLabels[i]);
  }
  display.display();
}


void drawWelcome() {
  display.clearDisplay();
  display.setTextSize(1);
  display.setTextColor(WHITE);
  display.setCursor(0,0);
  display.print("Conway's Game of Life");
  display.setCursor(0,10);
  display.print("Let the games begin!");
  display.setCursor(0,50);
  display.print("To view menu press Q0");
  display.display();
}

void setup() {
  if (!display.begin(SSD1306_SWITCHCAPVCC, 0x3C)) {
    for(;;);
  }
  
  drawWelcome();
  delay(5000);

  randomStartState(fill);
  // drawGrid();
  Serial.begin(115200);
  while (!Serial) { ; }
  
}
void loop() {
  update_touch();

  if (uiMode == MODE_GAME) {
    if (buttonPressed(BTN_SELECT)) {
      uiMode = MODE_MENU;
    } else {
      if (!paused) { step(); step_num += 1; }
      drawGrid();
      delay(1000 / fps);
    }
  }

  else if (uiMode == MODE_MENU) {
    if (buttonPressed(BTN_UP))   menuIndex = (menuIndex - 1 + MENU_ITEMS) % MENU_ITEMS;
    if (buttonPressed(BTN_DOWN)) menuIndex = (menuIndex + 1) % MENU_ITEMS;

    if (buttonPressed(BTN_SELECT)) {
      if (menuIndex == 0) {                    // Resume
        uiMode = MODE_GAME;
      } else if (menuIndex == 1) {             // Random Seed
        memset(grid, 0, sizeof(grid));
        randomStartState(fill);                 // use current fill
        step_num = 0;
        uiMode = MODE_GAME;
      } else if (menuIndex == 2) {             // Set FPS
        startPrompt(PROMPT_SPEED);
      } else if (menuIndex == 3) {             // Set Fill
        startPrompt(PROMPT_FILL);
      }
    }

    drawMenu();
    delay(150);
  }

  else if (uiMode == MODE_PROMPT) {
    // keep running EVERY frame until user submits/cancels
    pollPromptSerial();

    // optional: SELECT = cancel
    if (buttonPressed(BTN_SELECT)) {
      promptKind = PROMPT_NONE;
      uiMode = MODE_MENU;
    }

    if (inputDone) {
      applyPrompt();   // parses + applies + returns to MENU
    }

    drawPrompt();
    delay(50);
  }
}