#include <Arduino.h>
#include <AccelStepper.h>

#include "config.h"
#include "config_pins.h"
#include "common.h"
#include "lcd.h"
#include "data.h"
#include "states.h"
#include "steppers.h"

static AccelStepper stepper_x(AccelStepper::DRIVER, X_STEP_PIN, X_DIR_PIN, 0, 0, false);
static AccelStepper stepper_y(AccelStepper::DRIVER, Y_STEP_PIN, Y_DIR_PIN, 0, 0, false);
static AccelStepper stepper_z(AccelStepper::DRIVER, Z_STEP_PIN, Z_DIR_PIN, 0, 0, false);
static AccelStepper stepper_e(AccelStepper::DRIVER, E0_STEP_PIN, E0_DIR_PIN, 0, 0, false);

enum stepper_states {
    step_disabled,
    step_not_homed,
    step_homing_x_fast,
    step_homing_x_back,
    step_homing_x_slow,
    step_homing_y_fast,
    step_homing_y_back,
    step_homing_y_slow,
    step_homing_z_fast,
    step_homing_z_back,
    step_homing_z_slow,
    step_homing_e_fast,
    step_homing_e_back,
    step_homing_e_slow,
    step_homed
};

static volatile stepper_states state = step_disabled;
static volatile unsigned long steppers_home_move_start_time = 0;

void steppers_init(void) {
    pinMode(X_ENDSTOP_PIN, INPUT_PULLUP);
    pinMode(Y_ENDSTOP_PIN, INPUT_PULLUP);
    pinMode(Z_ENDSTOP_PIN, INPUT_PULLUP);
    pinMode(E_ENDSTOP_PIN, INPUT_PULLUP);

    pinMode(X_ENABLE_PIN, OUTPUT);
    pinMode(X_STEP_PIN, OUTPUT);
    pinMode(X_DIR_PIN, OUTPUT);
    pinMode(Y_ENABLE_PIN, OUTPUT);
    pinMode(Y_STEP_PIN, OUTPUT);
    pinMode(Y_DIR_PIN, OUTPUT);
    pinMode(Z_ENABLE_PIN, OUTPUT);
    pinMode(Z_STEP_PIN, OUTPUT);
    pinMode(Z_DIR_PIN, OUTPUT);
    pinMode(E0_ENABLE_PIN, OUTPUT);
    pinMode(E0_STEP_PIN, OUTPUT);
    pinMode(E0_DIR_PIN, OUTPUT);

    stepper_x.setEnablePin(X_ENABLE_PIN);
    stepper_x.setPinsInverted(false, false, true);
    steppers_set_speed_x(data_options()->speed_x);
    steppers_set_accel_x(data_options()->accel_x);

    stepper_y.setEnablePin(Y_ENABLE_PIN);
    stepper_y.setPinsInverted(false, false, true);
    steppers_set_speed_y(data_options()->speed_y);
    steppers_set_accel_y(data_options()->accel_y);

    stepper_z.setEnablePin(Z_ENABLE_PIN);
    stepper_z.setPinsInverted(false, false, true);
    steppers_set_speed_z(data_options()->speed_z);
    steppers_set_accel_z(data_options()->accel_z);

    stepper_e.setEnablePin(E0_ENABLE_PIN);
    stepper_e.setPinsInverted(false, false, true);
    steppers_set_speed_e(data_options()->speed_e);
    steppers_set_accel_e(data_options()->accel_e);
}

static void steppers_initiate_home(int axis, int phase) {
    Serial.print(F("steppers_initiate_home("));
    if (axis == 0) {
        Serial.print('X');
    } else if (axis == 1) {
        Serial.print('Y');
    } else if (axis == 2) {
        Serial.print('Z');
    } else if (axis == 3) {
        Serial.print('E');
    } else {
        Serial.print(axis);
    }
    Serial.print(F(", "));
    Serial.print(phase);
    Serial.println(F(")"));

    CLEAR_STORE_INTERRUPTS();

    steppers_home_move_start_time = millis();

    if (axis == 0) {
        // x
        if (phase == 0) {
            state = step_homing_x_fast;
            stepper_x.setSpeed(X_HOMING_DIR * X_AXIS_MOVEMENT_DIR * XY_FAST_HOME_SPEED * XY_STEPS_PER_MM);
        } else if (phase == 1) {
            state = step_homing_x_back;
            stepper_x.setSpeed(-1.0 * X_HOMING_DIR * X_AXIS_MOVEMENT_DIR * XY_FAST_HOME_SPEED * XY_STEPS_PER_MM);
        } else if (phase == 2) {
            state = step_homing_x_slow;
            stepper_x.setSpeed(X_HOMING_DIR * X_AXIS_MOVEMENT_DIR * XY_SLOW_HOME_SPEED * XY_STEPS_PER_MM);
        }
    } else if (axis == 1) {
        // y
        if (phase == 0) {
            state = step_homing_y_fast;
            stepper_y.setSpeed(Y_HOMING_DIR * Y_AXIS_MOVEMENT_DIR * XY_FAST_HOME_SPEED * XY_STEPS_PER_MM);
        } else if (phase == 1) {
            state = step_homing_y_back;
            stepper_y.setSpeed(-1.0 * Y_HOMING_DIR * Y_AXIS_MOVEMENT_DIR * XY_FAST_HOME_SPEED * XY_STEPS_PER_MM);
        } else if (phase == 2) {
            state = step_homing_y_slow;
            stepper_y.setSpeed(Y_HOMING_DIR * Y_AXIS_MOVEMENT_DIR * XY_SLOW_HOME_SPEED * XY_STEPS_PER_MM);
        }
    } else if (axis == 2) {
        // z
        if (phase == 0) {
            state = step_homing_z_fast;
            stepper_z.setSpeed(Z_HOMING_DIR * Z_AXIS_MOVEMENT_DIR * Z_FAST_HOME_SPEED * Z_STEPS_PER_MM);
        } else if (phase == 1) {
            state = step_homing_z_back;
            stepper_z.setSpeed(-1.0 * Z_HOMING_DIR * Z_AXIS_MOVEMENT_DIR * Z_FAST_HOME_SPEED * Z_STEPS_PER_MM);
        } else if (phase == 2) {
            state = step_homing_z_slow;
            stepper_z.setSpeed(Z_HOMING_DIR * Z_AXIS_MOVEMENT_DIR * Z_SLOW_HOME_SPEED * Z_STEPS_PER_MM);
        }
    } else if (axis == 3) {
        // e
        if (phase == 0) {
            state = step_homing_e_fast;
            stepper_e.setSpeed(E_HOMING_DIR * E_AXIS_MOVEMENT_DIR * E_FAST_HOME_SPEED * E_STEPS_PER_PERCENT);
        } else if (phase == 1) {
            state = step_homing_e_back;
            stepper_e.setSpeed(-1.0 * E_HOMING_DIR * E_AXIS_MOVEMENT_DIR * E_FAST_HOME_SPEED * E_STEPS_PER_PERCENT);
        } else if (phase == 2) {
            state = step_homing_e_slow;
            stepper_e.setSpeed(E_HOMING_DIR * E_AXIS_MOVEMENT_DIR * E_SLOW_HOME_SPEED * E_STEPS_PER_PERCENT);
        }
    } else {
        Serial.println(F("home_init error: invalid axis"));
    }

    RESTORE_INTERRUPTS();
}

bool steppers_home_switch(int axis) {
    bool r = true;
    CLEAR_STORE_INTERRUPTS();

    if (axis == 0) {
        r = digitalRead(X_ENDSTOP_PIN) == LOW;
    } else if (axis == 1) {
        r = digitalRead(Y_ENDSTOP_PIN) == LOW;
    } else if (axis == 2) {
        r = digitalRead(Z_ENDSTOP_PIN) == LOW;
    } else if (axis == 3) {
        r = digitalRead(E_ENDSTOP_PIN) == LOW;
    } else {
        Serial.println(F("home_switch error: invalid axis"));
    }

    RESTORE_INTERRUPTS();
    return r;
}

float steppers_current_pos(int axis) {
    float r = 0.0;
    CLEAR_STORE_INTERRUPTS();

    if (axis == 0) {
        float v = stepper_x.currentPosition();
        r = v / XY_STEPS_PER_MM / X_AXIS_MOVEMENT_DIR;
    } else if (axis == 1) {
        float v = stepper_y.currentPosition();
        r = v / XY_STEPS_PER_MM / Y_AXIS_MOVEMENT_DIR;
    } else if (axis == 2) {
        float v = stepper_z.currentPosition();
        r = v / Z_STEPS_PER_MM / Z_AXIS_MOVEMENT_DIR;
    } else if (axis == 3) {
        float v = stepper_e.currentPosition();
        r = E_PERCENT_TO_MM(v / E_STEPS_PER_PERCENT / E_AXIS_MOVEMENT_DIR);
    } else {
        Serial.println(F("steppers_pos error: invalid axis"));
    }

    RESTORE_INTERRUPTS();
    return r;
}

bool steppers_run(void) {
    if (state == step_homing_x_fast) {
        if (steppers_home_switch(0)) {
            steppers_initiate_home(0, 1);
        } else {
            stepper_x.runSpeed();
        }
    } else if (state == step_homing_x_back) {
        unsigned long end_time = steppers_home_move_start_time + XY_HOME_BACK_OFF_TIME;
        if ((!steppers_home_switch(0)) && (millis() >= end_time)) {
            steppers_initiate_home(0, 2);
        } else {
            stepper_x.runSpeed();
        }
    } else if (state == step_homing_x_slow) {
        if (steppers_home_switch(0)) {
            stepper_x.setSpeed(0);

#if (X_MIN_PIN == -1)
            stepper_x.setCurrentPosition(X_AXIS_MAX * XY_STEPS_PER_MM * X_AXIS_MOVEMENT_DIR);
#else
            stepper_x.setCurrentPosition(X_AXIS_MIN * XY_STEPS_PER_MM * X_AXIS_MOVEMENT_DIR);
#endif

            steppers_initiate_home(1, 0);
        } else {
            stepper_x.runSpeed();
        }
    } else if (state == step_homing_y_fast) {
        if (steppers_home_switch(1)) {
            steppers_initiate_home(1, 1);
        } else {
            stepper_y.runSpeed();
        }
    } else if (state == step_homing_y_back) {
        unsigned long end_time = steppers_home_move_start_time + XY_HOME_BACK_OFF_TIME;
        if ((!steppers_home_switch(1)) && (millis() >= end_time)) {
            steppers_initiate_home(1, 2);
        } else {
            stepper_y.runSpeed();
        }
    } else if (state == step_homing_y_slow) {
        if (steppers_home_switch(1)) {
            stepper_y.setSpeed(0);

#if (Y_MIN_PIN == -1)
            stepper_y.setCurrentPosition(Y_AXIS_MAX * XY_STEPS_PER_MM * Y_AXIS_MOVEMENT_DIR);
#else
            stepper_y.setCurrentPosition(Y_AXIS_MIN * XY_STEPS_PER_MM * Y_AXIS_MOVEMENT_DIR);
#endif

            //steppers_initiate_home(3, 0);
            // TODO
            state = step_homed;
            return false;
        } else {
            stepper_y.runSpeed();
        }
    } else if (state == step_homing_z_fast) {
        if (steppers_home_switch(2)) {
            steppers_initiate_home(2, 1);
        } else {
            stepper_z.runSpeed();
        }
    } else if (state == step_homing_z_back) {
        unsigned long end_time = steppers_home_move_start_time + Z_HOME_BACK_OFF_TIME;
        if ((!steppers_home_switch(2)) && (millis() >= end_time)) {
            steppers_initiate_home(2, 2);
        } else {
            stepper_z.runSpeed();
        }
    } else if (state == step_homing_z_slow) {
        if (steppers_home_switch(2)) {
            stepper_z.setSpeed(0);

#if (Z_MIN_PIN == -1)
            stepper_z.setCurrentPosition(Z_AXIS_MAX * Z_STEPS_PER_MM * Z_AXIS_MOVEMENT_DIR);
#else
            stepper_z.setCurrentPosition(Z_AXIS_MIN * Z_STEPS_PER_MM * Z_AXIS_MOVEMENT_DIR);
#endif

            steppers_initiate_home(0, 0);
        } else {
            stepper_z.runSpeed();
        }
    } else if (state == step_homing_e_fast) {
        if (steppers_home_switch(3)) {
            steppers_initiate_home(3, 1);
        } else {
            stepper_e.runSpeed();
        }
    } else if (state == step_homing_e_back) {
        unsigned long end_time = steppers_home_move_start_time + E_HOME_BACK_OFF_TIME;
        if ((!steppers_home_switch(3)) && (millis() >= end_time)) {
            steppers_initiate_home(3, 2);
        } else {
            stepper_e.runSpeed();
        }
    } else if (state == step_homing_e_slow) {
        if (steppers_home_switch(3)) {
            stepper_e.setSpeed(0);

#if (E_MIN_PIN == -1)
            stepper_e.setCurrentPosition(E_MM_TO_PERCENT(E_AXIS_MAX) * E_STEPS_PER_PERCENT * E_AXIS_MOVEMENT_DIR);
#else
            stepper_e.setCurrentPosition(E_MM_TO_PERCENT(E_AXIS_MIN) * E_STEPS_PER_PERCENT * E_AXIS_MOVEMENT_DIR);
#endif

            state = step_homed;
            return false;
        } else {
            stepper_e.runSpeed();
        }
    } else if (state != step_disabled) {
        if (state == step_homed) {
            for (int i = 0; i < 4; i++) {
                Serial.print(i);
                Serial.print(": ");
                Serial.println(steppers_current_pos(i));

                float compare = 0.0, epsilon = 0.0;
                if (i == 0) {
                    epsilon = X_AXIS_EPSILON;
#if (X_MIN_PIN == -1)
                    compare = X_AXIS_MAX;
#else
                    compare = X_AXIS_MIN;
#endif
                } else if (i == 1) {
                    epsilon = Y_AXIS_EPSILON;
#if (Y_MIN_PIN == -1)
                    compare = Y_AXIS_MAX;
#else
                    compare = Y_AXIS_MIN;
#endif
                } else if (i == 2) {
                    epsilon = Z_AXIS_EPSILON;
#if (Z_MIN_PIN == -1)
                    compare = Z_AXIS_MAX;
#else
                    compare = Z_AXIS_MIN;
#endif
                } else if (i == 3) {
                    epsilon = E_AXIS_EPSILON;
#if (E_MIN_PIN == -1)
                    compare = E_AXIS_MAX;
#else
                    compare = E_AXIS_MIN;
#endif
                }

                if (fabs(steppers_current_pos(i) - compare) > epsilon) {
                    if (steppers_home_switch(i)) {
                        Serial.print(F("Endstop hit at "));
                        Serial.println(steppers_current_pos(i));

                        steppers_kill();

                        if (i == 0) {
                            sm_error.setText("Enstop hit on X axis");
                        } else if (i == 1) {
                            sm_error.setText("Enstop hit on Y axis");
                        } else if (i == 2) {
                            sm_error.setText("Enstop hit on Z axis");
                        } else if (i == 3) {
                            sm_error.setText("Enstop hit on E axis");
                        } else {
                            sm_error.setText("Enstop hit on unknown axis");
                        }
                        states_go_to(&sm_error);
                    }
                }
            }
        }

        boolean x = stepper_x.run();
        boolean y = stepper_y.run();
        boolean z = stepper_z.run();
        boolean e = stepper_e.run();

        return x || y || z || e;
    }

    return true;
}

bool steppers_homed(void) {
    CLEAR_STORE_INTERRUPTS();
    bool r = (state == step_homed);
    RESTORE_INTERRUPTS();
    return r;
}

char steppers_homing_axis(void) {
    CLEAR_STORE_INTERRUPTS();

    char r = ' ';
    if ((state == step_homing_x_fast) || (state == step_homing_x_back) || (state == step_homing_x_slow)) {
        r = 'X';
    } else if ((state == step_homing_y_fast) || (state == step_homing_y_back) || (state == step_homing_y_slow)) {
        r = 'Y';
    } else if ((state == step_homing_z_fast) || (state == step_homing_z_back) || (state == step_homing_z_slow)) {
        r = 'Z';
    } else if ((state == step_homing_e_fast) || (state == step_homing_e_back) || (state == step_homing_e_slow)) {
        r = 'E';
    }

    RESTORE_INTERRUPTS();
    return r;
}

void steppers_start_homing(void) {
    CLEAR_STORE_INTERRUPTS();

    stepper_x.enableOutputs();
    stepper_y.enableOutputs();
    stepper_z.enableOutputs();
    stepper_e.enableOutputs();

    //steppers_initiate_home(2, 0);
    // TODO
    steppers_initiate_home(0, 0);

    RESTORE_INTERRUPTS();
}

static int steppers_move_axis(AccelStepper &axis, long pos) {
    CLEAR_STORE_INTERRUPTS();

    if (state == step_disabled) {
        Serial.println(F("Enabling stepper drivers"));

        stepper_x.enableOutputs();
        stepper_y.enableOutputs();
        stepper_z.enableOutputs();
        stepper_e.enableOutputs();
        state = step_not_homed;
    }

    if (!steppers_homed()) {
        Serial.println(F("WARNING: un-homed move!"));
    }

    axis.moveTo(pos);

    RESTORE_INTERRUPTS();
    return 0;
}

int steppers_move_x(float pos) {
    Serial.print(F("Moving X to "));
    Serial.print(pos);
    Serial.print(F(" mm ("));
    Serial.print(pos * XY_STEPS_PER_MM * X_AXIS_MOVEMENT_DIR);
    Serial.println(F(" steps)"));

    if (pos < X_AXIS_MIN) {
        Serial.println(F("Move below X_AXIS_MIN!"));
        return -1;
    }

    if (pos > X_AXIS_MAX) {
        Serial.println(F("Move above X_AXIS_MAX!"));
        return -1;
    }

    return steppers_move_axis(stepper_x, pos * XY_STEPS_PER_MM * X_AXIS_MOVEMENT_DIR);
}

int steppers_move_y(float pos) {
    Serial.print(F("Moving Y to "));
    Serial.print(pos);
    Serial.print(F(" mm ("));
    Serial.print(pos * XY_STEPS_PER_MM * Y_AXIS_MOVEMENT_DIR);
    Serial.println(F(" steps)"));

    if (pos < Y_AXIS_MIN) {
        Serial.println(F("Move below Y_AXIS_MIN!"));
        return -1;
    }

    if (pos > Y_AXIS_MAX) {
        Serial.println(F("Move above Y_AXIS_MAX!"));
        return -1;
    }

    return steppers_move_axis(stepper_y, pos * XY_STEPS_PER_MM * Y_AXIS_MOVEMENT_DIR);
}

int steppers_move_z(float pos) {
    Serial.print(F("Moving Z to "));
    Serial.print(pos);
    Serial.print(F(" mm ("));
    Serial.print(pos * Z_STEPS_PER_MM * Z_AXIS_MOVEMENT_DIR);
    Serial.println(F(" steps)"));

    if (pos < Z_AXIS_MIN) {
        Serial.println(F("Move below Z_AXIS_MIN!"));
        return -1;
    }

    if (pos > Z_AXIS_MAX) {
        Serial.println(F("Move above Z_AXIS_MAX!"));
        return -1;
    }

    return steppers_move_axis(stepper_z, pos * Z_STEPS_PER_MM * Z_AXIS_MOVEMENT_DIR);
}

int steppers_move_e(float pos) {
    Serial.print(F("Moving E to "));
    Serial.print(pos);
    Serial.print(F(" mm ("));
    Serial.print(E_MM_TO_PERCENT(pos));
    Serial.println(F("% = "));
    Serial.print(E_MM_TO_PERCENT(pos) * E_STEPS_PER_PERCENT * E_AXIS_MOVEMENT_DIR);
    Serial.println(F(" steps)"));

    if (pos < E_AXIS_MIN) {
        Serial.println(F("Move below E_AXIS_MIN!"));
        return -1;
    }

    if (pos > E_AXIS_MAX) {
        Serial.println(F("Move above E_AXIS_MAX!"));
        return -1;
    }

    return steppers_move_axis(stepper_e, E_MM_TO_PERCENT(pos) * E_STEPS_PER_PERCENT * E_AXIS_MOVEMENT_DIR);
}

float steppers_pos_x(void) {
    CLEAR_STORE_INTERRUPTS();
    float v = stepper_x.targetPosition();
    RESTORE_INTERRUPTS();
    return v / XY_STEPS_PER_MM / X_AXIS_MOVEMENT_DIR;
}

float steppers_pos_y(void) {
    CLEAR_STORE_INTERRUPTS();
    float v = stepper_y.targetPosition();
    RESTORE_INTERRUPTS();
    return v / XY_STEPS_PER_MM / Y_AXIS_MOVEMENT_DIR;
}

float steppers_pos_z(void) {
    CLEAR_STORE_INTERRUPTS();
    float v = stepper_z.targetPosition();
    RESTORE_INTERRUPTS();
    return v / Z_STEPS_PER_MM / Z_AXIS_MOVEMENT_DIR;
}

float steppers_pos_e(void) {
    CLEAR_STORE_INTERRUPTS();
    float v = stepper_e.targetPosition();
    RESTORE_INTERRUPTS();
    return E_PERCENT_TO_MM(v / E_STEPS_PER_PERCENT / E_AXIS_MOVEMENT_DIR);
}

void steppers_kill(void) {
    CLEAR_STORE_INTERRUPTS();

    stepper_x.setCurrentPosition(0);
    stepper_y.setCurrentPosition(0);
    stepper_z.setCurrentPosition(0);
    stepper_e.setCurrentPosition(0);

    stepper_x.disableOutputs();
    stepper_y.disableOutputs();
    stepper_z.disableOutputs();
    stepper_e.disableOutputs();

    state = step_not_homed;

    RESTORE_INTERRUPTS();
}

void steppers_set_speed_x(float speed) {
    if (speed < 0.0) {
        speed = 0.0;
    }
    if (speed > XY_MAX_SPEED) {
        speed = XY_MAX_SPEED;
    }

    CLEAR_STORE_INTERRUPTS();
    stepper_x.setMaxSpeed(speed * XY_STEPS_PER_MM);
    RESTORE_INTERRUPTS();
}

void steppers_set_speed_y(float speed) {
    if (speed < 0.0) {
        speed = 0.0;
    }
    if (speed > XY_MAX_SPEED) {
        speed = XY_MAX_SPEED;
    }

    CLEAR_STORE_INTERRUPTS();
    stepper_y.setMaxSpeed(speed * XY_STEPS_PER_MM);
    RESTORE_INTERRUPTS();
}

void steppers_set_speed_z(float speed) {
    if (speed < 0.0) {
        speed = 0.0;
    }
    if (speed > Z_MAX_SPEED) {
        speed = Z_MAX_SPEED;
    }

    CLEAR_STORE_INTERRUPTS();
    stepper_z.setMaxSpeed(speed * Z_STEPS_PER_MM);
    RESTORE_INTERRUPTS();
}

void steppers_set_speed_e(float speed) {
    if (speed < 0.0) {
        speed = 0.0;
    }
    if (speed > E_MAX_SPEED) {
        speed = E_MAX_SPEED;
    }

    CLEAR_STORE_INTERRUPTS();
    stepper_e.setMaxSpeed(speed * E_STEPS_PER_PERCENT);
    RESTORE_INTERRUPTS();
}

void steppers_set_accel_x(float accel) {
    if (accel < 0.0) {
        accel = 0.0;
    }
    if (accel > XY_MAX_ACCEL) {
        accel = XY_MAX_ACCEL;
    }

    CLEAR_STORE_INTERRUPTS();
    stepper_x.setAcceleration(accel * XY_STEPS_PER_MM);
    RESTORE_INTERRUPTS();
}

void steppers_set_accel_y(float accel) {
    if (accel < 0.0) {
        accel = 0.0;
    }
    if (accel > XY_MAX_ACCEL) {
        accel = XY_MAX_ACCEL;
    }

    CLEAR_STORE_INTERRUPTS();
    stepper_y.setAcceleration(accel * XY_STEPS_PER_MM);
    RESTORE_INTERRUPTS();
}

void steppers_set_accel_z(float accel) {
    if (accel < 0.0) {
        accel = 0.0;
    }
    if (accel > Z_MAX_ACCEL) {
        accel = Z_MAX_ACCEL;
    }

    CLEAR_STORE_INTERRUPTS();
    stepper_z.setAcceleration(accel * Z_STEPS_PER_MM);
    RESTORE_INTERRUPTS();
}

void steppers_set_accel_e(float accel) {
    if (accel < 0.0) {
        accel = 0.0;
    }
    if (accel > E_MAX_ACCEL) {
        accel = E_MAX_ACCEL;
    }

    CLEAR_STORE_INTERRUPTS();
    stepper_e.setAcceleration(accel * E_STEPS_PER_PERCENT);
    RESTORE_INTERRUPTS();
}