From e99b710d126c0afc4317447f8f656cfd7135fc96 Mon Sep 17 00:00:00 2001
From: Maxim Medvedev <redrathnure@gmail.com>
Date: Thu, 2 Jan 2025 01:20:40 +0100
Subject: [PATCH] feat: add screenshots capture

---
 src/NanoVNASaver/Hardware/litevna_64.py | 78 ++++++++++++++++++++++++-
 tests/Hardware/__init__.py              |  0
 tests/Hardware/test_litevna_64.py       | 25 ++++++++
 3 files changed, 101 insertions(+), 2 deletions(-)
 create mode 100644 tests/Hardware/__init__.py
 create mode 100644 tests/Hardware/test_litevna_64.py

diff --git a/src/NanoVNASaver/Hardware/litevna_64.py b/src/NanoVNASaver/Hardware/litevna_64.py
index 6b629a46..12fc4a99 100644
--- a/src/NanoVNASaver/Hardware/litevna_64.py
+++ b/src/NanoVNASaver/Hardware/litevna_64.py
@@ -1,9 +1,10 @@
 import logging
 import platform
-from struct import pack
+from struct import pack, unpack
 from time import sleep
 
-from serial import Serial
+from PyQt6.QtGui import QImage, QPixmap
+from serial import Serial, SerialException
 
 from NanoVNASaver.Hardware.Serial import Interface
 
@@ -32,6 +33,33 @@
 
 
 _ADDR_VBAT_MILIVOLTS = 0x5C
+_ADDR_SCREENSHOT = 0xEE
+
+
+SUPPORTED_PIXEL_FORMAT = 16
+
+
+class ScreenshotData:
+    header_size = 2 + 2 + 1
+
+    def __init__(self, width: int, height: int, pixel_size: int):
+        self.width = width
+        self.height = height
+        self.pixel_size = pixel_size
+        self.data = bytes()
+
+    def data_size(self) -> int:
+        return self.width * self.height * int(self.pixel_size / 8)
+
+    def __repr__(self) -> str:
+        return f"{self.width}x{self.height} {self.pixel_size}bits ({self.data_size()} Bytes)"
+
+    @staticmethod
+    def from_header(header_data: bytes) -> "ScreenshotData":
+        logger.debug("Screenshot header: %s", header_data)
+
+        width, height, depth = unpack("<HHB", header_data)
+        return ScreenshotData(width, height, depth)
 
 
 class LiteVNA64(NanoVNA_V2):
@@ -68,6 +96,7 @@ def init_features(self) -> None:
         # VBat state will be added dynamicly in get_features()
 
         self.features.add("Customizable data points")
+        self.features.add("Screenshots")
 
         # TODO: more than one dp per freq
         self.features.add("Multi data points")
@@ -177,3 +206,48 @@ def setSweep(self, start, stop):
             self.sweepStepHz,
         )
         self._updateSweep()
+
+    def getScreenshot(self) -> QPixmap:
+        logger.debug("Capturing screenshot...")
+        self.serial.timeout = 8
+        if self.connected():
+            try:
+                screenshot = self._get_screenshot()
+
+                if screenshot.pixel_size != SUPPORTED_PIXEL_FORMAT:
+                    logger.warning(
+                        "Unsupported %d screenshot pixel format!",
+                        screenshot.pixel_size,
+                    )
+                    return QPixmap()
+
+                image = QImage(
+                    screenshot.data,
+                    screenshot.width,
+                    screenshot.height,
+                    QImage.Format.Format_RGB16,
+                )
+                logger.debug("Screenshot was captured")
+                return QPixmap(image)
+            except SerialException as exc:
+                logger.exception(
+                    "Exception while capturing screenshot: %s", exc
+                )
+
+        logger.debug("Unable to get screenshot")
+        return QPixmap()
+
+    def _get_screenshot(self) -> ScreenshotData:
+        with self.serial.lock:
+            self.serial.write(pack("<BBB", _CMD_WRITE, _ADDR_SCREENSHOT, 0))
+            sleep(WRITE_SLEEP)
+
+            result = ScreenshotData.from_header(
+                self.serial.read(ScreenshotData.header_size)
+            )
+
+            logger.debug("Screenshot format: %s. Loading data...", result)
+
+            result.data = self.serial.read(result.data_size())
+
+            return result
diff --git a/tests/Hardware/__init__.py b/tests/Hardware/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/tests/Hardware/test_litevna_64.py b/tests/Hardware/test_litevna_64.py
new file mode 100644
index 00000000..cbafc552
--- /dev/null
+++ b/tests/Hardware/test_litevna_64.py
@@ -0,0 +1,25 @@
+
+from NanoVNASaver.Hardware.litevna_64 import ScreenshotData
+
+VALID_HEADER = b'\xe0\x01@\x01\x10'
+
+class TestScreenshotData:
+
+    @staticmethod
+    def test_from_header() -> None:
+        result = ScreenshotData.from_header(VALID_HEADER)
+
+        assert result.width == 480
+        assert result.height == 320
+        assert result.pixel_size == 16
+        assert len(result.data) == 0
+
+    @staticmethod
+    def test_data_size() -> None:
+        assert  ScreenshotData(0,0,0).data_size() == 0
+        assert  ScreenshotData(480,320,16).data_size() == 307200
+
+    @staticmethod
+    def test_repr() -> None:
+        assert  f"{ScreenshotData(0,0,0)}" == "0x0 0bits (0 Bytes)"
+        assert  f"{ScreenshotData(480,320,16)}" == "480x320 16bits (307200 Bytes)"