From 380bcb773e113cf5de33106990e003fb0ea8cc1f Mon Sep 17 00:00:00 2001
From: Filip Leonarski <filip.leonarski@psi.ch>
Date: Fri, 31 Oct 2025 16:33:07 +0100
Subject: [PATCH] GUI: status_bar now ahs logic depending on current user and
 state of beamline

---
 gui/src/aaregui/widgets/status_bar.py | 64 ++++++++++++++++++++++++---
 1 file changed, 59 insertions(+), 5 deletions(-)

diff --git a/gui/src/aaregui/widgets/status_bar.py b/gui/src/aaregui/widgets/status_bar.py
index dc92dd3e..e7513105 100644
--- a/gui/src/aaregui/widgets/status_bar.py
+++ b/gui/src/aaregui/widgets/status_bar.py
@@ -1,4 +1,4 @@
-from PySide6.QtCore import Signal, Slot, QPoint
+from PySide6.QtCore import Signal, Slot, QPoint, QTimer
 from PySide6.QtGui import QFont
 from PySide6.QtWidgets import QStatusBar, QDialog, QMenu, QLabel
 
@@ -27,6 +27,8 @@ class StatusBar(QStatusBar):
         super().__init__(parent)
         self.__decoded_token = token
         self.__status = None
+        self.__is_staff = self.__decoded_token.staff
+        self.__allowed_pgroups = self.__decoded_token.pgroups
 
         self.flux = ValueLabel("Flux", "x 10<sup>9</sup> ph/s", self)
         self.transmission = ValueLabel("Transmission", "", self)
@@ -143,10 +145,13 @@ class StatusBar(QStatusBar):
 
     def show_session_menu(self):
         menu = QMenu(self)
-
+        is_busy = self.__status and self.__status.busy
+        is_vacant = self.__status and self.__status.session.session == SessionsStateEnum.Vacant
         action_1 = menu.addAction("Grab")
-        action_1.triggered.connect(self.grab_session_clicked)
+        action_1.setEnabled(bool(not is_busy or self.__is_staff or is_vacant))
+        action_1.triggered.connect(self._on_grab_clicked)
         action_2 = menu.addAction("End")
+        action_2.setEnabled(bool(not is_busy or self.__is_staff))
         action_2.triggered.connect(self.end_session_clicked)
 
         label_geometry = self.session_label.geometry()
@@ -157,10 +162,23 @@ class StatusBar(QStatusBar):
         menu.exec()
 
     def show_pgroup_menu(self):
-        if self.__decoded_token.staff:
+        in_curr = self.__status and self.__status.session.current_pgroup in  (self.__allowed_pgroups or [])
+
+        logger.info(f"in_curr is {in_curr}")
+
+        if not in_curr:
+            print(f"is status None?: {self.__status}")
+
+        allow_pgroup_menu = self.__is_staff or in_curr
+
+        owned_by_you = self.__status and self.__status.session.session == SessionsStateEnum.OwnedByYou
+
+        if allow_pgroup_menu:
+            logger.info("showing pgroup menu")
             menu = QMenu(self)
 
             action_1 = menu.addAction("Change")
+            action_1.setEnabled(bool(owned_by_you))
             action_1.triggered.connect(self.show_change_dialog)
 
             label_geometry = self.pgroup_label.geometry()
@@ -226,6 +244,39 @@ class StatusBar(QStatusBar):
 
         menu.exec()
 
+    def _on_grab_clicked(self):
+        self.grab_session_clicked()
+
+        def after_grab():
+
+            if not self.__status:
+                logger.info("status is None when session is grabbed")
+                return
+
+            check_state = self.__status.state in (
+                BeamlineStateEnum.SampleAlignment,
+                BeamlineStateEnum.SampleExchange,
+                BeamlineStateEnum.DewarTransfer,
+                BeamlineStateEnum.Maintenance,
+            )
+
+            session_ownership = self.__status.session.session == SessionsStateEnum.OwnedByYou
+            have_pgroup = (self.__status.session.current_pgroup in (self.__allowed_pgroups or []))
+            allowed = (not self.__status.busy and check_state and session_ownership and have_pgroup) or self.__is_staff
+            logger.info(f"do I have the baton in on grab clicked? {session_ownership}")
+
+            if allowed:
+                self.show_change_dialog()
+
+            else:
+                logger.debug(
+                    "not allowed to change pgroup due to; "
+                    f"session ownership: {session_ownership}, allowed_pgroup: {have_pgroup}, "
+                    f"beamline busy: {self.__status.busy}, beamline state: {check_state}"
+                )
+
+        QTimer.singleShot(300, after_grab)
+
     def grab_session_clicked(self):
         self.force_session.emit()
 
@@ -248,7 +299,10 @@ class StatusBar(QStatusBar):
         self.open_shutter.emit()
 
     def show_change_dialog(self):
-        dialog = PGroupDialog()
+        logger.debug(self.__decoded_token.pgroups)
+        curr = self.__status.session.current_pgroup
+        pgroups = [str(p) for p in (self.__allowed_pgroups or []) if p is not None and str(p).strip()]
+        dialog = PGroupDialog(curr_pgroup=curr, pgroups=pgroups)
         if dialog.exec() == QDialog.DialogCode.Accepted:
             entered_text = dialog.get_input()
             self.set_pgroup.emit(entered_text)