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

#include "config.h"
#include "config_pins.h"
#include "lcd.h"
#include "data.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 stepper_states state = step_disabled;
static 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.setMaxSpeed(XY_MAX_SPEED * XY_STEPS_PER_MM);
    steppers_set_speed_x(data_options()->speed_x);
    steppers_set_accel_x(data_options()->accel_x);

    stepper_y.setEnablePin(Y_ENABLE_PIN);
    steppers_set_speed_y(data_options()->speed_y);
    steppers_set_accel_y(data_options()->accel_y);

    stepper_z.setEnablePin(Z_ENABLE_PIN);
    steppers_set_speed_z(data_options()->speed_z);
    steppers_set_accel_z(data_options()->accel_z);

    stepper_e.setEnablePin(E0_ENABLE_PIN);
    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) {
    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"));
    }
}

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

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);
            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);
            steppers_initiate_home(3, 0);
        } 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);
            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);
            state = step_homed;
            return false;
        } else {
            stepper_e.runSpeed();
        }
    } else if (state != step_disabled) {
        for (int i = 0; i < 4; i++) {
            if (steppers_home_switch(i)) {
                Serial.print(F("ERROR: endstop hit on "));
                Serial.println(i);

                // TODO proper error handling?
                lcd_clear();
                lcd_set_heading("ERROR");
                lcd_set_text("Endstop hit. Aborting!");
                while (1) { }
            }
        }
        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) {
    return (state == step_homed);
}

void steppers_start_homing(void) {
    steppers_initiate_home(2, 0);
}

static int steppers_move_axis(AccelStepper &axis, long pos) {
    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);
    return 0;
}

int steppers_move_x(long 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(long 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(long 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(long 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);
}

void steppers_kill(void) {
    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;
}

void steppers_set_speed_x(float speed) {
    if (speed < 0.0) {
        speed = 0.0;
    }
    if (speed > XY_MAX_SPEED) {
        speed = XY_MAX_SPEED;
    }
    stepper_x.setMaxSpeed(speed * XY_STEPS_PER_MM);
}

void steppers_set_speed_y(float speed) {
    if (speed < 0.0) {
        speed = 0.0;
    }
    if (speed > XY_MAX_SPEED) {
        speed = XY_MAX_SPEED;
    }
    stepper_y.setMaxSpeed(speed * XY_STEPS_PER_MM);
}

void steppers_set_speed_z(float speed) {
    if (speed < 0.0) {
        speed = 0.0;
    }
    if (speed > Z_MAX_SPEED) {
        speed = Z_MAX_SPEED;
    }
    stepper_z.setMaxSpeed(speed * Z_STEPS_PER_MM);
}

void steppers_set_speed_e(float speed) {
    if (speed < 0.0) {
        speed = 0.0;
    }
    if (speed > E_MAX_SPEED) {
        speed = E_MAX_SPEED;
    }
    stepper_e.setMaxSpeed(speed * E_STEPS_PER_PERCENT);
}

void steppers_set_accel_x(float accel) {
    if (accel < 0.0) {
        accel = 0.0;
    }
    if (accel > XY_MAX_ACCEL) {
        accel = XY_MAX_ACCEL;
    }
    stepper_x.setAcceleration(accel * XY_STEPS_PER_MM);
}

void steppers_set_accel_y(float accel) {
    if (accel < 0.0) {
        accel = 0.0;
    }
    if (accel > XY_MAX_ACCEL) {
        accel = XY_MAX_ACCEL;
    }
    stepper_y.setAcceleration(accel * XY_STEPS_PER_MM);
}

void steppers_set_accel_z(float accel) {
    if (accel < 0.0) {
        accel = 0.0;
    }
    if (accel > Z_MAX_ACCEL) {
        accel = Z_MAX_ACCEL;
    }
    stepper_z.setAcceleration(accel * Z_STEPS_PER_MM);
}

void steppers_set_accel_e(float accel) {
    if (accel < 0.0) {
        accel = 0.0;
    }
    if (accel > E_MAX_ACCEL) {
        accel = E_MAX_ACCEL;
    }
    stepper_e.setAcceleration(accel * E_STEPS_PER_PERCENT);
}