From af94edae60cd0cac86071d0b24ff3a6780c0574f Mon Sep 17 00:00:00 2001
From: Alex Davies <traverse.da@gmail.com>
Date: Sat, 23 Nov 2024 16:33:39 -0400
Subject: [PATCH] Refactor to make launch part of Robot. Auto detect compose
 files

---
 guiTools/spiri_sdk_guitools/launcher.py  | 27 ++----------------
 guiTools/spiri_sdk_guitools/sim_drone.py | 36 +++++++++++++++++++++++-
 2 files changed, 38 insertions(+), 25 deletions(-)

diff --git a/guiTools/spiri_sdk_guitools/launcher.py b/guiTools/spiri_sdk_guitools/launcher.py
index 53e6082..3deeb3c 100644
--- a/guiTools/spiri_sdk_guitools/launcher.py
+++ b/guiTools/spiri_sdk_guitools/launcher.py
@@ -27,7 +27,7 @@ def launch_app(command):
 ui.label("Spiri Robotics SDK").style('font-size: 40px; margin-bottom: 10px;').classes('w-full text-center')
 
 robots = []
-from spiri_sdk_guitools.sim_drone import Robot
+from spiri_sdk_guitools.sim_drone import robot_types
 import aiodocker
 import asyncio
 
@@ -65,29 +65,8 @@ async def main():
             #Add a new robot
             with new_robot_widget:
                 ui.label("Add new robot").classes("text-3xl")
-                newRobotParams = defaultdict(binding.BindableProperty)
-                ui.number(value=1, label="SysID", min=1, max=254,
-                          ).bind_value(newRobotParams, 'sysid')
-                default_robot_compose = (
-                        "/robots/spiri-mu/core/docker-compose.yaml\n"
-                        "#/robots/spiri-mu/virtual_camera/docker-compose.yaml --build"
-                )
-                ui.label("Compose files").classes("text-xl")
-                ui.codemirror(value=default_robot_compose, language="bash", theme="basicDark").bind_value(newRobotParams, 'compose_files')
-                async def new_robot():
-                    compose_files = []
-                    #Split on comma or newline, and remove comments
-                    for line in newRobotParams['compose_files'].split('\n'):
-                        line = line.split('#')[0].strip()
-                        if line:
-                            compose_files.append(line)
-                    current_robot = newRobotParams.copy()
-                    current_robot['compose_files'] = compose_files
-                    robot = Robot(**current_robot)
-                    asyncio.tasks.create_task(robot.ui(robots_widget))
-
-                    newRobotParams['sysid'] += 1
-                ui.button("Add", on_click=new_robot)
+                for name, robot in robot_types.items():
+                    await robot.launch_widget(robots_widget)
 
 # Start the NiceGUI application
 ui.run(title="Spiri SDK Launcher", port=8923, dark=None)
diff --git a/guiTools/spiri_sdk_guitools/sim_drone.py b/guiTools/spiri_sdk_guitools/sim_drone.py
index b46b15d..44e08da 100644
--- a/guiTools/spiri_sdk_guitools/sim_drone.py
+++ b/guiTools/spiri_sdk_guitools/sim_drone.py
@@ -5,13 +5,14 @@ from typing import List
 import os
 import sh
 import subprocess
-from nicegui import ui, run, app
+from nicegui import ui, run, app, binding
 import yaml
 
 import docker
 import aiodocker
 import asyncio
 from spiri_sdk_guitools.video_button import EnableStreamingButton
+from collections import defaultdict
 
 docker_client = docker.from_env()
 
@@ -73,7 +74,13 @@ async def container_logs(container, element):
                 # ui.html(conv.convert(bytes(log,'utf-8').decode('utf-8', 'xmlcharrefreplace'), full=False))
 
 
+robot_types = {}
+
 class Robot:
+    def __init_subclass__(self):
+        robot_types[self.robot_type] = self
+
+class Spirimu(Robot):
     robot_type = "spiri_mu"
 
     def __init__(self, sysid: int, compose_files: List[Path] | str):
@@ -92,6 +99,33 @@ class Robot:
         self.robot_name = f"{self.robot_type}_{self.sysid}".replace("-","_")
         self.world_name = "citadel_hill"
 
+    @classmethod
+    async def launch_widget(cls, robots_widget):
+        newRobotParams = defaultdict(binding.BindableProperty)
+        ui.number(value=1, label="SysID", min=1, max=254,
+                  ).bind_value(newRobotParams, 'sysid')
+        default_robot_compose = ""
+        for compose_file in Path("/robots").glob("**/docker-compose.yaml"):
+            default_robot_compose += f"{compose_file} --build \n"
+
+        ui.label("Compose files").classes("text-xl")
+        ui.codemirror(value=default_robot_compose, language="bash", theme="basicDark").bind_value(newRobotParams, 'compose_files')
+        async def new_robot():
+            compose_files = []
+            #Split on comma or newline, and remove comments
+            for line in newRobotParams['compose_files'].split('\n'):
+                line = line.split('#')[0].strip()
+                if line:
+                    compose_files.append(line)
+            current_robot = newRobotParams.copy()
+            current_robot['compose_files'] = compose_files
+            robot = cls(**current_robot)
+            asyncio.tasks.create_task(robot.ui(robots_widget))
+
+            newRobotParams['sysid'] += 1
+        ui.button("Add", on_click=new_robot)
+
+
     async def ui_containers(self, element):
         docker_elements = {}
         container_status = {}