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

#define TRIG_PIN 9
#define ECHO_PIN 8
#define SERVO_PIN 6

#define SCREEN_WIDTH 128
#define SCREEN_HEIGHT 64
#define OLED_RESET -1  // Not used with 128x64 I2C OLED

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

Servo punchServo;

long duration;
int distance;
bool triggered = false;
unsigned long triggerTime = 0;

void setup() {
  pinMode(TRIG_PIN, OUTPUT);
  pinMode(ECHO_PIN, INPUT);
  punchServo.attach(SERVO_PIN);
  punchServo.write(90);  // Rest position

  // Initialize OLED
  if (!display.begin(SSD1306_SWITCHCAPVCC, 0x3C)) {
    // OLED not found
    while (true);
  }

  display.clearDisplay();
  display.setTextSize(2);
  display.setTextColor(WHITE);
  display.setCursor(0, 20);
  display.print("Wave for");
  display.setCursor(0, 40);
  display.print("punch");
  display.display();
}

void loop() {
  // Measure distance
  digitalWrite(TRIG_PIN, LOW);
  delayMicroseconds(2);
  digitalWrite(TRIG_PIN, HIGH);
  delayMicroseconds(10);
  digitalWrite(TRIG_PIN, LOW);

  duration = pulseIn(ECHO_PIN, HIGH);
  distance = duration * 0.034 / 2;

  // Trigger action if close
  if (distance <= 2 && !triggered) {
    punchServo.write(0);  // Rotate counterclockwise
    triggerTime = millis();
    triggered = true;

    // Update OLED to "Punching"
    display.clearDisplay();
    display.setTextSize(2);
    display.setCursor(10, 25);
    display.print("Punching");
    display.display();
  }

  // Return to rest after 2 sec
  if (triggered && millis() - triggerTime >= 2000) {
    punchServo.write(120);  // Back to rest
    triggered = false;

    // Update OLED to "Wave for punch"
    display.clearDisplay();
    display.setTextSize(2);
    display.setCursor(0, 20);
    display.print("Wave for");
    display.setCursor(0, 40);
    display.print("punch");
    display.display();
  }

  delay(100);
}