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

/********* OLED *********/
#define SCREEN_WIDTH    128
#define SCREEN_HEIGHT   64
#define SCREEN_ADDRESS  0x3C
Adafruit_SSD1306 display(SCREEN_WIDTH, SCREEN_HEIGHT, &Wire, -1, 1700000UL, 1700000UL);

/********* Touch pins (official pinout) *********/
// CH0..CH5 → P03,P04,P02,P27,P01,P26  (3,4,2,27,1,26)
#define N_TOUCH 6
int touch_pins[N_TOUCH] = {3, 4, 2, 27, 1, 26};

/********* Touch timing *********/
const uint32_t SAFE_WINDOW_MS   = 2000;
volatile uint16_t T_MAX_US      = 500;     // enough for your board
const uint16_t DISCHARGE_US     = 25;

/********* Game timing *********/
const uint32_t FRAME_MS         = 50;      // 20 FPS

/********* Touch state *********/
uint16_t touch_us[N_TOUCH]      = {0};
bool     pressed[N_TOUCH]       = {0};     // nonzero => pressed

/********* Utils *********/
static inline void usb_friendly_yield() { yield(); }

static inline uint16_t measure_touch_us(int pin) {
  pinMode(pin, OUTPUT);
  digitalWriteFast(pin, LOW);
  delayMicroseconds(DISCHARGE_US);
  pinMode(pin, INPUT_PULLUP);
  unsigned long t0 = micros();
  while (!digitalReadFast(pin)) {
    if ((uint16_t)(micros() - t0) >= T_MAX_US) break;
    usb_friendly_yield();
  }
  return (uint16_t)(micros() - t0);
}

/********* Input mapping *********/
// CH2,3,4,5 = down,left,right,up ; CH0,1 = func right/left
#define CH_DOWN   2
#define CH_LEFT   3
#define CH_RIGHT  4
#define CH_UP     5
#define CH_FUNC_R 0   // pause/resume
#define CH_FUNC_L 1   // restart

/********* Game state *********/
enum GameState { GS_MENU, GS_PLAY, GS_PAUSE, GS_GAMEOVER };
GameState gstate = GS_MENU;

struct Dino {
  int x = 12;
  int y = 48;          // feet on ground (groundY = 56; dino h ~ 10)
  float vy = 0;
  bool crouch = false;
} dino;

const int groundY = 56;    // baseline y of ground
const int dinoW = 10;
const int dinoH = 10;
const int dinoH_crouch = 7;

struct Rock {
  int x = 140;
  int y = groundY - 6;    // small rock
  int w = 6;
  int h = 6;
  bool active = false;
};

const int MAX_ROCKS = 3;
Rock rocks[MAX_ROCKS];

int scrollSpeed = 2;       // px per frame
uint32_t frameTimer = 0;
uint32_t spawnTimer = 0;
uint32_t spawnInterval = 1200;  // ms; randomized per spawn
uint32_t score = 0;
uint32_t best  = 0;

bool wasPressed[N_TOUCH] = {0}; // edge detection

/********* helpers *********/
bool justPressed(int ch) { bool jp = (pressed[ch] && !wasPressed[ch]); wasPressed[ch] = pressed[ch]; return jp; }
void updateWasPressed()  { for (int i=0;i<N_TOUCH;i++) wasPressed[i] = pressed[i]; }

void reset_game(bool keep_best) {
  dino.x = 12;
  dino.y = groundY - dinoH;
  dino.vy = 0;
  dino.crouch = false;

  for (int i=0;i<MAX_ROCKS;i++) { rocks[i].active = false; rocks[i].x = 140; rocks[i].y = groundY - 6; rocks[i].w = 6; rocks[i].h = 6; }
  scrollSpeed = 2;
  if (!keep_best) best = 0;
  score = 0;
  spawnTimer = 0;
  spawnInterval = 900; // first spawn quicker
}

void spawn_rock() {
  // find free slot
  for (int i=0;i<MAX_ROCKS;i++) if (!rocks[i].active) {
    rocks[i].active = true;
    rocks[i].x = SCREEN_WIDTH + random(0, 16);       // spawn just off-screen
    // randomize type/height/size a bit
    int type = random(0,3);
    if (type == 0) { rocks[i].w = 6;  rocks[i].h = 6;  rocks[i].y = groundY - rocks[i].h; }
    if (type == 1) { rocks[i].w = 8;  rocks[i].h = 10; rocks[i].y = groundY - rocks[i].h; }
    if (type == 2) { rocks[i].w = 4;  rocks[i].h = 4;  rocks[i].y = groundY - rocks[i].h; }
    break;
  }
  // next interval: faster with score
  uint32_t base = 800;
  uint32_t var  = 600;
  if (score > 100) base = 650;
  if (score > 200) base = 550;
  spawnInterval = base + random(0, var);
}

/********* collision *********/
bool collideRect(int x1,int y1,int w1,int h1, int x2,int y2,int w2,int h2) {
  return !(x2 > x1+w1-1 || x2+w2-1 < x1 || y2 > y1+h1-1 || y2+h2-1 < y1);
}

/********* draw *********/
void draw_dino() {
  int h = dino.crouch ? dinoH_crouch : dinoH;
  // body
  display.fillRect(dino.x, dino.y, dinoW, h, SSD1306_WHITE);
  // eye
  display.fillRect(dino.x + dinoW - 3, dino.y + 2, 2, 2, SSD1306_BLACK);
}

void draw_ground() {
  display.drawLine(0, groundY+1, SCREEN_WIDTH-1, groundY+1, SSD1306_WHITE);
}

void draw_rocks() {
  for (int i=0;i<MAX_ROCKS;i++) if (rocks[i].active) {
    display.fillRect(rocks[i].x, rocks[i].y, rocks[i].w, rocks[i].h, SSD1306_WHITE);
  }
}

void draw_hud() {
  display.setTextSize(1);
  display.setTextColor(SSD1306_WHITE);
  display.setCursor(0,0);
  display.print("S:");
  display.print(score);
  display.setCursor(64,0);
  display.print("B:");
  display.print(best);
  // hints
  display.setCursor(0,56);
  if (gstate==GS_MENU)      display.print("CH0:Start  CH1:Reset");
  else if (gstate==GS_PAUSE)display.print("PAUSE CH0:Resume");
  else if (gstate==GS_PLAY) display.print("UP:Jump  L/R:Move  CH0:Pause");
  else if (gstate==GS_GAMEOVER) display.print("GameOver CH1:Restart");
}

/********* input *********/
void poll_touch() {
  for (int i=0;i<N_TOUCH;i++) {
    touch_us[i] = measure_touch_us(touch_pins[i]);
    pressed[i]  = (touch_us[i] != 0);  // nonzero => pressed
  }
}

/********* setup/loop *********/
unsigned long t_boot = 0;

void setup() {
  pinMode(LED_BUILTIN, OUTPUT);
  Serial.begin(115200);
  t_boot = millis();

  // OLED init (keep your working Wire settings)
  Wire.begin();
  Wire.setClock(400000);
  delay(30);
  if (!display.begin(SSD1306_SWITCHCAPVCC, SCREEN_ADDRESS)) {
    // still continue game logic even if OLED fails
  } else {
    display.clearDisplay();
    display.setTextSize(1);
    display.setTextColor(SSD1306_WHITE);
    display.setCursor(8, 26);
    display.println("Tiny Dino!");
    display.setCursor(8, 38);
    display.println("CH0 Start / CH1 Reset");
    display.display();
  }

  randomSeed(micros());
  reset_game(true);
  gstate = GS_MENU;
}

void do_menu() {
  // wait for CH0 to start; CH1 resets best
  if (pressed[CH_FUNC_R]) { gstate = GS_PLAY; }
  if (pressed[CH_FUNC_L]) { reset_game(false); }
}

void do_pause() {
  if (pressed[CH_FUNC_R]) gstate = GS_PLAY;
  if (pressed[CH_FUNC_L]) { reset_game(true); gstate = GS_MENU; }
}

void do_gameover() {
  if (pressed[CH_FUNC_L]) { reset_game(true); gstate = GS_MENU; }
}

void do_play() {
  // input
  bool jump  = pressed[CH_UP];
  bool left  = pressed[CH_LEFT];
  bool right = pressed[CH_RIGHT];
  bool down  = pressed[CH_DOWN];

  // pause
  if (pressed[CH_FUNC_R]) { gstate = GS_PAUSE; return; }

  // crouch reduces hitbox
  dino.crouch = down;

  // jump: only if "on ground"
  int curH = dino.crouch ? dinoH_crouch : dinoH;
  bool onGround = (dino.y >= groundY - curH);
  if (onGround) { dino.y = groundY - curH; dino.vy = 0; }
  if (jump && onGround) {
    dino.vy = -6.2;  // jump impulse
  }

  // horizontal nudge
  if (left)  dino.x -= 2;
  if (right) dino.x += 2;
  if (dino.x < 2) dino.x = 2;
  if (dino.x > 50) dino.x = 50; // keep near left half

  // gravity & vertical motion
  dino.vy += 0.55;          // gravity
  dino.y  += (int)dino.vy;
  curH = dino.crouch ? dinoH_crouch : dinoH;
  if (dino.y > groundY - curH) { dino.y = groundY - curH; dino.vy = 0; }

  // spawn rocks
  spawnTimer += FRAME_MS;
  if (spawnTimer >= spawnInterval) {
    spawnTimer = 0;
    spawn_rock();
  }

  // move rocks & scoring/collision
  for (int i=0;i<MAX_ROCKS;i++) if (rocks[i].active) {
    rocks[i].x -= scrollSpeed;
    if (rocks[i].x + rocks[i].w < 0) {
      rocks[i].active = false;
      score += 10;
      if (score > best) best = score;
      if (score % 50 == 0 && scrollSpeed < 4) scrollSpeed++; // ramp difficulty
    } else {
      // collision
      int dH = dino.crouch ? dinoH_crouch : dinoH;
      if (collideRect(dino.x, dino.y, dinoW, dH,
                      rocks[i].x, rocks[i].y, rocks[i].w, rocks[i].h)) {
        gstate = GS_GAMEOVER;
      }
    }
  }
}

void draw_scene() {
  display.clearDisplay();
  draw_ground();
  draw_rocks();
  draw_dino();
  draw_hud();

  if (gstate == GS_MENU) {
    display.setCursor(20, 24);
    display.println("Tiny Dino");
    display.setCursor(10, 36);
    display.println("CH0 Start  CH1 Reset");
  } else if (gstate == GS_PAUSE) {
    display.setCursor(44, 24);
    display.println("PAUSED");
  } else if (gstate == GS_GAMEOVER) {
    display.setCursor(28, 24);
    display.println("GAME OVER");
    display.setCursor(10, 36);
    display.println("CH1: Restart");
  }

  display.display();
}

void loop() {
  // Safe window: allow USB enumeration & opening monitor
  if (millis() - t_boot < SAFE_WINDOW_MS) {
    digitalWrite(LED_BUILTIN, (millis()/250)%2);
    delay(50);
    return;
  }

  // frame pacing
  static uint32_t last = 0;
  uint32_t now = millis();
  if (now - last < FRAME_MS) return;
  last = now;

  // poll input
  poll_touch();

  // state machine
  switch (gstate) {
    case GS_MENU:     do_menu();     break;
    case GS_PLAY:     do_play();     break;
    case GS_PAUSE:    do_pause();    break;
    case GS_GAMEOVER: do_gameover(); break;
  }

  // render
  draw_scene();
}