From 29d02346accbe12cffbb42d75253f5e7f7bdfa1f Mon Sep 17 00:00:00 2001
From: gac-x05la <gac-x05la@x05la-bec-001.psi.ch>
Date: Thu, 9 Jan 2025 17:22:58 +0100
Subject: [PATCH] Closer to proper stopping

---
 tomcat_bec/devices/gigafrost/stddaq_client.py | 94 +++++++++++--------
 1 file changed, 56 insertions(+), 38 deletions(-)

diff --git a/tomcat_bec/devices/gigafrost/stddaq_client.py b/tomcat_bec/devices/gigafrost/stddaq_client.py
index 20173ec..4cd144b 100644
--- a/tomcat_bec/devices/gigafrost/stddaq_client.py
+++ b/tomcat_bec/devices/gigafrost/stddaq_client.py
@@ -83,35 +83,7 @@ class StdDaqMixin(CustomDeviceMixin):
             self.parent.configure(d=d)
 
         # Try to start a new run
-        file_path = self.parent.file_path.get()
-        num_images = self.parent.num_images.get()
-        message = {"command": "start", "path": file_path, "n_image": num_images, }
-        ii = 0
-        while True:
-            self.parent.connect()
-            reply = self.parent.message(message)
-
-            if reply is not None:
-                reply = json.loads(reply)
-                self.parent.status.set(reply["status"], force=True).wait()
-                logger.info(f"[{self.parent.name}] Start DAQ reply:  {reply}")
-                # Give it more time to reconfigure
-                if reply["status"] in ("rejected"):
-                    # FIXME: running exposure is a nogo
-                    if reply['reason'] == "gerhtrhjfjf":
-                        raise RuntimeError(f"[{self.parent.name}] Start StdDAQ command rejected: already running")
-                    else:
-                        # Give it more time to restart
-                        sleep(2)
-                else:
-                    break
-            ii += 1
-            if ii == 5:
-                break
-        if reply is not None and reply["status"] in ("rejected"):
-            raise RuntimeError(
-                f"Start StdDAQ command rejected (might be already running): {reply['reason']}"
-            )
+        self.parent.bluestage()
         # And start status monitoring
         self._mon = Thread(target=self.poll, daemon=True)
         self._mon.start()
@@ -165,7 +137,7 @@ class StdDaqClient(PSIDeviceBase):
     """
     # pylint: disable=too-many-instance-attributes
     custom_prepare_cls = StdDaqMixin
-    USER_ACCESS = ["set_daq_config", "get_daq_config", "surestop", "nuke", "connect", "message", "state"]
+    USER_ACCESS = ["set_daq_config", "get_daq_config", "surestop", "nuke", "connect", "message", "state", "bluestage"]
     _wsclient = None
 
     # Status attributes
@@ -211,13 +183,6 @@ class StdDaqClient(PSIDeviceBase):
         except Exception as ex:
             logger.error(f"Failed to connect to the stdDAQ REST API\n{ex}")
 
-    def __del__(self) -> None:
-        try:
-            self._wsclient.close()
-        except TypeError:
-            pass
-        return super().__del__()
-
     def connect(self):
         """Connect to the StdDAQ's websockets interface
 
@@ -352,6 +317,59 @@ class StdDaqClient(PSIDeviceBase):
             cfg=self.read_daq_config()
             print(cfg)
 
+    def bluestage(self):
+
+        file_path = self.file_path.get()
+        num_images = self.num_images.get()
+        message = {"command": "start", "path": file_path, "n_image": num_images, }
+        ii = 0
+        while True:
+            self.connect()
+            reply = self.message(message)
+
+            if reply is not None:
+                reply = json.loads(reply)
+                self.status.set(reply["status"], force=True).wait()
+                logger.info(f"[{self.name}] Start DAQ reply:  {reply}")
+                # Give it more time to reconfigure
+                if reply["status"] in ("rejected"):
+                    # FIXME: running exposure is a nogo
+                    if reply['reason'] == "gerhtrhjfjf":
+                        raise RuntimeError(f"[{self.name}] Start StdDAQ command rejected: already running")
+                    else:
+                        # Give it more time to restart
+                        sleep(2)
+                else:
+                    break
+            ii += 1
+            if ii == 5:
+                break
+        if reply is not None and reply["status"] in ("rejected"):
+            raise RuntimeError(
+                f"Start StdDAQ command rejected (might be already running): {reply['reason']}"
+            )
+        
+
+    def blueunstage(self):
+
+        ii = 0
+        while True:
+            self.connect()
+            self.message({"command": "stop_all"}, wait_reply=False)
+            # Reply is always "success"
+
+            self.connect()
+            reply = self.message({"command": "status"})
+            if reply is not None:
+                logger.info(f"[{self.name}] DAQ status reply: {reply}")
+                reply = json.loads(reply)
+                # Give it more time to reconfigure
+                if reply["status"] in ("idle") or ii >= 5:
+                    break
+            ii += 1
+        if reply is not None and reply["status"] not in ("idle"):
+            raise RuntimeError(f"Failed to stop StdDAQ: {reply}")
+
     def get_daq_config(self) -> dict:
         """Read the current configuration from the DAQ
         """
@@ -461,5 +479,5 @@ class StdDaqClient(PSIDeviceBase):
 
 # Automatically connect to microXAS testbench if directly invoked
 if __name__ == "__main__":
-    daq = StdDaqClient(name="daq", ws_url="ws://xbl-daq-29:8080", rest_url="http://xbl-daq-29:5000")
+    daq = StdDaqClient(name="daq", ws_url="ws://sls-daq-001:8080", rest_url="http://sls-daq-001:5000")
     daq.wait_for_connection()