const int greenLEDPin = A0;
const int redLEDPin = A1;
const int buttonPin = A2;
const int speakerPin = A3;

unsigned long ledOnTime;
bool waitingForReaction = false;
bool isGreen = false;

unsigned long bestTime = 9999;
unsigned long worstTime = 0;

void playHappyTune() {
  tone(speakerPin, 523, 200); delay(200);
  tone(speakerPin, 659, 200); delay(200);
  tone(speakerPin, 784, 300); delay(300);
  noTone(speakerPin);
}

void playNewBestTune() {
  tone(speakerPin, 784, 200); delay(200);
  tone(speakerPin, 880, 200); delay(200);
  tone(speakerPin, 988, 300); delay(300);
  tone(speakerPin, 1047, 400); delay(400);
  noTone(speakerPin);
}

void playSadTune() {
  tone(speakerPin, 440, 300); delay(300); // A4
  tone(speakerPin, 392, 300); delay(300); // G4
  tone(speakerPin, 348, 400); delay(400); // F4
  noTone(speakerPin);
}

void setup() {
  pinMode(greenLEDPin, OUTPUT);
  pinMode(redLEDPin, OUTPUT);
  pinMode(buttonPin, INPUT_PULLUP);
  pinMode(speakerPin, OUTPUT);
}

void loop() {
  digitalWrite(greenLEDPin, LOW);
  digitalWrite(redLEDPin, LOW);

  // Random delay before showing signal
  unsigned long waitStart = millis();
  unsigned long waitDuration = random(2000, 4500);

  // Check for early press during the wait
  while (millis() - waitStart < waitDuration) {
    if (digitalRead(buttonPin) == LOW) {
      // Early press detected
      digitalWrite(redLEDPin, HIGH);
      playSadTune();
      delay(2000);
      digitalWrite(redLEDPin, LOW);
      return;  // Restart the loop
    }
  }

  // Now choose which LED to light
  isGreen = random(0, 2) == 1;

  if (isGreen) {
    digitalWrite(greenLEDPin, HIGH);
  } else {
    digitalWrite(redLEDPin, HIGH);
  }

  ledOnTime = millis();
  waitingForReaction = true;

  bool redLEDTimeoutHandled = false;

  while (waitingForReaction) {
    unsigned long currentTime = millis();

    if (digitalRead(buttonPin) == LOW) {
      waitingForReaction = false;
      digitalWrite(greenLEDPin, LOW);
      digitalWrite(redLEDPin, LOW);

      if (isGreen) {
        unsigned long reactionTime = currentTime - ledOnTime;

        bool isNewBest = reactionTime < bestTime;
        if (isNewBest) bestTime = reactionTime;
        if (reactionTime > worstTime) worstTime = reactionTime;

        delay(1000);

        if (isNewBest) {
          playNewBestTune();
        } else {
          playHappyTune();
        }

      } else {
        playSadTune();
      }

      delay(2000);
    }

    if (!isGreen && !redLEDTimeoutHandled && (currentTime - ledOnTime >= 3000)) {
      waitingForReaction = false;
      redLEDTimeoutHandled = true;

      digitalWrite(redLEDPin, LOW);
      delay(1000);
    }
  }
}