From 515a42f0981009217ef581fdda8027cca221e6ca Mon Sep 17 00:00:00 2001
From: dw-0 <th33xitus@gmail.com>
Date: Sun, 5 Nov 2023 16:15:19 +0100
Subject: [PATCH] feat(klipper): implement update function

Signed-off-by: Dominik Willner <th33xitus@gmail.com>
---
 kiauh/instance_manager/instance_manager.py | 10 +++++++++
 kiauh/menus/update_menu.py                 |  3 ++-
 kiauh/modules/klipper/klipper_dialogs.py   | 17 ++++++++++++++
 kiauh/modules/klipper/klipper_setup.py     | 26 ++++++++++++++++++++++
 kiauh/repo_manager/repo_manager.py         | 18 +++++++++++++++
 5 files changed, 73 insertions(+), 1 deletion(-)

diff --git a/kiauh/instance_manager/instance_manager.py b/kiauh/instance_manager/instance_manager.py
index f08f43e..5a775c1 100644
--- a/kiauh/instance_manager/instance_manager.py
+++ b/kiauh/instance_manager/instance_manager.py
@@ -91,6 +91,11 @@ class InstanceManager:
             Logger.print_error(f"Error starting service {self.instance_name}.service:")
             Logger.print_error(f"{e}")
 
+    def start_all_instance(self) -> None:
+        for instance in self.instances:
+            self.set_current_instance(instance)
+            self.start_instance()
+
     def stop_instance(self) -> None:
         Logger.print_info(f"Stopping {self.instance_name}.service ...")
         try:
@@ -102,6 +107,11 @@ class InstanceManager:
             Logger.print_error(f"{e}")
             raise
 
+    def stop_all_instance(self) -> None:
+        for instance in self.instances:
+            self.set_current_instance(instance)
+            self.stop_instance()
+
     def reload_daemon(self) -> None:
         Logger.print_info("Reloading systemd manager configuration ...")
         try:
diff --git a/kiauh/menus/update_menu.py b/kiauh/menus/update_menu.py
index 0e9a78a..f617fe9 100644
--- a/kiauh/menus/update_menu.py
+++ b/kiauh/menus/update_menu.py
@@ -12,6 +12,7 @@
 import textwrap
 
 from kiauh.menus.base_menu import BaseMenu
+from kiauh.modules.klipper.klipper_setup import update_klipper
 from kiauh.utils.constants import COLOR_GREEN, RESET_FORMAT
 
 
@@ -74,7 +75,7 @@ class UpdateMenu(BaseMenu):
         print("update_all")
 
     def update_klipper(self):
-        print("update_klipper")
+        update_klipper()
 
     def update_moonraker(self):
         print("update_moonraker")
diff --git a/kiauh/modules/klipper/klipper_dialogs.py b/kiauh/modules/klipper/klipper_dialogs.py
index 204ea43..cc8d15c 100644
--- a/kiauh/modules/klipper/klipper_dialogs.py
+++ b/kiauh/modules/klipper/klipper_dialogs.py
@@ -96,3 +96,20 @@ def print_missing_usergroup_dialog(missing_groups) -> None:
         f"| {COLOR_YELLOW}Relog required for group assignments to take effect!{RESET_FORMAT}  |"
     )
     print("\\=======================================================/")
+
+
+def print_update_warn_dialog():
+    print("/=======================================================\\")
+    print(
+        f"| {COLOR_YELLOW}WARNING:                                              {RESET_FORMAT}|"
+    )
+    print(
+        f"| {COLOR_YELLOW}Do NOT continue if there are ongoing prints running!  {RESET_FORMAT}|"
+    )
+    print(
+        f"| {COLOR_YELLOW}All Klipper instances will be restarted during the    {RESET_FORMAT}|"
+    )
+    print(
+        f"| {COLOR_YELLOW}update process and any ongoing prints WILL FAIL.      {RESET_FORMAT}|"
+    )
+    print("\\=======================================================/")
diff --git a/kiauh/modules/klipper/klipper_setup.py b/kiauh/modules/klipper/klipper_setup.py
index 7b343b8..d38b89f 100644
--- a/kiauh/modules/klipper/klipper_setup.py
+++ b/kiauh/modules/klipper/klipper_setup.py
@@ -27,6 +27,7 @@ from kiauh.modules.klipper.klipper import Klipper
 from kiauh.modules.klipper.klipper_dialogs import (
     print_instance_overview,
     print_select_instance_count_dialog,
+    print_update_warn_dialog,
 )
 from kiauh.modules.klipper.klipper_utils import (
     handle_convert_single_to_multi_instance_names,
@@ -249,3 +250,28 @@ def remove_multi_instance(instance_manager: InstanceManager) -> None:
         instance_manager.delete_instance(del_remnants=False)
 
     instance_manager.reload_daemon()
+
+
+def update_klipper() -> None:
+    print_update_warn_dialog()
+
+    if not get_confirm("Update Klipper now?"):
+        return
+
+    instance_manager = InstanceManager(Klipper)
+    instance_manager.get_instances()
+    instance_manager.stop_all_instance()
+
+    cm = ConfigManager()
+    cm.read_config()
+
+    repo = str(cm.get_value("klipper", "repository_url") or DEFAULT_KLIPPER_REPO_URL)
+    branch = str(cm.get_value("klipper", "branch") or "master")
+
+    repo_manager = RepoManager(
+        repo=repo,
+        branch=branch,
+        target_dir=KLIPPER_DIR,
+    )
+    repo_manager.pull_repo()
+    instance_manager.start_all_instance()
diff --git a/kiauh/repo_manager/repo_manager.py b/kiauh/repo_manager/repo_manager.py
index b1501d2..4555920 100644
--- a/kiauh/repo_manager/repo_manager.py
+++ b/kiauh/repo_manager/repo_manager.py
@@ -77,6 +77,15 @@ class RepoManager:
             Logger.print_error(f"Error removing existing repository: {e.strerror}")
             return
 
+    def pull_repo(self) -> None:
+        Logger.print_info(f"Updating repository '{self.repo}' ...")
+        try:
+            self._pull()
+        except subprocess.CalledProcessError:
+            log = "An unexpected error occured during updating the repository."
+            Logger.print_error(log)
+            return
+
     def _clone(self):
         try:
             command = ["git", "clone", self.repo, self.target_dir]
@@ -99,5 +108,14 @@ class RepoManager:
             Logger.print_error(log)
             raise
 
+    def _pull(self) -> None:
+        try:
+            command = ["git", "pull"]
+            subprocess.run(command, cwd=self.target_dir, check=True)
+        except subprocess.CalledProcessError as e:
+            log = f"Error on git pull: {e.stderr.decode()}"
+            Logger.print_error(log)
+            raise
+
     def _get_method(self) -> str:
         return "ssh" if self.repo.startswith("git") else "https"