From b020f1096768983c4fcfd132e4f30a727608a188 Mon Sep 17 00:00:00 2001
From: dw-0 <th33xitus@gmail.com>
Date: Thu, 18 Apr 2024 21:54:34 +0200
Subject: [PATCH] feat: implement repo rollback feature

Signed-off-by: Dominik Willner <th33xitus@gmail.com>
---
 kiauh/core/menus/advanced_menu.py | 14 ++++++++++++
 kiauh/utils/git_utils.py          | 36 ++++++++++++++++++++++++++++++-
 2 files changed, 49 insertions(+), 1 deletion(-)

diff --git a/kiauh/core/menus/advanced_menu.py b/kiauh/core/menus/advanced_menu.py
index f86788e..fb4a3ea 100644
--- a/kiauh/core/menus/advanced_menu.py
+++ b/kiauh/core/menus/advanced_menu.py
@@ -10,6 +10,8 @@
 import textwrap
 from typing import Type, Optional
 
+from components.klipper import KLIPPER_DIR
+from components.klipper.klipper import Klipper
 from components.klipper_firmware.menus.klipper_build_menu import (
     KlipperBuildFirmwareMenu,
 )
@@ -17,12 +19,16 @@ from components.klipper_firmware.menus.klipper_flash_menu import (
     KlipperFlashMethodMenu,
     KlipperSelectMcuConnectionMenu,
 )
+from components.moonraker import MOONRAKER_DIR
+from components.moonraker.moonraker import Moonraker
 from core.menus import Option
 from core.menus.base_menu import BaseMenu
 from utils.constants import COLOR_YELLOW, RESET_FORMAT
+from utils.git_utils import rollback_repository
 
 
 # noinspection PyUnusedLocal
+# noinspection PyMethodMayBeStatic
 class AdvancedMenu(BaseMenu):
     def __init__(self, previous_menu: Optional[Type[BaseMenu]] = None):
         super().__init__()
@@ -37,6 +43,8 @@ class AdvancedMenu(BaseMenu):
 
     def set_options(self):
         self.options = {
+            "1": Option(method=self.klipper_rollback, menu=True),
+            "2": Option(method=self.moonraker_rollback, menu=True),
             "3": Option(method=self.build, menu=True),
             "4": Option(method=self.flash, menu=False),
             "5": Option(method=self.build_flash, menu=False),
@@ -65,6 +73,12 @@ class AdvancedMenu(BaseMenu):
         )[1:]
         print(menu, end="")
 
+    def klipper_rollback(self, **kwargs):
+        rollback_repository(KLIPPER_DIR, Klipper)
+
+    def moonraker_rollback(self, **kwargs):
+        rollback_repository(MOONRAKER_DIR, Moonraker)
+
     def build(self, **kwargs):
         KlipperBuildFirmwareMenu(previous_menu=self.__class__).run()
 
diff --git a/kiauh/utils/git_utils.py b/kiauh/utils/git_utils.py
index 2ed9a4d..4feeaf3 100644
--- a/kiauh/utils/git_utils.py
+++ b/kiauh/utils/git_utils.py
@@ -2,8 +2,12 @@ import json
 import urllib.request
 from http.client import HTTPResponse
 from json import JSONDecodeError
-from typing import List
+from subprocess import CalledProcessError, PIPE, run
+from typing import List, Type
 
+from core.instance_manager.base_instance import BaseInstance
+from core.instance_manager.instance_manager import InstanceManager
+from utils.input_utils import get_number_input, get_confirm
 from utils.logger import Logger
 
 
@@ -55,3 +59,33 @@ def get_latest_unstable_tag(repo_path: str) -> str:
     except Exception:
         Logger.print_error("Error while getting the latest unstable tag")
         raise
+
+
+def rollback_repository(repo_dir: str, instance: Type[BaseInstance]) -> None:
+    q1 = "How many commits do you want to roll back"
+    amount = get_number_input(q1, 1, allow_go_back=True)
+
+    im = InstanceManager(instance)
+
+    Logger.print_warn("Do not continue if you have ongoing prints!", start="\n")
+    Logger.print_warn(
+        f"All currently running {im.instance_type.__name__} services will be stopped!"
+    )
+    if not get_confirm(
+        f"Roll back {amount} commit{'s' if amount > 1 else ''}",
+        default_choice=False,
+        allow_go_back=True,
+    ):
+        Logger.print_info("Aborting roll back ...")
+        return
+
+    im.stop_all_instance()
+
+    try:
+        cmd = ["git", "reset", "--hard", f"HEAD~{amount}"]
+        run(cmd, cwd=repo_dir, check=True, stdout=PIPE, stderr=PIPE)
+        Logger.print_ok(f"Rolled back {amount} commits!", start="\n")
+    except CalledProcessError as e:
+        Logger.print_error(f"An error occured during repo rollback:\n{e}")
+
+    im.start_all_instance()