From 303f73433fdea0fa09b27e5b4eeeececac4d0771 Mon Sep 17 00:00:00 2001
From: Mathieu Broillet <mathieu@broillet.ch>
Date: Thu, 29 Aug 2024 12:50:18 +0200
Subject: [PATCH] improve start/update/uninstall

---
 core/stack.py   | 32 ++++++++++++++++++++++++++------
 main.py         |  2 +-
 ui/interface.py | 15 +++++++++------
 3 files changed, 36 insertions(+), 13 deletions(-)

diff --git a/core/stack.py b/core/stack.py
index 58ebacd..5ee8687 100644
--- a/core/stack.py
+++ b/core/stack.py
@@ -41,12 +41,15 @@ class Stack:
 
 
     def install(self):
-        self.check_for_broken_install()
-        self.create_venv()
-        self._install()
+        if self.is_installed():
+            self.update()
+        else:
+            self.check_for_broken_install()
+            self.create_venv()
+            self._install()
 
-        self.create_file('.installed', 'true')
-        logger.info(f"Installed {self.name}")
+            self.create_file('.installed', 'true')
+            logger.info(f"Installed {self.name}")
 
     def _install(self):
         pass
@@ -66,17 +69,33 @@ class Stack:
 
     def update(self, folder: str = 'webui'):
         if self.is_installed():
+            status = self.status()
+            if status:
+                self.stop()
+
             logger.info(f"Updating {self.name}")
             self.git_pull(folder)
+            self._update()
+
+            if status:
+                self.start()
         else:
             logger.warning(f"Could not update {self.name} as {self.name} is not installed")
             choices.any_key.ask()
 
+    def _update(self):
+        pass
+
     def uninstall(self):
         logger.info(f"Uninstalling {self.name}")
+        if self.status():
+            self.stop()
         self.bash(f"rm -rf {self.path}")
 
     def start(self):
+        if self.status():
+            logger.warning(f"{self.name} is already running")
+
         if self.is_installed():
             self._start()
         else:
@@ -90,6 +109,8 @@ class Stack:
         if self.status():
             logger.debug(f"Killing {self.name} with PID: {self.pid}")
             psutil.Process(self.pid).kill()
+        else:
+            logger.warning(f"{self.name} is not running")
 
         self.set_pid(None)
 
@@ -158,7 +179,6 @@ class Stack:
                     return
                 else:
                     # TODO: attach to subprocess / redirect logs?
-                    logger.info("Continuing without restarting...")
                     return
             else:
                 logger.debug(f"Running command as daemon: {cmd}")
diff --git a/main.py b/main.py
index 94e55a7..b77f94f 100644
--- a/main.py
+++ b/main.py
@@ -27,7 +27,7 @@ def load_services():
 
 
 if __name__ == '__main__':
-    setup_logger(logging.DEBUG)
+    setup_logger(logging.INFO)
     logger.info("Starting AI Suite for ROCM")
 
     setup_config()
diff --git a/ui/interface.py b/ui/interface.py
index 7bba706..0c716d5 100644
--- a/ui/interface.py
+++ b/ui/interface.py
@@ -19,22 +19,25 @@ def handle_services(action, service):
 
     service = loaded_services[service]
     if action == "start":
-        logger.info(f"Starting service: {service.name}")
         service.start()
     elif action == "stop":
-        logger.info(f"Stopping service: {service.name}")
         service.stop()
     elif action == "install":
         confirmation = choices.are_you_sure.ask()
         if confirmation:
-            logger.info(f"Installing/updating service: {service.name}")
             service.install()
     elif action == "uninstall":
         confirmation = choices.are_you_sure.ask()
         if confirmation:
-            type_confirmation = questionary.text(f"Please type {service.id} to confirm uninstallation:")
-            if type_confirmation.ask() == service.id:
-                logger.info(f"Uninstalling service: {service.name}")
+            type_confirmation = questionary.text(f"Please type {service.id} to confirm uninstallation (or type cancel):")
+
+            value = type_confirmation.ask()
+
+            if value == "cancel":
+                questionary.print("Canceled", style="fg:ansired")
+            elif value != service.id:
+                questionary.print("Invalid input, please try again", style="fg:ansired")
+            elif value == service.id:
                 service.uninstall()
 
     choices.any_key.ask()