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

// ---------- OLED ----------
#define SCREEN_WIDTH 128
#define SCREEN_HEIGHT 64
#define OLED_RESET -1
#define OLED_ADDR 0x3C
#define SDA_PIN 8
#define SCL_PIN 9

Adafruit_SSD1306 display(SCREEN_WIDTH, SCREEN_HEIGHT, &Wire, OLED_RESET);

// ---------- Servos ----------
Servo leftArm;
Servo rightArm;
Servo headServo;
Servo bodyServo;

#define LEFT_ARM_PIN 2
#define RIGHT_ARM_PIN 3
#define HEAD_PIN 4
#define BODY_PIN 5

// ---------- Robot States ----------
enum RobotState { IDLE, HAPPY, ALERT };
RobotState currentState = IDLE;
RobotState lastState = IDLE;

unsigned long lastStateChange = 0;
unsigned long lastMotionTime = 0;
uint8_t motionStep = 0;

// ---------- Setup ----------
void setup() {
  Wire.begin(SDA_PIN, SCL_PIN);

  display.begin(SSD1306_SWITCHCAPVCC, OLED_ADDR);
  display.setTextColor(SSD1306_WHITE);

  leftArm.attach(LEFT_ARM_PIN, 500, 2500);
  rightArm.attach(RIGHT_ARM_PIN, 500, 2500);
  headServo.attach(HEAD_PIN, 500, 2500);
  bodyServo.attach(BODY_PIN, 500, 2500);

  resetPose();
  showFace("READY");
}

// ---------- Loop ----------
void loop() {
  unsigned long now = millis();

  if (now - lastStateChange > 4000) {
    nextState();
    lastStateChange = now;
    motionStep = 0;
  }

  if (currentState != lastState) {
    updateFace();
    lastState = currentState;
  }

  runBehavior(now);
}

// ---------- State Logic ----------
void nextState() {
  currentState = (RobotState)((currentState + 1) % 3);
}

void runBehavior(unsigned long now) {
  switch (currentState) {
    case IDLE:  idleMode(); break;
    case HAPPY: happyMode(now); break;
    case ALERT: alertMode(now); break;
  }
}

// ---------- Behaviors ----------
void idleMode() {
  leftArm.write(80);
  rightArm.write(100);
  headServo.write(90);
  bodyServo.write(90);
}

void happyMode(unsigned long now) {
  if (now - lastMotionTime < 300) return;
  lastMotionTime = now;

  if (motionStep == 0) {
    leftArm.write(30);
    rightArm.write(150);
  } else {
    leftArm.write(90);
    rightArm.write(90);
  }
  motionStep ^= 1;
}

void alertMode(unsigned long now) {
  if (now - lastMotionTime < 200) return;
  lastMotionTime = now;

  bodyServo.write(motionStep == 0 ? 70 : 110);
  motionStep ^= 1;
}

// ---------- Helpers ----------
void resetPose() {
  leftArm.write(90);
  rightArm.write(90);
  headServo.write(90);
  bodyServo.write(90);
}

void updateFace() {
  display.clearDisplay();
  display.setTextSize(2);
  display.setCursor(0, 20);

  if (currentState == IDLE)  display.println("IDLE");
  if (currentState == HAPPY) display.println("HAPPY");
  if (currentState == ALERT) display.println("ALERT");

  display.display();
}

void showFace(const char* msg) {
  display.clearDisplay();
  display.setTextSize(2);
  display.setCursor(0, 20);
  display.println(msg);
  display.display();
}