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

#define SCREEN_WIDTH 128 // OLED display width, in pixels
#define SCREEN_HEIGHT 64 // OLED display height, in pixels

#define OLED_RESET -1 // Reset pin # (or -1 if sharing Arduino reset pin)
Adafruit_SSD1306 display(SCREEN_WIDTH, SCREEN_HEIGHT, &Wire, OLED_RESET);

#define BUZZER_PIN 12
#define LED_PIN 10
#define BUTTON1_PIN 2
#define BUTTON2_PIN 3
#define BUTTON3_PIN 4

enum Route { STRAIGHT, SINE, TANGENT };
Route currentRoute = STRAIGHT;
bool startRouteFlag = false;
bool changeRouteFlag = false;

struct Obstacle {
  int x;
  int y;
};

Obstacle obstacles[3];
bool obstacleDetected = false;
int originalY;

void setup() {
  pinMode(BUZZER_PIN, OUTPUT);
  pinMode(LED_PIN, OUTPUT);
  pinMode(BUTTON1_PIN, INPUT_PULLUP);
  pinMode(BUTTON2_PIN, INPUT_PULLUP);
  pinMode(BUTTON3_PIN, INPUT_PULLUP);

  Serial.begin(9600);

  if (!display.begin(SSD1306_SWITCHCAPVCC, 0x3C)) { // Address 0x3C for 128x64
    Serial.println(F("SSD1306 allocation failed"));
    for (;;);
  }
  display.display();
  delay(500);
  display.clearDisplay();

  randomSeed(analogRead(0)); // Initialize random seed
}

void loop() {
  if (digitalRead(BUTTON1_PIN) == LOW) {
    currentRoute = STRAIGHT;
    changeRouteFlag = true;
    startRouteFlag = true;
  }
  if (digitalRead(BUTTON2_PIN) == LOW) {
    currentRoute = SINE;
    changeRouteFlag = true;
    startRouteFlag = true;
  }
  if (digitalRead(BUTTON3_PIN) == LOW) {
    currentRoute = TANGENT;
    changeRouteFlag = true;
    startRouteFlag = true;
  }

  if (changeRouteFlag) {
    initializeRoute();
    changeRouteFlag = false;
  }

  if (startRouteFlag) {
    switch (currentRoute) {
      case STRAIGHT:
        drawStraightRoute();
        break;
      case SINE:
        drawSineRoute();
        break;
      case TANGENT:
        drawTangentRoute();
        break;
    }
  }
}

void initializeRoute() {
  display.clearDisplay();
  bool obstacleOnRoute = false;
  for (int i = 0; i < 3; i++) {
    bool validPosition;
    do {
      validPosition = true;
      obstacles[i] = { random(20, 110), random(10, 50) }; // Full width for obstacles
      for (int j = 0; j < i; j++) {
        if (obstacles[i].x + 6 >= obstacles[j].x && obstacles[i].x <= obstacles[j].x + 6 &&
            obstacles[i].y + 6 >= obstacles[j].y && obstacles[i].y <= obstacles[j].y + 6) {
          validPosition = false;
          break;
        }
      }
    } while (!validPosition);

    if (detectObstacle(0, 30)) {
      obstacleOnRoute = true;
    }
  }

  // Ensure at least one obstacle on the route
  if (!obstacleOnRoute) {
    obstacles[0] = { random(20, 110), 30 };
  }

  drawObstacles();
}

void drawObstacles() {
  for (int i = 0; i < 3; i++) {
    display.fillRect(obstacles[i].x, obstacles[i].y, 6, 6, SSD1306_WHITE);
  }
  display.display();
}

bool detectObstacle(int x, int y) {
  for (int i = 0; i < 3; i++) {
    if (x + 10 >= obstacles[i].x && x <= obstacles[i].x + 6 &&
        y + 2 >= obstacles[i].y && y <= obstacles[i].y + 6) {
      return true;
    }
  }
  return false;
}

void clearPreviousRocketPosition(int x, int y) {
  display.drawLine(x, y, x + 10, y, SSD1306_BLACK); // Clear the previous horizontal line
  display.display();
}

void drawRocket(int x, int y) {
  display.drawLine(x, y, x + 10, y, SSD1306_WHITE); // Draw the rocket as a horizontal line
  display.display();
}

void drawStraightRoute() {
  int y = 30; // Initial y position of the rocket
  originalY = y;
  bool bypassingObstacle = false;
  int bypassDirection = 0; // 0 for no bypass, 1 for up, -1 for down

  for (int x = 0; x < SCREEN_WIDTH; x++) {
    clearPreviousRocketPosition(x - 1, y); // Clear the previous position

    int obstacleIndex = findCurrentObstacleIndex(x, y);
    if (obstacleIndex != -1) {
      digitalWrite(LED_PIN, HIGH); // Turn on the warning LED

      if (!bypassingObstacle) {
        int aboveDistance = y - (obstacles[obstacleIndex].y + 6);
        int belowDistance = (obstacles[obstacleIndex].y - 2) - y;

        if (aboveDistance > 0 && belowDistance < 0) {
          bypassDirection = 1; // Move up if obstacle is below
        } else if (aboveDistance < 0 && belowDistance > 0) {
          bypassDirection = -1; // Move down if obstacle is above
        } else {
          bypassDirection = (abs(aboveDistance) < abs(belowDistance)) ? 1 : -1; // Choose the shortest path
        }
        bypassingObstacle = true;
      }

      y += bypassDirection * 2; // Move up or down
    } else {
      digitalWrite(LED_PIN, LOW); // Turn off the warning LED

      if (bypassingObstacle && x > obstacles[0].x + 10 && x > obstacles[1].x + 10 && x > obstacles[2].x + 10) {
        y = originalY; // Return to the original route
        bypassingObstacle = false;
        bypassDirection = 0;
      }
    }

    drawRocket(x, y);
    drawObstacles(); // Redraw the obstacles
    delay(2); // Faster speed
  }

  playVictorySound();
  startRouteFlag = false;
}

void drawSineRoute() {
  originalY = 30;
  bool bypassingObstacle = false;
  int bypassDirection = 0; // 0 for no bypass, 1 for up, -1 for down

  for (int x = 0; x < SCREEN_WIDTH; x++) {
    int y = 30 + 10 * sin(x * PI / 64);
    clearPreviousRocketPosition(x - 1, originalY); // Clear the previous position

    int obstacleIndex = findCurrentObstacleIndex(x, y);
    if (obstacleIndex != -1) {
      digitalWrite(LED_PIN, HIGH); // Turn on the warning LED

      if (!bypassingObstacle) {
        int aboveDistance = y - (obstacles[obstacleIndex].y + 6);
        int belowDistance = (obstacles[obstacleIndex].y - 2) - y;

        if (aboveDistance > 0 && belowDistance < 0) {
          bypassDirection = 1; // Move up if obstacle is below
        } else if (aboveDistance < 0 && belowDistance > 0) {
          bypassDirection = -1; // Move down if obstacle is above
        } else {
          bypassDirection = (abs(aboveDistance) < abs(belowDistance)) ? 1 : -1; // Choose the shortest path
        }
        bypassingObstacle = true;
      }

      y += bypassDirection * 2; // Move up or down
    } else {
      digitalWrite(LED_PIN, LOW); // Turn off the warning LED

      if (bypassingObstacle && x > obstacles[0].x + 10 && x > obstacles[1].x + 10 && x > obstacles[2].x + 10) {
        y = originalY; // Return to the original route
        bypassingObstacle = false;
        bypassDirection = 0;
      }
    }

    drawRocket(x, y);
    drawObstacles(); // Redraw the obstacles
    originalY = y;
    delay(2); // Faster speed
  }

  playVictorySound();
  startRouteFlag = false;
}

void drawTangentRoute() {
  int y = 30; // Initial y position of the rocket
  bool goingUp = true;
  originalY = y;
  bool bypassingObstacle = false;
  int bypassDirection = 0; // 0 for no bypass, 1 for up, -1 for down

  for (int x = 0; x < SCREEN_WIDTH; x++) {
    clearPreviousRocketPosition(x - 1, y); // Clear the previous position

    if (x % 16 == 0) {
      goingUp = !goingUp; // Change direction every 16 pixels
    }

    int obstacleIndex = findCurrentObstacleIndex(x, y);
    if (obstacleIndex != -1) {
      digitalWrite(LED_PIN, HIGH); // Turn on the warning LED

      if (!bypassingObstacle) {
        int aboveDistance = y - (obstacles[obstacleIndex].y + 6);
        int belowDistance = (obstacles[obstacleIndex].y - 2) - y;

        if (aboveDistance > 0 && belowDistance < 0) {
          bypassDirection = 1; // Move up if obstacle is below
        } else if (aboveDistance < 0 && belowDistance > 0) {
          bypassDirection = -1; // Move down if obstacle is above
        } else {
          bypassDirection = (abs(aboveDistance) < abs(belowDistance)) ? 1 : -1; // Choose the shortest path
        }
        bypassingObstacle = true;
      }

      y += bypassDirection * 2; // Move up or down
    } else {
      digitalWrite(LED_PIN, LOW); // Turn off the warning LED

      if (bypassingObstacle && x > obstacles[0].x + 10 && x > obstacles[1].x + 10 && x > obstacles[2].x + 10) {
        y = originalY; // Return to the original route
        bypassingObstacle = false;
        bypassDirection = 0;
      }
    }

    y += goingUp ? -1 : 1;

    drawRocket(x, y);
    drawObstacles(); // Redraw the obstacles
    originalY = y;
    delay(2); // Faster speed
  }

  playVictorySound();
  startRouteFlag = false;
}

void playVictorySound() {
  int melody[] = {262, 330, 392, 523}; // Victory sound notes
  int noteDurations[] = {4, 4, 4, 4}; // Duration of each note

  for (int thisNote = 0; thisNote < 4; thisNote++) {
    int noteDuration = 1000 / noteDurations[thisNote];
    tone(BUZZER_PIN, melody[thisNote], noteDuration);
    int pauseBetweenNotes = noteDuration * 1.30;
    delay(pauseBetweenNotes);
    noTone(BUZZER_PIN);
  }
}

int findCurrentObstacleIndex(int x, int y) {
  for (int i = 0; i < 3; i++) {
    if (x + 10 >= obstacles[i].x && x <= obstacles[i].x + 6 &&
        y + 2 >= obstacles[i].y && y <= obstacles[i].y + 6) {
      return i;
    }
  }
  return -1;
}