From 6c7a174424c264dab67cf42813f76d127657ae8d Mon Sep 17 00:00:00 2001
From: Chris Gerth <gerth2@users.noreply.github.com>
Date: Fri, 3 Jan 2025 12:23:27 -0600
Subject: [PATCH] Adding timeSyncServer for Python (#1675)

---
 photon-lib/py/photonlibpy/photonCamera.py     |  4 +
 .../py/photonlibpy/timesync/__init__.py       |  1 +
 .../py/photonlibpy/timesync/timeSyncServer.py | 94 +++++++++++++++++++
 3 files changed, 99 insertions(+)
 create mode 100644 photon-lib/py/photonlibpy/timesync/__init__.py
 create mode 100644 photon-lib/py/photonlibpy/timesync/timeSyncServer.py

diff --git a/photon-lib/py/photonlibpy/photonCamera.py b/photon-lib/py/photonlibpy/photonCamera.py
index bcbc70d90..512e3633b 100644
--- a/photon-lib/py/photonlibpy/photonCamera.py
+++ b/photon-lib/py/photonlibpy/photonCamera.py
@@ -27,6 +27,7 @@ from wpilib import RobotController, Timer
 
 from .packet import Packet
 from .targeting.photonPipelineResult import PhotonPipelineResult
+from .timesync.timeSyncServer import inst
 from .version import PHOTONLIB_VERSION  # type: ignore[import-untyped]
 
 
@@ -104,6 +105,9 @@ class PhotonCamera:
         self._prevHeartbeat = 0
         self._prevHeartbeatChangeTime = Timer.getFPGATimestamp()
 
+        # Start the time sync server
+        inst.start()
+
     def getAllUnreadResults(self) -> List[PhotonPipelineResult]:
         """
         The list of pipeline results sent by PhotonVision since the last call to getAllUnreadResults().
diff --git a/photon-lib/py/photonlibpy/timesync/__init__.py b/photon-lib/py/photonlibpy/timesync/__init__.py
new file mode 100644
index 000000000..91438c9ee
--- /dev/null
+++ b/photon-lib/py/photonlibpy/timesync/__init__.py
@@ -0,0 +1 @@
+# no one but us chickens
diff --git a/photon-lib/py/photonlibpy/timesync/timeSyncServer.py b/photon-lib/py/photonlibpy/timesync/timeSyncServer.py
new file mode 100644
index 000000000..f7a860dd9
--- /dev/null
+++ b/photon-lib/py/photonlibpy/timesync/timeSyncServer.py
@@ -0,0 +1,94 @@
+import logging
+import socket
+import struct
+import threading
+from typing import Callable, Optional
+
+from wpilib import Timer
+
+
+class TspPing:
+    def __init__(self, version: int, message_id: int, client_time: int):
+        self.version = version
+        self.message_id = message_id
+        self.client_time = client_time
+
+    @staticmethod
+    def unpack(data: bytes) -> "TspPing":
+        # Unpack using struct.unpack
+        version, message_id, client_time = struct.unpack("<BBQ", data)
+        return TspPing(version, message_id, client_time)
+
+    def pack(self) -> bytes:
+        # Pack using struct.pack
+        return struct.pack("<BBQ", self.version, self.message_id, self.client_time)
+
+
+class TspPong:
+    def __init__(self, ping: "TspPing", server_time: int):
+        self.version = ping.version
+        self.message_id = 2  # Pong message ID
+        self.client_time = ping.client_time
+        self.server_time = server_time
+
+    def pack(self) -> bytes:
+        # Pack using struct.pack
+        return struct.pack(
+            "<BBQQ", self.version, self.message_id, self.client_time, self.server_time
+        )
+
+    @staticmethod
+    def unpack(data: bytes) -> "TspPong":
+        # Unpack using struct.unpack
+        version, message_id, client_time, server_time = struct.unpack("<BBQQ", data)
+        ping = TspPing(version, message_id, client_time)
+        return TspPong(ping, server_time)
+
+
+class TimeSyncServer:
+    """This class is a python re-write of the UDP time sync server protocol
+    which runs on a roboRIO to establish a timebase for all PhotonVision coprocessors.
+    """
+
+    PORT = 5810
+
+    def __init__(self, time_provider: Optional[Callable[[], int]] = None):
+        self.time_provider = time_provider or Timer.getFPGATimestamp
+        self._process: Optional[threading.Thread] = None
+        self.logger = logging.getLogger("PhotonVision-TimeSyncServer")
+
+    def _udp_server(self):
+        with socket.socket(socket.AF_INET, socket.SOCK_DGRAM) as udp_socket:
+            udp_socket.bind(("0.0.0.0", self.PORT))
+            while True:
+                data, addr = udp_socket.recvfrom(1024)  # Buffer size of 1024 bytes
+
+                if len(data) < 10:
+                    self.logger.error("Too few bytes")
+                    continue  # Ignore incomplete packets
+
+                ping = TspPing.unpack(data)
+                if ping.version != 1 or ping.message_id != 1:
+                    self.logger.error("Invalid Version/ID")
+                    continue  # Ignore invalid pings
+
+                server_time = int(self.time_provider() * 1e6)  # Convert to microseconds
+                pong = TspPong(ping, server_time)
+                udp_socket.sendto(pong.pack(), addr)
+
+    def start(self):
+        if self._process is not None and self._process.is_alive():
+            return  # Nothing to do
+
+        self._process = threading.Thread(target=self._udp_server, daemon=True)
+        self._process.start()
+        self.logger.info("Server Started")
+
+    def stop(self):
+        if self._process is not None:
+            self._process.join()
+            self._process = None
+            self.logger.info("Server Stopped")
+
+
+inst = TimeSyncServer()