#include <Wire.h>
#include <Adafruit_MPU6050.h>
#include <Adafruit_Sensor.h>
#include <Adafruit_SSD1306.h>
#include "HX711.h"

// ---------- Pin definitions ----------
#define BUTTON_PIN   4
#define LED_PIN      3
#define SDA_PIN      5
#define SCL_PIN      6
#define HX711_DT     A8
#define HX711_SCK    A9

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

// ---------- Sensors ----------
Adafruit_MPU6050 mpu;
HX711 scale;
float calibration_factor = 2280.0f;

// ---------- Timing ----------
unsigned long lastRead = 0;
const unsigned long READ_INTERVAL = 500; // ms
bool collecting = false;

void setup() {
  Serial.begin(115200);
  Serial.println("Initializing ESP32-S3 system...");

  pinMode(BUTTON_PIN, INPUT_PULLUP);
  pinMode(LED_PIN, OUTPUT);
  digitalWrite(LED_PIN, LOW);

  // --- I2C bus ---
  Wire.begin(SDA_PIN, SCL_PIN);
  Wire.setClock(100000);
  delay(100);

  // --- OLED init ---
  if (!display.begin(SSD1306_SWITCHCAPVCC, SCREEN_ADDRESS)) {
    Serial.println("OLED init failed!");
  } else {
    display.clearDisplay();
    display.setTextSize(1);
    display.setTextColor(SSD1306_WHITE);
    display.setCursor(0, 0);
    display.println("System initializing...");
    display.display();
  }

  // --- Initialize MPU6050 ---
  if (!mpu.begin(0x68)) {
    Serial.println("MPU6050 not detected at 0x68, trying 0x69...");
    if (!mpu.begin(0x69)) {
      Serial.println("MPU6050 not found! Check wiring.");
      display.clearDisplay();
      display.setCursor(0, 0);
      display.println("MPU6050 not found!");
      display.display();
    } else {
      Serial.println("MPU6050 found at 0x69");
    }
  } else {
    Serial.println("MPU6050 found at 0x68");
  }

  // --- Initialize HX711 ---
  scale.begin(HX711_DT, HX711_SCK);
  scale.set_scale(calibration_factor);
  scale.tare();
  Serial.println("HX711 initialized and tared.");

  display.clearDisplay();
  display.setCursor(0, 0);
  display.println("Setup complete!");
  display.display();
  delay(1000);
}

void loop() {
  bool pressed = (digitalRead(BUTTON_PIN) == LOW);

  digitalWrite(LED_PIN, pressed ? HIGH : LOW);
  static bool lastPressed = false;
  if (pressed != lastPressed) {
    if (pressed) Serial.println("Button pressed → LED ON");
    else         Serial.println("Button released → LED OFF");
    lastPressed = pressed;
  }

  if (pressed) {
    if (!collecting) {
      collecting = true;
      Serial.println("=== START DATA COLLECTION ===");
    }

    unsigned long now = millis();
    if (now - lastRead >= READ_INTERVAL) {
      lastRead = now;

      sensors_event_t a, g, temp;
      bool ok = mpu.getEvent(&a, &g, &temp);
      long raw = 0;
      float weight = 0;

      if (scale.is_ready()) {
        raw = scale.read();
        weight = scale.get_units(5);
      }

      // --- Serial output ---
      if (ok) {
        Serial.printf("Accel X=%.2f Y=%.2f Z=%.2f m/s²\n", a.acceleration.x, a.acceleration.y, a.acceleration.z);
        Serial.printf("Gyro X=%.2f Y=%.2f Z=%.2f rad/s\n", g.gyro.x, g.gyro.y, g.gyro.z);
        Serial.printf("Temp: %.2f °C\n", temp.temperature);
      }
      Serial.printf("HX711 raw: %ld   Weight: %.2f g\n", raw, weight);
      Serial.println("---------------------------");

      // --- OLED output ---
      display.clearDisplay();
      display.setCursor(0, 0);
      display.setTextSize(1);
      if (ok) {
        display.printf("Ax: %.1f  Ay: %.1f\n", a.acceleration.x, a.acceleration.y);
        display.printf("Az: %.1f  T: %.1fC\n", a.acceleration.z, temp.temperature);
        display.printf("Gx: %.1f  Gy: %.1f\n", g.gyro.x, g.gyro.y);
        display.printf("W: %.1f g\n", weight);
      } else {
        display.println("MPU read fail");
      }
      display.display();
    }
  } else if (collecting) {
    collecting = false;
    Serial.println("=== STOP DATA COLLECTION ===");
    display.clearDisplay();
    display.setCursor(0, 0);
    display.println("Data collection stopped.");
    display.display();
  }

  delay(10);
}