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

#define SCREEN_WIDTH 128
#define SCREEN_HEIGHT 64
Adafruit_SSD1306 display(SCREEN_WIDTH, SCREEN_HEIGHT, &Wire, -1);

// กำหนดพินมอเตอร์สเต็ป
#define STEP_X 4
#define DIR_X 5
#define STEP_Y 6
#define DIR_Y 7
#define STEP_Z 8
#define DIR_Z 9

// พิน Home switch (ใช้ Push Button)
#define HOME_X A0
#define HOME_Y A1
#define HOME_Z A2

// พิน Joystick
#define JOY_X A3
#define JOY_Y A4
#define JOY_SW 2  // ปุ่มกดของจอยสติ๊ก

// กำหนดมอเตอร์สเต็ป (โหมด Step & Dir)
AccelStepper stepperX(AccelStepper::DRIVER, STEP_X, DIR_X);
AccelStepper stepperY(AccelStepper::DRIVER, STEP_Y, DIR_Y);
AccelStepper stepperZ(AccelStepper::DRIVER, STEP_Z, DIR_Z);

// ตำแหน่งปัจจุบันของมอเตอร์
long posX = 0, posY = 0, posZ = 0;

// สถานะ Homing
bool isHomed = false;

void setup() {
    Serial.begin(115200);
    pinMode(HOME_X, INPUT_PULLUP);
    pinMode(HOME_Y, INPUT_PULLUP);
    pinMode(HOME_Z, INPUT_PULLUP);
    pinMode(JOY_SW, INPUT_PULLUP);
    
    // ตั้งค่ามอเตอร์
    stepperX.setMaxSpeed(1000);
    stepperX.setAcceleration(500);
    stepperY.setMaxSpeed(1000);
    stepperY.setAcceleration(500);
    stepperZ.setMaxSpeed(1000);
    stepperZ.setAcceleration(500);

    // ตั้งค่า OLED
    if (!display.begin(SSD1306_SWITCHCAPVCC, 0x3C)) {
        Serial.println("OLED failed!");
        while (1);
    }
    display.clearDisplay();
    display.setTextSize(1);
    display.setTextColor(WHITE);

    // ไม่ต้องทำการ Homing เมื่อเริ่มต้น
    Serial.println("System ready. Press home switch to return to home position.");
}

void loop() {
    // ตรวจสอบว่ามีการกดปุ่ม Home หรือไม่
    if (digitalRead(HOME_X) == LOW || digitalRead(HOME_Y) == LOW || digitalRead(HOME_Z) == LOW) {
        returnToHome();
    }

    if (Serial.available()) {
        String command = Serial.readStringUntil('\n');
        command.trim();
        if (command.startsWith("M ")) {
            moveToPosition(command);
        }
    }

    manualControl();
    updateOLED();
}

void returnToHome() {
    Serial.println("Returning to home position...");

    // เคลื่อนที่แกน X กลับไปที่ตำแหน่ง Home
    stepperX.moveTo(0);
    while (stepperX.distanceToGo() != 0) {
        stepperX.run();
    }

    // เคลื่อนที่แกน Y กลับไปที่ตำแหน่ง Home
    stepperY.moveTo(0);
    while (stepperY.distanceToGo() != 0) {
        stepperY.run();
    }

    // เคลื่อนที่แกน Z กลับไปที่ตำแหน่ง Home
    stepperZ.moveTo(0);
    while (stepperZ.distanceToGo() != 0) {
        stepperZ.run();
    }

    posX = posY = posZ = 0;
    Serial.println("Returned to home position.");
}

void moveToPosition(String cmd) {
    cmd.remove(0, 2);
    int x, y, z;
    sscanf(cmd.c_str(), "%d %d %d", &x, &y, &z);

    Serial.print("Moving to: X=");
    Serial.print(x);
    Serial.print(" Y=");
    Serial.print(y);
    Serial.print(" Z=");
    Serial.println(z);
    
    stepperX.moveTo(x);
    stepperY.moveTo(y);
    stepperZ.moveTo(z);

    while (stepperX.distanceToGo() != 0 || stepperY.distanceToGo() != 0 || stepperZ.distanceToGo() != 0) {
        stepperX.run();
        stepperY.run();
        stepperZ.run();
    }

    posX = x;
    posY = y;
    posZ = z;
}

void manualControl() {
    int joyX = analogRead(JOY_X);
    int joyY = analogRead(JOY_Y);
    bool button = digitalRead(JOY_SW) == LOW;

    if (joyX < 400) stepperX.move(-5);
    if (joyX > 600) stepperX.move(5);
    if (joyY < 400) stepperY.move(-5);
    if (joyY > 600) stepperY.move(5);
    if (button) stepperZ.move(5);

    stepperX.run();
    stepperY.run();
    stepperZ.run();
}

void updateOLED() {
    display.clearDisplay();
    display.setCursor(0, 0);
    display.print("X: "); display.println(posX);
    display.print("Y: "); display.println(posY);
    display.print("Z: "); display.println(posZ);
    display.display();
}