write shit

2026-01-06 19:50:45 +08:00
parent 4cdc56a32d
commit 6f4d23868c
6 changed files with 252 additions and 48 deletions
--- a/BasaltTrainer/command_server.py
+++ b/BasaltTrainer/command_server.py
@@ -0,0 +1,113 @@
+import struct
+from enum import IntEnum, Enum, auto
+from dataclasses import dataclass
+from typing import Tuple, Optional
+from pipe_operator import PipeOperator
+
+PIPE_NAME: str = "ed0e3f1f-d214-4880-9562-640bce15e72e"
+
+class ProtocolCode(IntEnum):
+    HANDSHAKE_REQUEST = 0x61  # Trainer -> Presenter
+    HANDSHAKE_RESPONSE = 0x62  # Presenter -> Trainer
+    DATA_READY = 0x01  # Presenter -> Trainer
+    DATA_RECEIVED = 0x02  # Trainer -> Presenter
+    REQUEST_STOP = 0x71  # Presenter -> Trainer (request stop)
+    STOP = 0x71  # Trainer -> Presenter (confirm stop)
+
+class PixelKind(IntEnum):
+    GRAY_FLOAT32 = 0x01  # Grayscale represented by one float32
+    GRAY_U8 = 0x02  # Grayscale represented by one u8
+    RGB_FLOAT32 = 0x03  # RGB represented by three float32
+    RGB_U8 = 0x04  # RGB represented by three u8
+
+@dataclass
+class ImageProperties:
+    pixel_kind: PixelKind
+    width: int
+    height: int
+
+class ServerStatus(Enum):
+    Ready = auto()
+    Running = auto()
+    Stop = auto()
+
+
+CODE_PACKER: struct.Struct = struct.Struct("=B")
+HANDSHAKE_RESPONSE_PACKER: struct.Struct = struct.Struct("=BII")
+
+class CommandServer:
+    """
+    Command server implementation for the Trainer side according to the protocol.
+    """
+
+    pipe_operator: PipeOperator
+    statue: ServerStatus
+
+    def __init__(self):
+        self.pipe_operator = PipeOperator(PIPE_NAME)
+        self.handshaked = ServerStatus.Ready
+
+    def __del__(self):
+        """Cleanup resources when object is destroyed."""
+        self.pipe_operator.close()
+
+    def wait_handshake(self) -> ImageProperties:
+        """
+        Wait for handshake from Presenter, send request first and wait for response with data properties.
+        Returns a tuple of (pixel_kind, width, height) from the Presenter.
+        """
+        if self.statue != ServerStatus.Ready:
+            raise RuntimeError("unexpected server status")
+
+        # Send handshake request to Presenter (code 0x61)
+        self.pipe_operator.write(CODE_PACKER.pack(ProtocolCode.HANDSHAKE_REQUEST))
+
+        # Wait for handshake response from Presenter (code 0x62)
+        code_bytes = self.pipe_operator.read(CODE_PACKER.size)
+        (code, ) = CODE_PACKER.unpack(code_bytes)
+        if ProtocolCode(code) != ProtocolCode.HANDSHAKE_RESPONSE:
+            raise RuntimeError("expect handshake response code, but got another")
+
+        # Read data properties from Presenter (pixel_kind, width, height)
+        handshake_response_payload = self.pipe_operator.read(HANDSHAKE_RESPONSE_PACKER.size)
+        (raw_pixel_kind, width, height) = HANDSHAKE_RESPONSE_PACKER.unpack(handshake_response_payload)
+
+        # Set status and return
+        self.statue = ServerStatus.Running
+        return ImageProperties(PixelKind(raw_pixel_kind), width, height)
+
+    def tick(self, request_stop: bool) -> bool:
+        """
+        Tick function called every frame to wait for data ready from Presenter and send response.
+        Returns True if a stop code was received (meaning the process should stop), False otherwise.
+        """
+        if self.statue != ServerStatus.Running:
+            raise RuntimeError("unexpected server status")
+
+        # If there is stop requested, we post it first and return
+        if request_stop:
+            self.pipe_operator.write(CODE_PACKER.pack(ProtocolCode.STOP))
+            self.statue = ServerStatus.Stop
+            return True
+
+        # Wait for code from Presenter
+        code_bytes = self.pipe_operator.read(CODE_PACKER.size)
+        (code, ) = CODE_PACKER.unpack(code_bytes)
+        
+        # Analyse code
+        match ProtocolCode(code):
+            case ProtocolCode.DATA_READY:
+                # Receive data
+                print('Data received')
+                # Send data received symbol
+                self.pipe_operator.write(CODE_PACKER.pack(ProtocolCode.DATA_RECEIVED))
+            case ProtocolCode.REQUEST_STOP:
+                # Presenter requested stop.
+                # Agree with it, send code and return.
+                self.pipe_operator.write(CODE_PACKER.pack(ProtocolCode.STOP))
+                self.statue = ServerStatus.Stop
+                return True
+            case _:
+                raise RuntimeError("unexpected protocol code when running")
+
+        return False