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

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

// XIAO RP2040 Onboard NeoPixel Setup
#define NUM_PIXELS 1
#define PIN_NEO_DATA 12
Adafruit_NeoPixel strip(NUM_PIXELS, PIN_NEO_DATA, NEO_GRB + NEO_KHZ800);


#define N_TOUCH 6
#define THRESHOLD 6     

// Mapping: [Left, Down, Up, Right, Q1 (Pin A), Q0 (Pin B)]
const uint8_t touch_pins[N_TOUCH] = { D1, D8, D0, D7, D9, D10 };
int touch_values[N_TOUCH] = { 0, 0, 0, 0, 0, 0 };

bool pin_touched_now[N_TOUCH]  = { false, false, false, false, false, false };
bool pin_touched_past[N_TOUCH] = { false, false, false, false, false, false };

void update_touch() {
  int t;
  const int t_max = 200;
  int p;

  for (int i = 0; i < N_TOUCH; i++) {
    p = touch_pins[i];

    pinMode(p, OUTPUT);
    digitalWrite(p, LOW);
    delayMicroseconds(25);

    pinMode(p, INPUT_PULLUP);
    t = 0;
    while (!digitalRead(p) && t < t_max) {
      t++;
    }
    touch_values[i] = t;

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

bool touchPressed(int idx) { return pin_touched_now[idx] && !pin_touched_past[idx]; }

// Game State Variables
bool gameStarted = false;
bool gameOver = false;
bool paused = false;
int score = 0;
int maxCombo = 0;
int currentCombo = 0;
int health = 100;
float fallSpeed = 1.5;

// Feedback variables
String currentFeedback = "";
unsigned long feedbackTimer = 0;
unsigned long hitEffectTimers[4] = {0, 0, 0, 0};

struct Note {
  int lane; // 0: Left, 1: Down, 2: Up, 3: Right
  float yPos;
  bool active;
};

#define MAX_NOTES 4
Note notes[MAX_NOTES];

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

  #ifdef NEOPIXEL_POWER
    pinMode(NEOPIXEL_POWER, OUTPUT);
    digitalWrite(NEOPIXEL_POWER, HIGH);
  #endif

  strip.begin();
  strip.setBrightness(30); 
  strip.show();

  if(!display.begin(SSD1306_SWITCHCAPVCC, 0x3C)) {
    Serial.println(F("SSD1306 allocation failed"));
    while(1);
  }

  display.clearDisplay();
}

void loop() {
  update_touch();

  if (!gameStarted) {
    setLedColor(0, 0, 0);
    displayStartScreen();
    
    if (touchPressed(4)) { 
      gameStarted = true;
      resetGame();
      delay(300);
    }
    delay(20); 
    return;
  }

  // Pause with Q0 
  if (gameStarted && !gameOver && touchPressed(5)) {
    paused = !paused;
    delay(200);
  }

  if (paused) {
    displayPauseScreen();
    delay(20);
    return;
  }

  if (gameOver) {
    setLedColor(0, 0, 0); 
    displayGameOverScreen();
    if (touchPressed(4)) { 
      gameStarted = false;
      gameOver = false;
      delay(300);
    }
    return;
  }

  updateGame();
  renderGame();
  
  if (millis() - feedbackTimer > 200) {
    setLedColor(0, 0, 0);
  }

  delay(20);
}

void resetGame() {
  gameOver = false;
  paused = false;
  score = 0;
  maxCombo = 0;
  currentCombo = 0;
  health = 100;
  fallSpeed = 1.5;
  setLedColor(0, 0, 0);

  for (int i = 0; i < 4; i++) {
    hitEffectTimers[i] = 0;
  }

  for (int i = 0; i < MAX_NOTES; i++) {
    notes[i].lane = i % 4; 
    notes[i].yPos = -20 - (i * 35); 
    notes[i].active = true;
  }
}

void updateGame() {
  fallSpeed += 0.001; 
  if (fallSpeed > 4.5) fallSpeed = 4.5; 

  for (int i = 0; i < MAX_NOTES; i++) {
    if (notes[i].active) {
      notes[i].yPos += fallSpeed;

      if (notes[i].yPos > 58) {
        notes[i].active = false;
        currentCombo = 0;
        health -= 15; 
        if (health <= 0) {
          health = 0;
          gameOver = true;
        }
        showFeedback("MISS", 100, 0, 0); 
        respawnNote(i);
      }
    }
  }

  for (int i = 0; i < 4; i++) {
    if (touchPressed(i)) {
      checkNoteHit(i);
    }
  }
}

void respawnNote(int index) {
  notes[index].lane = random(0, 4);
  float highestY = 0;
  for (int i = 0; i < MAX_NOTES; i++) {
    if (notes[i].active && notes[i].yPos < highestY) {
      highestY = notes[i].yPos;
    }
  }
  notes[index].yPos = min((float)-30, highestY - 35);
  notes[index].active = true;
}

void checkNoteHit(int lane) {
  bool hitRegistered = false;
  int targetZoneY = 50; 

  for (int n = 0; n < MAX_NOTES; n++) {
    if (notes[n].active && notes[n].lane == lane) {
      int distance = abs((int)notes[n].yPos - targetZoneY);

      if (distance <= 6) {
        score += 300;
        currentCombo++;
        showFeedback("PERF", 0, 100, 0); 
        hitRegistered = true;
      } 
      else if (distance <= 14) {
        score += 100;
        currentCombo++;
        showFeedback("GOOD", 100, 100, 0); 
        hitRegistered = true;
      } 
      else if (distance <= 22) {
        score += 50;
        currentCombo++;
        health = max(0, health - 5); // Lose health on BAD
        showFeedback("BAD", 100, 40, 0); 
        hitRegistered = true;
      }
      else if (distance <= 28) {
        currentCombo = 0;
        health -= 15; // Lose health on MISS
        if (health <= 0) {
          health = 0;
          gameOver = true;
        }
        showFeedback("MISS", 100, 0, 0); 
        hitRegistered = true;
      }

      if (hitRegistered) {
        if (currentFeedback != "MISS") {
          hitEffectTimers[lane] = millis(); 
        }
        notes[n].active = false;
        if (currentCombo > maxCombo) maxCombo = currentCombo;
        respawnNote(n);
        break;
      }
    }
  }
}

void showFeedback(String fb, uint8_t r, uint8_t g, uint8_t b) {
  currentFeedback = fb;
  feedbackTimer = millis();
  setLedColor(r, g, b);
}

void setLedColor(uint8_t r, uint8_t g, uint8_t b) {
  strip.setPixelColor(0, strip.Color(r, g, b));
  strip.show();
}

void drawTargetBoxWithArrow(int x, int y, int lane) {
  display.drawRect(x, y, 10, 10, SSD1306_WHITE);
  drawPlainArrow(x + 2, y + 2, lane);
}

void drawPlainArrow(int x, int y, int lane) {
  if (lane == 0) { // Left arrow (<)
    display.drawLine(x + 3, y + 1, x + 1, y + 3, SSD1306_WHITE);
    display.drawLine(x + 1, y + 3, x + 3, y + 5, SSD1306_WHITE);
    display.drawLine(x + 1, y + 3, x + 6, y + 3, SSD1306_WHITE);
  } else if (lane == 1) { // Down arrow (v)
    display.drawLine(x + 1, y + 3, x + 3, y + 5, SSD1306_WHITE);
    display.drawLine(x + 3, y + 5, x + 5, y + 3, SSD1306_WHITE);
    display.drawLine(x + 3, y + 1, x + 3, y + 5, SSD1306_WHITE);
  } else if (lane == 2) { // Up arrow (^)
    display.drawLine(x + 1, y + 3, x + 3, y + 1, SSD1306_WHITE);
    display.drawLine(x + 3, y + 1, x + 5, y + 3, SSD1306_WHITE);
    display.drawLine(x + 3, y + 1, x + 3, y + 5, SSD1306_WHITE);
  } else if (lane == 3) { // Right arrow (>)
    display.drawLine(x + 3, y + 1, x + 5, y + 3, SSD1306_WHITE);
    display.drawLine(x + 5, y + 3, x + 3, y + 5, SSD1306_WHITE);
    display.drawLine(x, y + 3, x + 5, y + 3, SSD1306_WHITE);
  }
}

void drawHitEffect(int x, int y, unsigned long elapsed) {
  int size = 12 - (elapsed * 9 / 300);
  if (size < 4) size = 4;
  size = (size / 2) * 2;
  
  int cx = x + 5;
  int cy = y + 5;
  int half = size / 2;

  display.drawLine(cx - half, cy - half, cx + half, cy + half, SSD1306_WHITE);
  display.drawLine(cx - half, cy + half, cx + half, cy - half, SSD1306_WHITE);
}

void displayStartScreen() {
  display.clearDisplay();
  display.setTextSize(2);
  display.setTextColor(SSD1306_WHITE);
  
  int xWiggle = (int)(sin(millis() / 120.0) * 2.0);
  display.setCursor(34 + xWiggle, 8);
  display.println(F("DANCE"));
  display.setCursor(34 - xWiggle, 26);
  display.println(F("DANCE"));

  display.setTextSize(1);
  display.setCursor(13, 52);
  display.println(F("press Q1 to start"));
  display.display();
}

void displayPauseScreen() {
  display.clearDisplay();
  display.setTextSize(1);
  display.setTextColor(SSD1306_WHITE);
  display.setCursor(45, 15);
  display.println(F("PAUSED"));

  display.setCursor(10, 40);
  display.println(F("press q0 to resume"));
  display.display();
}

void renderGame() {
  display.clearDisplay();

  // Health Bar Container
  display.drawRect(4, 10, 6, 45, SSD1306_WHITE);
  int fillHeight = (health * 43) / 100;
  display.fillRect(6, 55 - fillHeight, 2, fillHeight, SSD1306_WHITE);

  // Score Counter
  display.setTextSize(1);
  display.setTextColor(SSD1306_WHITE);
  display.setCursor(4, 0);
  display.print(score);

  // Target Boxes
  unsigned long currentTime = millis();
  for (int i = 0; i < 4; i++) {
    int boxX = 45 + (i * 12);
    int boxY = 50;
    
    drawTargetBoxWithArrow(boxX, boxY, i);

    unsigned long elapsed = currentTime - hitEffectTimers[i];
    if (elapsed < 300) {
      drawHitEffect(boxX, boxY, elapsed);
    }
  }

  // Falling Notes
  for (int i = 0; i < MAX_NOTES; i++) {
    if (notes[i].active) {
      int xCoord = 45 + (notes[i].lane * 12);
      drawPlainArrow(xCoord + 2, (int)notes[i].yPos + 2, notes[i].lane);
    }
  }

  // Combo count
  display.setCursor(104, 0);
  display.print("c:");
  display.print(currentCombo);

  // Feedback text
  if (millis() - feedbackTimer < 250) {
    int feedbackWidth = currentFeedback.length() * 6;
    display.setCursor(128 - feedbackWidth, 20);
    display.print(currentFeedback);
  }

  display.display();
}

void displayGameOverScreen() {
  display.clearDisplay();
  display.setTextSize(1);
  display.setTextColor(SSD1306_WHITE);
  
  String headerStr = "GAME OVER :(";
  int headerWidth = headerStr.length() * 6;
  display.setCursor((128 - headerWidth) / 2, 14);
  display.print(headerStr);
  
  String comboStr = "max combo: " + String(maxCombo);
  int comboWidth = comboStr.length() * 6;
  display.setCursor((128 - comboWidth) / 2, 28);
  display.print(comboStr);
  
  String restartStr = "press Q1 to restart";
  int restartWidth = restartStr.length() * 6;
  display.setCursor((128 - restartWidth) / 2, 42);
  display.print(restartStr);
  
  display.display();
}