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

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

// ── 74HC595 x2 (LED output) ──────────────────
#define DS    23   // Serial data
#define SHCP  18   // Shift clock
#define STCP   5   // Latch

// ── 74HC165 x2 (Button input) ────────────────
#define SH_LD  4   // Parallel Load (PL)
#define CLK   16   // Clock
#define QH    19   // Serial data out (from last 165)

// ── Game state ───────────────────────────────
uint16_t litLEDs    = 0;   // bitmask of currently glowing LEDs
int      score      = 0;
int      misses     = 0;
bool     gameOver   = false;

// ─────────────────────────────────────────────
void setup() {
  pinMode(DS,    OUTPUT);
  pinMode(SHCP,  OUTPUT);
  pinMode(STCP,  OUTPUT);
  pinMode(SH_LD, OUTPUT);
  pinMode(CLK,   OUTPUT);
  pinMode(QH,    INPUT);

  digitalWrite(SH_LD, HIGH);
  digitalWrite(CLK,   LOW);

  Serial.begin(115200);
  randomSeed(analogRead(0));

  // OLED init
  if (!display.begin(SSD1306_SWITCHCAPVCC, 0x3C)) {
    Serial.println("OLED not found!");
    while (true);
  }
  display.clearDisplay();
  display.setTextColor(WHITE);

  writeLEDs(0);
  showSplash();
  delay(2000);
}

// ── Write 16 bits to two chained 74HC595 ─────
void writeLEDs(uint16_t val) {
  digitalWrite(STCP, LOW);
  // Send high byte first (LEDs 9-16), then low byte (LEDs 1-8)
  shiftOut(DS, SHCP, MSBFIRST, (val >> 8) & 0xFF);
  shiftOut(DS, SHCP, MSBFIRST, val & 0xFF);
  digitalWrite(STCP, HIGH);
}

// ── Read 16 bits from two chained 74HC165 ────
uint16_t readButtons() {
  // Latch / sample all inputs
  digitalWrite(SH_LD, LOW);
  delayMicroseconds(5);
  digitalWrite(SH_LD, HIGH);

  uint16_t data = 0;
  // u4 shifts out first (buttons 9-16), then u3 (buttons 1-8)
  for (int i = 15; i >= 0; i--) {
    if (digitalRead(QH)) data |= (1 << i);
    digitalWrite(CLK, HIGH);
    delayMicroseconds(2);
    digitalWrite(CLK, LOW);
    delayMicroseconds(2);
  }
  return data;
}

// ── OLED helpers ─────────────────────────────
void showSplash() {
  display.clearDisplay();
  display.setTextSize(2);
  display.setCursor(10, 10);
  display.println("LED GAME");
  display.setTextSize(1);
  display.setCursor(15, 40);
  display.println("Press any button");
  display.setCursor(20, 52);
  display.println("to start!");
  display.display();
}

void updateOLED(int ledNum, bool correct) {
  display.clearDisplay();

  // Title
  display.setTextSize(1);
  display.setCursor(0, 0);
  display.println("--- LED GAME ---");

  // Active LED
  display.setTextSize(2);
  display.setCursor(0, 16);
  display.print("LED: ");
  display.println(ledNum);

  // Score / misses
  display.setTextSize(1);
  display.setCursor(0, 40);
  display.print("Score: ");
  display.println(score);
  display.setCursor(0, 52);
  display.print("Misses: ");
  display.println(misses);

  // Feedback
  display.setTextSize(1);
  display.setCursor(80, 40);
  if (correct) {
    display.println("CORRECT!");
  } else {
    display.println("WRONG!");
  }

  display.display();
}

void showGameOver() {
  display.clearDisplay();
  display.setTextSize(2);
  display.setCursor(10, 5);
  display.println("GAME");
  display.setCursor(10, 28);
  display.println("OVER!");
  display.setTextSize(1);
  display.setCursor(0, 52);
  display.print("Score: ");
  display.print(score);
  display.print("  Miss: ");
  display.println(misses);
  display.display();
}

void showWaiting(int ledNum) {
  display.clearDisplay();
  display.setTextSize(1);
  display.setCursor(0, 0);
  display.println("--- LED GAME ---");
  display.setTextSize(2);
  display.setCursor(0, 16);
  display.print("LED ");
  display.println(ledNum);
  display.setTextSize(1);
  display.setCursor(0, 40);
  display.println("Press matching");
  display.setCursor(0, 52);
  display.println("button!");
  display.setCursor(80, 40);
  display.print("S:");
  display.print(score);
  display.print(" M:");
  display.println(misses);
  display.display();
}

// ─────────────────────────────────────────────
void loop() {
  if (gameOver) {
    showGameOver();
    writeLEDs(litLEDs);   // keep wrong LEDs glowing
    delay(100);
    return;
  }

  // Pick a random LED that is NOT already lit
  int newLED;
  do {
    newLED = random(0, 16);   // 0–15
  } while (litLEDs & (1 << newLED));

  uint16_t newMask = (1 << newLED);
  litLEDs |= newMask;          // add to glowing set
  writeLEDs(litLEDs);

  Serial.printf("LED %d ON → press Button %d\n", newLED + 1, newLED + 1);
  showWaiting(newLED + 1);

  // Wait for a button press
  unsigned long start = millis();
  bool answered = false;

  while (!answered) {
    uint16_t raw     = readButtons();
    uint16_t pressed = ~raw;          // pull-up wiring → invert

    // Check if correct button pressed
    if (pressed & newMask) {
      // ✅ Correct!
      litLEDs &= ~newMask;            // turn off this LED
      writeLEDs(litLEDs);
      score++;
      Serial.printf("✅ Correct! Score: %d\n", score);
      updateOLED(newLED + 1, true);
      delay(400);
      answered = true;
    }
    // Check if WRONG button pressed
    else if (pressed != 0) {
      // ❌ Wrong button — LED stays on forever, misses++
      misses++;
      Serial.printf("❌ Wrong button! LED %d stays ON. Misses: %d\n",
                    newLED + 1, misses);
      updateOLED(newLED + 1, false);
      delay(600);
      answered = true;

      // Game over after 3 misses
      if (misses >= 3) {
        gameOver = true;
        return;
      }
    }

    // All 16 LEDs permanently on = game over
    if (litLEDs == 0xFFFF) {
      gameOver = true;
      return;
    }

    delay(10);
  }
}