From 7cbbf7942afd18b8afa2c6f64cefe7d5dd0bcbdf Mon Sep 17 00:00:00 2001
From: appleb_m <martin.appleby@psi.ch>
Date: Fri, 22 May 2026 15:23:32 +0200
Subject: [PATCH] DAQ: added screenshot service handler

---
 .../daq/operations/screenshot/__init__.py     |   0
 src/aare/daq/operations/screenshot/service.py | 122 ++++++++++++++++++
 2 files changed, 122 insertions(+)
 create mode 100644 src/aare/daq/operations/screenshot/__init__.py
 create mode 100644 src/aare/daq/operations/screenshot/service.py

diff --git a/src/aare/daq/operations/screenshot/__init__.py b/src/aare/daq/operations/screenshot/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/src/aare/daq/operations/screenshot/service.py b/src/aare/daq/operations/screenshot/service.py
new file mode 100644
index 00000000..5906aa49
--- /dev/null
+++ b/src/aare/daq/operations/screenshot/service.py
@@ -0,0 +1,122 @@
+from dataclasses import dataclass
+from datetime import datetime
+from pathlib import Path
+import time
+from typing import Protocol
+
+import cv2
+import numpy as np
+
+from aare.common.models import SampleShortInfo
+from aare.daq.aaredb import AareWrapper
+from aare.daq.mlbox import MlBox
+from aare.devices.mx_lib import clean_filename
+
+
+class NonCriticalRunner(Protocol):
+    def __call__(
+        self,
+        action,
+        *,
+        description: str,
+        sample: SampleShortInfo | None = None,
+    ) -> object | None: ...
+
+
+class SampleProvider(Protocol):
+    @property
+    def sample(self) -> SampleShortInfo | None: ...
+
+
+class PGroupProvider(Protocol):
+    @property
+    def pgroup(self) -> str | None: ...
+
+
+@dataclass
+class ScreenshotService:
+    mlbox: MlBox
+    aare: AareWrapper
+    logger: object
+    run_noncritical: NonCriticalRunner
+    sample_provider: SampleProvider
+    pgroup_provider: PGroupProvider
+    output_dir: str = "/sls/mx/applications/logs"
+    photos_root: str = "/sls/mx/data"
+
+    def _get_inference_image(self) -> np.ndarray:
+        image = self.mlbox.get_latest_image()
+        if image is None:
+            raise RuntimeError("No inference image available from aarelc-infer")
+        return image
+
+    def _current_sample(self) -> SampleShortInfo | None:
+        provider = self.sample_provider
+        if callable(provider):
+            return provider()
+        return provider.sample
+
+    def _current_pgroup(self) -> str | None:
+        provider = self.pgroup_provider
+        if callable(provider):
+            return provider()
+        return provider.pgroup
+
+    def save_local(self, filename: str, settle_time_s: float = 0.2) -> None:
+        time.sleep(settle_time_s)
+        bgr_image = self._get_inference_image()
+        self.logger.debug(f"saving screenshot {filename} from inference image")
+        cv2.imwrite(f"{self.output_dir}/{filename}.jpg", bgr_image)
+
+    def save_to_db(self, sample_id: int, filename: str, settle_time_s: float = 0.2) -> None:
+        time.sleep(settle_time_s)
+        sample = self._current_sample()
+        bgr_image = self._get_inference_image()
+
+        def _upload() -> None:
+            self.aare.upload_image(sample_id, filename, bgr_image)
+
+        self.run_noncritical(
+            _upload,
+            description=f"screenshot upload '{filename}'",
+            sample=sample,
+        )
+
+    def send_to_db(
+        self,
+        *,
+        filename: str | None = None,
+        message: str | None = None,
+        default_message: str,
+    ) -> None:
+        sample = self._current_sample()
+        if sample is None or sample.db_id is None or sample.db_id < 0:
+            raise ValueError("No sample with a valid sample_id is mounted.")
+
+        sample_id = sample.db_id
+        pgroup = self._current_pgroup()
+        bgr_image = self._get_inference_image()
+
+        if filename:
+            safe_filename = clean_filename(filename)
+            if not pgroup:
+                raise ValueError("No active pgroup set; cannot save screenshot to photos directory.")
+
+            photos_dir = Path(self.photos_root) / pgroup / "raw" / "photos" / str(sample_id)
+            photos_dir.mkdir(parents=True, exist_ok=True)
+            photo_path = photos_dir / f"{safe_filename}.jpeg"
+            cv2.imwrite(str(photo_path), bgr_image)
+            upload_name = safe_filename
+        else:
+            upload_name = f"{sample_id}_{datetime.now().strftime('%Y%m%d_%H%M%S')}"
+
+        final_message = (message or "").strip() or default_message
+
+        def _upload() -> None:
+            self.aare.upload_image(sample_id, upload_name, bgr_image, message=final_message)
+
+        self.run_noncritical(
+            _upload,
+            description=f"send screenshot '{upload_name}'",
+            sample=sample,
+        )
\ No newline at end of file