#include <Arduino.h>

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

//#define DISABLE_HOMING_STATE
#define ENABLE_MOVEMENT_TEST_STATE
#define ENABLE_INPUT_TEST_STATE

// --------------------------------------

StateText sm_ask_homing = StateText();
StateText sm_do_homing = StateText(&sm_ask_homing);
StateMenu sm_menu = StateMenu(&sm_do_homing, false);

StateMenu sm_auto = StateMenu(&sm_menu);
StateDynamicMenu sm_presets = StateDynamicMenu(&sm_auto);
// TODO dispense

StateText sm_new_preset = StateText(&sm_auto);
StateText sm_wiz_move = StateText(&sm_new_preset);
StateValues<uint8_t, 2> sm_wiz_count = StateValues<uint8_t, 2>(&sm_wiz_move);
StateValues<float, 2> sm_wiz_first_pos = StateValues<float, 2>(&sm_wiz_count);
StateValues<float, 2> sm_wiz_last_pos = StateValues<float, 2>(&sm_wiz_first_pos);
// TODO move z above container
// TODO move z inside container
// TODO move to extruder depth
StateText sm_new_preset_done = StateText(&sm_wiz_last_pos);

StateDynamicMenu sm_mod_preset = StateDynamicMenu(&sm_auto);
// TODO modify

StateDynamicMenu sm_del_preset = StateDynamicMenu(&sm_auto);
// TODO delete

float move_pos_x = 0.0, move_pos_y = 0.0, move_pos_z = 0.0, move_pos_e = 0.0;
StateMenu sm_move = StateMenu(&sm_menu);
StateValue<float> sm_move_x = StateValue<float>(&sm_move, move_pos_x, X_AXIS_MIN, X_AXIS_MAX);
StateValue<float> sm_move_y = StateValue<float>(&sm_move, move_pos_y, Y_AXIS_MIN, Y_AXIS_MAX);
StateValue<float> sm_move_z = StateValue<float>(&sm_move, move_pos_z, Z_AXIS_MIN, Z_AXIS_MAX);
StateValue<float> sm_move_e = StateValue<float>(&sm_move, move_pos_e, E_AXIS_MIN, E_AXIS_MAX);

StateMenu sm_config = StateMenu(&sm_menu);
StateText sm_conf_load = StateText(&sm_config);
StateText sm_conf_save = StateText(&sm_config);
StateValue<float> sm_conf_speed_xy = StateValue<float>(&sm_config, data_options()->speed_x, 1.0, XY_MAX_SPEED);
StateValue<float> sm_conf_speed_z = StateValue<float>(&sm_config, data_options()->speed_z, 1.0, Z_MAX_SPEED);
StateValue<float> sm_conf_speed_e = StateValue<float>(&sm_config, data_options()->speed_e, 1.0, E_MAX_SPEED);
StateValue<float> sm_conf_accel_xy = StateValue<float>(&sm_config, data_options()->accel_x, 1.0, XY_MAX_ACCEL);
StateValue<float> sm_conf_accel_z = StateValue<float>(&sm_config, data_options()->accel_z, 1.0, Z_MAX_ACCEL);
StateValue<float> sm_conf_accel_e = StateValue<float>(&sm_config, data_options()->accel_e, 1.0, E_MAX_ACCEL);

StateText sm_error = StateText(NULL);

#ifdef ENABLE_MOVEMENT_TEST_STATE
StateMenu sm_movement_test = StateMenu(&sm_menu);
StateText sm_move_x_test = StateText(&sm_movement_test);
StateText sm_move_y_test = StateText(&sm_movement_test);
#endif // ENABLE_MOVEMENT_TEST_STATE

#ifdef ENABLE_INPUT_TEST_STATE
StateText sm_input_test = StateText(&sm_menu);
#endif // ENABLE_INPUT_TEST_STATE

// --------------------------------------

State *current_state = NULL;
String strbuf;
struct data_config_preset wizard_preset = { .count_x = 1, .count_y = 1 };
float wizard_first_x = 0.0, wizard_first_y = 0.0;
float wizard_last_x = 0.0, wizard_last_y = 0.0;

void states_init(void) {
    // ----------------------------------

    sm_ask_homing.setTitle("Ask for homing");
    sm_ask_homing.setHeading("Homing Required!");
    sm_ask_homing.setText("Click to home all four axes.");

    // ----------------------------------

    sm_do_homing.setTitle("Home all axes");
    sm_do_homing.setHeading("Homing...");

    sm_do_homing.onEnter([]() {
#ifndef DISABLE_HOMING_STATE
        steppers_start_homing();
#endif // DISABLE_HOMING_STATE
    });

    sm_do_homing.whenIn([](StateMachineInput smi) {
#ifndef DISABLE_HOMING_STATE
        if (smi.motors_done) {
            if (steppers_homed()) {
#endif // DISABLE_HOMING_STATE

                async_beep(HOMING_BEEP_TIME, HOMING_BEEP_FREQ);
                states_go_to(&sm_menu);

#ifndef DISABLE_HOMING_STATE
            }
        }
#endif // DISABLE_HOMING_STATE

        static char last_axis = ' ';
        char axis = steppers_homing_axis();
        if ((axis != ' ') && (axis != last_axis)) {
            last_axis = axis;
            strbuf = String(F("Currently homing ")) + axis + String(F(" axis"));
            sm_do_homing.setText(strbuf.c_str());
            sm_do_homing.updateText();
        }
    });

    // ----------------------------------

    sm_menu.setTitle("Main Menu");
    sm_menu.addChild(&sm_do_homing, 1);

    // ----------------------------------

    sm_move.setTitle("Move Axis");

    sm_move_x.setTitle("X Axis");
    sm_move_x.onEnter([]() {
        move_pos_x = steppers_pos_x();
    });
    sm_move_x.onLiveUpdate([](float v) {
        steppers_move_x(v);
    });

    sm_move_y.setTitle("Y Axis");
    sm_move_y.onEnter([]() {
        move_pos_y = steppers_pos_y();
    });
    sm_move_y.onLiveUpdate([](float v) {
        steppers_move_y(v);
    });

    sm_move_z.setTitle("Z Axis");
    sm_move_z.onEnter([]() {
        move_pos_z = steppers_pos_z();
    });
    sm_move_z.onLiveUpdate([](float v) {
        steppers_move_z(v);
    });

    sm_move_e.setTitle("E Axis");
    sm_move_e.onEnter([]() {
        move_pos_e = steppers_pos_e();
    });
    sm_move_e.onLiveUpdate([](float v) {
        steppers_move_e(v);
    });

    // ----------------------------------

    sm_auto.setTitle("Filling Menu");

    // TODO
    sm_presets.setTitle("Use Preset");
    sm_presets.setPrefix("Preset ");
    sm_presets.dataCount([]() {
        return (int)data_preset_count();
    });

    sm_new_preset.setTitle("Add new Preset");
    sm_new_preset.setHeading("Preset Wizard");
    sm_new_preset.setText("Moving axes into initial positions. Click to continue.");

    sm_wiz_move.setTitle("Initial Movement");
    sm_wiz_move.onEnter([]() {
        steppers_start_homing();
    });
    sm_wiz_move.whenIn([](StateMachineInput smi) {
#ifndef DISABLE_HOMING_STATE
        if (smi.motors_done) {
            if (steppers_homed()) {
#endif // DISABLE_HOMING_STATE
                states_go_to(&sm_wiz_count);
#ifndef DISABLE_HOMING_STATE
            }
        }
#endif // DISABLE_HOMING_STATE

        static char last_axis = ' ';
        char axis = steppers_homing_axis();
        if ((axis != ' ') && (axis != last_axis)) {
            last_axis = axis;
            strbuf = String(F("Currently homing ")) + axis + String(F(" axis"));
            sm_wiz_move.setText(strbuf.c_str());
            sm_wiz_move.updateText();
        }
    });

    sm_wiz_count.setTitle("Container Count");
    sm_wiz_count.setData(0, "X: ", &wizard_preset.count_x, 1, 50);
    sm_wiz_count.setData(1, "Y: ", &wizard_preset.count_y, 1, 100);

    sm_wiz_first_pos.setTitle("First Container Position");
    sm_wiz_first_pos.setData(0, "X: ", &wizard_first_x, X_AXIS_MIN, X_AXIS_MAX);
    sm_wiz_first_pos.setData(1, "Y: ", &wizard_first_y, Y_AXIS_MIN, Y_AXIS_MAX);
    sm_wiz_first_pos.onEnter([]() {
        wizard_first_x = steppers_current_pos(0);
        wizard_first_y = steppers_current_pos(1);
    });
    sm_wiz_first_pos.onLiveUpdate([](size_t i, float v) {
        if (i == 0) {
            steppers_move_x(v);
        } else {
            steppers_move_y(v);
        }
    });
    sm_wiz_first_pos.onUpdate([](size_t i, float v) {
        if (i == 0) {
            wizard_preset.offset_x = v;
        } else {
            wizard_preset.offset_y = v;
        }
    });

    sm_wiz_last_pos.setTitle("Last Container Position");
    sm_wiz_last_pos.setData(0, "X: ", &wizard_last_x, X_AXIS_MIN, X_AXIS_MAX);
    sm_wiz_last_pos.setData(1, "Y: ", &wizard_last_y, Y_AXIS_MIN, Y_AXIS_MAX);
    sm_wiz_last_pos.onEnter([]() {
        wizard_last_x = steppers_current_pos(0);
        wizard_last_y = steppers_current_pos(1);
    });
    sm_wiz_last_pos.onLiveUpdate([](size_t i, float v) {
        if (i == 0) {
            steppers_move_x(v);
        } else {
            steppers_move_y(v);
        }
    });
    sm_wiz_last_pos.onUpdate([](size_t i, float v) {
        if (i == 0) {
            wizard_preset.distance_x = (v - wizard_preset.offset_x) / wizard_preset.count_x;
        } else {
            wizard_preset.distance_y = (v - wizard_preset.offset_y) / wizard_preset.count_y;
        }
    });

    sm_new_preset_done.setTitle("Preset Wizard Done");
    sm_new_preset_done.setHeading("Preset Wizard Done");
    sm_new_preset_done.onEnter([]() {
        if (!data_preset_add(wizard_preset)) {
            strbuf = String(data_eeprom_error()) + F(" Error adding preset!");
            sm_new_preset_done.setText(strbuf.c_str());
        } else {
            strbuf = F("Preset ");
            strbuf += String(data_preset_count());
            strbuf += F(" has been added! Don't forget to store config to EEPROM to keep it.");
            sm_new_preset_done.setText(strbuf.c_str());
        }
    });
    sm_new_preset_done.whenIn([](StateMachineInput smi) {
        if (smi.click) {
            states_go_to(&sm_auto);
        }
    });

    // TODO
    sm_mod_preset.setTitle("Modify Preset");
    sm_mod_preset.setPrefix("Preset ");
    sm_mod_preset.dataCount([]() {
        return (int)data_preset_count();
    });

    // TODO
    sm_del_preset.setTitle("Delete Preset");
    sm_del_preset.setPrefix("Preset ");
    sm_del_preset.dataCount([]() {
        return (int)data_preset_count();
    });

    // ----------------------------------

    sm_config.setTitle("Configuration");

    sm_conf_load.setTitle("Load from EEPROM");
    sm_conf_load.setHeading("Load from EEPROM");
    sm_conf_load.onEnter([]() {
        if (!data_eeprom_read()) {
            strbuf = String(data_eeprom_error()) + " Error reading configuration!";
            sm_conf_load.setText(strbuf.c_str());
        } else {
            sm_conf_load.setText("Configuration read successful!");

            steppers_set_speed_x(data_options()->speed_x);
            steppers_set_accel_x(data_options()->accel_x);
            steppers_set_speed_y(data_options()->speed_y);
            steppers_set_accel_y(data_options()->accel_y);
            steppers_set_speed_z(data_options()->speed_z);
            steppers_set_accel_z(data_options()->accel_z);
            steppers_set_speed_e(data_options()->speed_e);
            steppers_set_accel_e(data_options()->accel_e);
        }
    });

    sm_conf_save.setTitle("Store to EEPROM");
    sm_conf_save.setHeading("Store to EEPROM");
    sm_conf_save.setText("Configuration written to EEPROM!");
    sm_conf_save.onEnter([]() {
        data_eeprom_write();
    });

    sm_conf_speed_xy.setTitle("XY Speed");
    sm_conf_speed_xy.onUpdate([](float v) {
        steppers_set_speed_x(v);
        steppers_set_speed_y(v);
    });

    sm_conf_speed_z.setTitle("Z Speed");
    sm_conf_speed_z.onUpdate([](float v) {
        steppers_set_speed_z(v);
    });

    sm_conf_speed_e.setTitle("E Speed");
    sm_conf_speed_e.onUpdate([](float v) {
        steppers_set_speed_e(v);
    });

    sm_conf_accel_xy.setTitle("XY Acceleration");
    sm_conf_accel_xy.onUpdate([](float v) {
        steppers_set_accel_x(v);
        steppers_set_accel_y(v);
    });

    sm_conf_accel_z.setTitle("Z Acceleration");
    sm_conf_accel_z.onUpdate([](float v) {
        steppers_set_accel_z(v);
    });

    sm_conf_accel_e.setTitle("E Acceleration");
    sm_conf_accel_e.onUpdate([](float v) {
        steppers_set_accel_e(v);
    });

    sm_error.setChild(&sm_ask_homing);
    sm_error.setTitle("Axis Error");
    sm_error.setHeading("Axis Error");
    sm_error.setText("Endstop has been hit!");

#ifdef ENABLE_MOVEMENT_TEST_STATE
    sm_movement_test.setTitle("Movement Test");

    sm_move_x_test.setTitle("X Axis");
    sm_move_x_test.setHeading("X Move Test");
    sm_move_x_test.whenIn([](StateMachineInput smi) {
        static bool s = false;
        if (smi.motors_done) {
            s = !s;
            steppers_move_x(s ? (X_AXIS_MAX - 10.0) : (X_AXIS_MIN + 10.0));
            if (smi.click) {
                states_go_to(&sm_movement_test);
            }
        }
    });

    sm_move_y_test.setTitle("Y Axis");
    sm_move_y_test.setHeading("Y Move Test");
    sm_move_y_test.whenIn([](StateMachineInput smi) {
        static bool s = false;
        if (smi.motors_done) {
            s = !s;
            steppers_move_y(s ? (Y_AXIS_MAX - 10.0) : (Y_AXIS_MIN + 10.0));
            if (smi.click) {
                states_go_to(&sm_movement_test);
            }
        }
    });
#endif // ENABLE_MOVEMENT_TEST_STATE

#ifdef ENABLE_INPUT_TEST_STATE
    sm_input_test.setTitle("Input Test");
    sm_input_test.setHeading("Input Test");

    sm_input_test.whenIn([](StateMachineInput smi) {
        String s = "Endstops: ";
        for (int i = 0; i < 4; i++) {
            if (steppers_home_switch(i)) {
                s += '1';
            } else {
                s += '0';
            }
            if (i < 3) {
                s += ' ';
            }
        }
        if (s != strbuf) {
            strbuf = s;
            sm_input_test.setText(strbuf.c_str());
            sm_input_test.updateText();
        }
    });
#endif // ENABLE_INPUT_TEST_STATE

    // ----------------------------------

    states_go_to(&sm_ask_homing);
}

void states_run(StateMachineInput smi) {
    if (smi.click) {
        async_beep(ENCODER_CLICK_BEEP_TIME, ENCODER_CLICK_BEEP_FREQ);
    }

    if (smi.kill) {
        steppers_kill();
        Serial.println(F("Kill button pressed!"));
        blocking_beep(KILL_BEEP_TIME, KILL_BEEP_FREQ, KILL_BEEP_REPEAT - 1);
        states_go_to(&sm_ask_homing);
        return;
    }

    if (current_state != NULL) {
        current_state->inState(smi);
    }
}

State *states_get(void) {
    return current_state;
}

void states_go_to(State *state) {
    current_state = state;

    if (current_state != NULL) {
        current_state->enterState();
    }
}