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

#define SCREEN_WIDTH 128 
#define SCREEN_HEIGHT 64
#define VERT_PIN A0
#define HOR_PIN A1
#define SELECT_PIN 2
#define BACK_PIN 3
#define SERVO_PIN 5
#define STEPS_PER_REV 200

Adafruit_SSD1306 display(SCREEN_WIDTH, SCREEN_HEIGHT, &Wire, -1);
Stepper myStepper(STEPS_PER_REV, 8, 9, 10, 11);
Servo tiltServo;

// --- Navigation States ---
String currentOption = "MAIN_MENU"; 
int hoverIndex = 0; 
const int TOTAL_OPTIONS = 3;
String menuOptions[TOTAL_OPTIONS] = {
    "Control Drill",
    "View Readings",
    "Send Data"
};

// --- Drill Telemetry Tracking Variables ---
long currentDepth = 0;           // Dynamic rotation angle derived from steps
int tiltAngle = 90;          // Servos idle at 90 degrees by default
String drillDirection = "Stationary";
float savedTiltAngle = 90.0; // Persistent state tracking (starts at center)
const float SERVO_SPEED_FACTOR = 0.3;

// --- Joystick State Control ---
bool joystickReturnedCenter = true; 

// --- Non-blocking Timing Variables ---
unsigned long lastScreenUpdate = 0;
unsigned long lastJoystickPoll = 0;
const unsigned long SCREEN_INTERVAL = 50;   
const unsigned long JOYSTICK_INTERVAL = 20; 

// --- Interrupt & Debounce Flags ---
volatile bool selectPressed = false;
volatile bool backPressed = false;
unsigned long lastSelectInterrupt = 0;
unsigned long lastBackInterrupt = 0;
const unsigned long DEBOUNCE_TIME = 200; 

void handleSelect() {
    unsigned long now = millis();
    if (now - lastSelectInterrupt > DEBOUNCE_TIME) {
        selectPressed = true;
        lastSelectInterrupt = now;
    }
}

void handleBack() {
    unsigned long now = millis();
    if (now - lastBackInterrupt > DEBOUNCE_TIME) {
        backPressed = true;
        lastBackInterrupt = now;
    }
}

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

    pinMode(VERT_PIN, INPUT);
    pinMode(HOR_PIN, INPUT);
    pinMode(SELECT_PIN, INPUT_PULLUP);
    pinMode(BACK_PIN, INPUT_PULLUP);

    attachInterrupt(digitalPinToInterrupt(SELECT_PIN), handleSelect, FALLING);
    attachInterrupt(digitalPinToInterrupt(BACK_PIN), handleBack, FALLING);

    tiltServo.attach(SERVO_PIN);
    tiltServo.write(tiltAngle); // Pre-align mechanism to center

    if(!display.begin(SSD1306_SWITCHCAPVCC, 0x3C)) { 
        Serial.println(F("SSD1306 allocation failed"));
        for(;;);
    }

    display.clearDisplay();
    display.display();
}

void loop() {
    unsigned long currentMillis = millis();

    // Poll Joystick
    if (currentMillis - lastJoystickPoll >= JOYSTICK_INTERVAL) {
        lastJoystickPoll = currentMillis;
        if (currentOption == "MAIN_MENU") {
            JoyStickMenuControl();
        } else if (currentOption == "CONTROL_DRILL") {
            JoyStickStepperAndServoControl(); 
        }
    }

    handleButtonActions();

    // Screen refresh handler
    if (currentMillis - lastScreenUpdate >= SCREEN_INTERVAL) {
        lastScreenUpdate = currentMillis;
        if (currentOption != "CONTROL_DRILL") {
            renderMenu(currentOption, menuOptions[hoverIndex]);
        } else {
            renderControlDrill();
        }
    }
}

// --- Smooth Single-Step Joystick & Servo Processing ---
void JoyStickStepperAndServoControl() {
    int rawVert = analogRead(VERT_PIN);
    int rawHor = analogRead(HOR_PIN);

    // 1. Incremental Servo Tilt Control (Lowered Speed & Intact State)
    int horDeatzone = 40;
    int horCenter = 512;

    if (rawHor > (horCenter + horDeatzone)) {
        // Push Right -> Increase angle slowly based on joystick deflection displacement
        float speedModifier = (rawHor - horCenter) * SERVO_SPEED_FACTOR;
        savedTiltAngle += (speedModifier / 100.0); 
    } 
    else if (rawHor < (horCenter - horDeatzone)) {
        // Push Left -> Decrease angle slowly
        float speedModifier = (horCenter - rawHor) * SERVO_SPEED_FACTOR;
        savedTiltAngle -= (speedModifier / 100.0);
    }

    // Enforce strict physical structural bounds for a standard servo
    if (savedTiltAngle > 180.0) savedTiltAngle = 180.0;
    if (savedTiltAngle < 0.0)   savedTiltAngle = 0.0;
    // Convert float tracking state to discrete integer values for the servo hardware
    tiltAngle = (int)savedTiltAngle;
    tiltServo.write(tiltAngle);


    // 2. Stepper Depth Speed & Step Engine
    if (rawVert > 550) {
        if (currentDepth <= 0) {
            currentDepth = 0; 
        } else {
            int speed = (int)(((rawVert - 512) / 512.0) * 60.0);
            if (speed < 1) speed = 1; 
            
            myStepper.setSpeed(speed);
            myStepper.step(-1);
            
            currentDepth--; 
            drillDirection = "Upwards";
        }
    } 
    else if (rawVert < 480) {
        if (currentDepth >= 100) {
            currentDepth = 100;
        } else {
            int speed = (int)(((512 - rawVert) / 512.0) * 60.0);
            if (speed < 1) speed = 1;
            
            myStepper.setSpeed(speed);
            myStepper.step(1);
            
            currentDepth++; 
            drillDirection = "Downwards";
        }
    } 
    else {
        drillDirection = "Stationary";
    }
}

void JoyStickMenuControl() {
    int vertValue = analogRead(VERT_PIN);
    
    if (vertValue < 300) { 
        if (joystickReturnedCenter) {
            if (hoverIndex > 0) hoverIndex--;
            joystickReturnedCenter = false; 
        }
    } 
    else if (vertValue > 700) { 
        if (joystickReturnedCenter) {
            if (hoverIndex < TOTAL_OPTIONS - 1) hoverIndex++;
            joystickReturnedCenter = false; 
        }
    } 
    else {
        joystickReturnedCenter = true; 
    }
}

void handleButtonActions() {
    if (selectPressed) {
        selectPressed = false; 
        
        if (currentOption == "MAIN_MENU") {
            if (hoverIndex == 0) currentOption = "CONTROL_DRILL";
            else if (hoverIndex == 1) currentOption = "VIEW_READINGS";
            else if (hoverIndex == 2) currentOption = "SEND_DATA";
        }
    }

    if (backPressed) {
        backPressed = false; 
        if (currentOption != "MAIN_MENU") {
            currentOption = "MAIN_MENU"; 
        }
    }
}

void renderMenu(String current, String hover) {
    display.clearDisplay();
    display.setTextSize(1);
    display.setTextColor(SSD1306_WHITE);

    if (current == "MAIN_MENU") {
        display.setCursor(0, 0);
        display.println(F("Thalassic Excavator"));
        display.drawFastHLine(0, 10, SCREEN_WIDTH, SSD1306_WHITE);
        
        for (int i = 0; i < TOTAL_OPTIONS; i++) {
            display.setCursor(10, 16 + (i * 14));
            
            if (menuOptions[i] == hover) {
                display.print(F("+ "));
            } else {
                display.print(F("  "));
            }
            
            display.print(i + 1);
            display.print(F(". "));
            display.println(menuOptions[i]);
        }
    } 
    else if (current == "VIEW_READINGS") {
        display.setCursor(0, 0);
        display.println(F("Sensor Telemetry"));
        display.drawFastHLine(0, 10, SCREEN_WIDTH, SSD1306_WHITE);
        display.setCursor(0, 24);
        display.println(F("Displaying readings.."));
        
        // Consistent Bottom Right Corner layout alignment
        display.setCursor(84, 55);
        display.println(F("< main"));
    }
    else if (current == "SEND_DATA") {
        display.setCursor(0, 0);
        display.println(F("Comms Module"));
        display.drawFastHLine(0, 10, SCREEN_WIDTH, SSD1306_WHITE);
        display.setCursor(0, 24);
        display.println(F("Transmitting logs..."));
        
        // Consistent Bottom Right Corner layout alignment
        display.setCursor(84, 55);
        display.println(F("< main"));
    }

    display.display();
}

// --- Dynamic Telemetry Dashboard Function ---
void renderControlDrill() {
    display.clearDisplay();
    display.setTextSize(1);
    display.setTextColor(SSD1306_WHITE);

    // Title Header
    display.setCursor(0, 0);
    display.println(F("DRILL MANUAL MODE"));
    display.drawFastHLine(0, 10, SCREEN_WIDTH, SSD1306_WHITE);

    // Depth Angle data conversion display
    display.setCursor(0, 16);
    display.print(F("Depth Step  : "));
    display.print(currentDepth, 1); 
    display.println(F(" deg"));

    // Dynamic Tilt Angle Data (Servo angle calculation output)
    display.setCursor(0, 28);
    display.print(F("Tilt Angle : "));
    display.print(tiltAngle);
    display.println(F(" deg"));

    // Directional Status
    display.setCursor(0, 40);
    display.print(F("Status     : "));
    display.println(drillDirection);

    // Return Notice Footer aligned uniformly to the bottom right
    display.drawFastHLine(0, 52, SCREEN_WIDTH, SSD1306_WHITE);
    display.setCursor(84, 55);
    display.println(F("< main"));

    display.display();
}