From a4eb54ae6c218ba6359966982a85b5560a7ff29b Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Daniel=20=C4=8Cejchan?= <daniel.cejchan@prusa3d.cz>
Date: Thu, 5 Sep 2024 14:00:04 +0200
Subject: [PATCH] Belt tuning: Add gantry align phase

Add a new phase as instructed by the research team.

BFW-6029
---
 .../belt_tuning/belt_tuning_wizard.cpp        | 24 ++++++++++++++++++-
 .../marlin_server_types/client_response.hpp   |  4 ++++
 src/gui/wizard/screen_belt_tuning_wizard.cpp  |  3 +++
 3 files changed, 30 insertions(+), 1 deletion(-)

diff --git a/src/common/feature/belt_tuning/belt_tuning_wizard.cpp b/src/common/feature/belt_tuning/belt_tuning_wizard.cpp
index bea27be1a6..1a20c4ecb0 100644
--- a/src/common/feature/belt_tuning/belt_tuning_wizard.cpp
+++ b/src/common/feature/belt_tuning/belt_tuning_wizard.cpp
@@ -1,5 +1,6 @@
 #include "belt_tuning_wizard.hpp"
 
+#include <module/stepper.h>
 #include <module/prusa/accelerometer.h>
 #include <Marlin/src/gcode/gcode.h>
 #include <gcode/calibrate/M958.hpp>
@@ -31,7 +32,27 @@ class FSMBeltTuning final {
 
 private:
     PhaseOpt phase_init(const Meta::LoopCallbackArgs &) {
-        return Phase::preparing;
+        return Phase::ask_for_gantry_align;
+    }
+
+    void phase_ask_for_gantry_align_init(const Meta::InitCallbackArgs &) {
+        // Disable the XY gantry - the user is prompted to move it by hand in this phase
+        planner.synchronize();
+        disable_XY();
+    }
+
+    PhaseOpt phase_ask_for_gantry_align(const Meta::LoopCallbackArgs &args) {
+        switch (args.response.value_or<Response>(Response::_none)) {
+
+        case Response::Done:
+            return Phase::preparing;
+
+        case Response::Abort:
+            return Phase::finish;
+
+        default:
+            return std::nullopt;
+        }
     }
 
     PhaseOpt phase_preparing(const Meta::LoopCallbackArgs &) {
@@ -195,6 +216,7 @@ class FSMBeltTuning final {
     using C = FSMBeltTuning;
     static constexpr Config config {
         { Phase::init, { &C::phase_init } },
+        { Phase::ask_for_gantry_align, { &C::phase_ask_for_gantry_align, &C::phase_ask_for_gantry_align_init } },
         { Phase::preparing, { &C::phase_preparing } },
         { Phase::ask_for_dampeners_installation, { &C::phase_ask_for_dampeners_installation } },
         { Phase::calibrating_accelerometer, { &C::phase_calibrating_accelerometer } },
diff --git a/src/common/marlin_server_types/client_response.hpp b/src/common/marlin_server_types/client_response.hpp
index 2a5fa4717e..79d7a6f58b 100644
--- a/src/common/marlin_server_types/client_response.hpp
+++ b/src/common/marlin_server_types/client_response.hpp
@@ -499,6 +499,9 @@ constexpr inline ClientFSM client_fsm_from_phase(PhasesSerialPrinting) { return
 enum class PhaseBeltTuning : PhaseUnderlyingType {
     init,
 
+    /// The user is prompted to loose the belts so that the gantry is aligned
+    ask_for_gantry_align,
+
     /// Homing, selecting a proper tool and moving it to the measuring position
     preparing,
 
@@ -880,6 +883,7 @@ class ClientResponses {
 #if HAS_BELT_TUNING()
     static constexpr EnumArray<PhaseBeltTuning, PhaseResponses, CountPhases<PhaseBeltTuning>()> belt_tuning_responses {
         { PhaseBeltTuning::init, {} },
+        { PhaseBeltTuning::ask_for_gantry_align, { Response::Done, Response::Abort } },
         { PhaseBeltTuning::preparing, { Response::Abort } },
         { PhaseBeltTuning::ask_for_dampeners_installation, { Response::Done, Response::Abort } },
         { PhaseBeltTuning::calibrating_accelerometer, { Response::Abort } },
diff --git a/src/gui/wizard/screen_belt_tuning_wizard.cpp b/src/gui/wizard/screen_belt_tuning_wizard.cpp
index ac90e55003..1e15eae841 100644
--- a/src/gui/wizard/screen_belt_tuning_wizard.cpp
+++ b/src/gui/wizard/screen_belt_tuning_wizard.cpp
@@ -8,6 +8,7 @@
 #include <display.hpp>
 
 #include <gui/standard_frame/frame_prompt.hpp>
+#include <gui/standard_frame/frame_qr_prompt.hpp>
 #include <feature/belt_tuning/belt_tuning_wizard.hpp>
 #include <feature/belt_tuning/printer_belt_parameters.hpp>
 
@@ -15,6 +16,7 @@ using Phase = PhaseBeltTuning;
 
 namespace {
 
+using FrameAskForGantryAlign = WithConstructorArgs<FrameQRPrompt, Phase::ask_for_gantry_align, N_("Please follow the XY gantry alignment process in the manual."), "belt-tuning"_tstr>;
 using FramePreparing = WithConstructorArgs<FramePrompt, Phase::preparing, N_("Preparing"), N_("Setting the printer up for the calibration.\n\nPlease wait.")>;
 
 // TODO: Maybe add progress bar?
@@ -189,6 +191,7 @@ class FrameResults : public FramePrompt {
 };
 
 using Frames = FrameDefinitionList<ScreenBeltTuningWizard::FrameStorage,
+    FrameDefinition<Phase::ask_for_gantry_align, FrameAskForGantryAlign>,
     FrameDefinition<Phase::preparing, FramePreparing>,
     FrameDefinition<Phase::ask_for_dampeners_installation, FrameAskForDampenersInstallation>,
     FrameDefinition<Phase::calibrating_accelerometer, FrameCalibratingAccelerometer>,