#include <Wire.h>
#include <MPU6050.h>          // <-- Use this instead of Adafruit
#include <Adafruit_SSD1306.h>
#include <Adafruit_GFX.h>
#include <EEPROM.h>
#include <Fonts/FreeMono9pt7b.h>

Adafruit_SSD1306 display(128, 64, &Wire, -1);
MPU6050 mpu;                  // <-- Changed

const byte MENU_BTN = 3;
const byte ENTER_BTN = 4;

unsigned long menuStartTime = 0;
byte menuItem = 0;
const int menuTimeout = 10000;
bool precisionMode = false;

int16_t ax_offset = 0, ay_offset = 0, az_offset = 0;   // offsets as int16_t (matches original library)

void getCalibration() {
  EEPROM.get(0, ax_offset);
  EEPROM.get(2, ay_offset);
  EEPROM.get(4, az_offset);
}

void setCalibrationOffsetsToEEPROM() {
  EEPROM.put(0, ax_offset);
  EEPROM.put(2, ay_offset);
  EEPROM.put(4, az_offset);
}

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

  pinMode(MENU_BTN, INPUT_PULLUP);
  pinMode(ENTER_BTN, INPUT_PULLUP);

  display.begin(SSD1306_SWITCHCAPVCC, 0x3C);
  display.clearDisplay();
  display.setTextColor(SSD1306_WHITE);
  display.setTextSize(1);
  display.setCursor(0, 20);
  display.println(F("Starting..."));
  display.display();
  delay(800);

  Serial.println("Initializing MPU6050...");

  mpu.initialize();                     // Simple and reliable in Wokwi

  if (!mpu.testConnection()) {
    display.clearDisplay();
    display.setCursor(0, 20);
    display.println(F("MPU6050"));
    display.println(F("FAILED!"));
    display.display();
    Serial.println("MPU6050 connection failed");
    while (1) delay(10);
  }

  Serial.println("MPU6050 ready!");
  getCalibration();

  display.clearDisplay();
  display.setCursor(0, 20);
  display.println(F("MPU Ready!"));
  display.display();
  delay(1000);
}

// ====================== MENU & CALIBRATION (same as before) ======================
void dispMenu(byte item) { /* same as previous full code */ }
void dispCalibrate(byte item) { /* same */ }
void menuMainWait() { /* same */ }
void menuCalibrateWait() { /* same */ }

void setCalibration() {
  int16_t ax, ay, az;
  long sum_ax = 0, sum_ay = 0, sum_az = 0;
  const int samples = 200;

  display.clearDisplay();
  display.setCursor(0, 30);
  display.println(F("CALIBRATING"));
  display.display();

  for (int i = 0; i < samples; i++) {
    mpu.getAcceleration(&ax, &ay, &az);
    sum_ax += ax;
    sum_ay += ay;
    sum_az += az;
    delay(10);
  }

  ax_offset = sum_ax / samples;
  ay_offset = sum_ay / samples;
  az_offset = (sum_az / samples) - 16384;   // ~1g on Z when face up (16384 = 1g at ±2g range)

  setCalibrationOffsetsToEEPROM();
}

// ====================== DISPLAY ======================
void formatDisplay(double angleVal, byte dispRotate) {
  // same as in previous full code
  display.clearDisplay();
  display.setFont();
  display.setRotation(dispRotate);

  if (!precisionMode) {
    display.setTextSize(2);
    if (dispRotate == 0 || dispRotate == 2) display.setCursor(40, 10);
    else display.setCursor(3, 30);
    display.println(round(abs(angleVal)));
  } else {
    if (dispRotate == 0 || dispRotate == 2) {
      display.setTextSize(1);
      display.setCursor(20, 10);
    } else {
      display.setTextSize(1);
      display.setCursor(6, 40);
    }
    display.println(abs(angleVal), 1);
  }

  if (round(abs(angleVal)) <= 1) {
    display.setTextSize(1);
    display.setTextColor(BLACK, WHITE);
    if (dispRotate == 0 || dispRotate == 2) {
      display.setCursor(37, 50);
      display.println(F("   LEVEL  "));
    } else {
      display.setCursor(0, 90);
      display.println(F("   PLUMB  "));
    }
  }
  display.setTextColor(WHITE);
  display.display();
}

// ====================== MAIN LOOP ======================
void loop() {
  if (digitalRead(MENU_BTN) == LOW) {
    delay(500);
    menuMainWait();
  }

  int16_t ax, ay, az;
  mpu.getAcceleration(&ax, &ay, &az);

  // Apply offsets
  ax -= ax_offset;
  ay -= ay_offset;
  az -= az_offset;

  // Convert to angles (same math)
  float roll  = atan2(ay, az) * 180.0 / PI;
  float pitch = atan2(-ax, sqrt(ay * ay + az * az)) * 180.0 / PI;

  double angleVal;
  byte dispRotate = 0;

  if (abs(pitch) < abs(roll)) {
    angleVal = roll;
    dispRotate = (pitch > 0) ? 0 : 2;
  } else {
    angleVal = pitch;
    dispRotate = (roll > 0) ? 1 : 3;
  }

  formatDisplay(angleVal, dispRotate);
  delay(50);
}