from nicegui import ui, binding, app, run 
import subprocess
from collections import defaultdict
import docker
import time

docker_client = docker.from_env()

# Dictionary of applications: key is the button text, value is the command to execute
applications = {
    "Terminal": ["qterminal"],
    "rqt": ["rqt"],
    "rviz2": ["rviz2"],
    "Gazebo": ["/gz_entrypoint.sh"],
    "Gazebo Standalone": "gz sim -v4".split(),
    # Add more applications here if needed
}

# Function to launch an application
def launch_app(command):
    try:
        subprocess.Popen(command)
    except FileNotFoundError:
        print(f"{command[0]} not found. Make sure it's installed and accessible in the PATH.")

# Create the NiceGUI interface
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
import aiodocker
import asyncio

@ui.page('/')
async def main():
    with ui.tabs().classes('w-full') as tabs:
        tab_tools = ui.tab('Tools')
        tab_robots = ui.tab('Robots')
        # two = ui.tab('Two')
    with ui.tab_panels(tabs, value=tab_tools).classes("w-full"):
        with ui.tab_panel(tab_tools):
            # Create and place buttons dynamically based on the dictionary
            with ui.grid(columns=3):
                for app_name, command in applications.items():
                    ui.button(app_name, on_click=lambda cmd=command: launch_app(cmd)).style('width: 150px; height: 50px; margin: 5px;')
        with ui.tab_panel(tab_robots):
            new_robot_widget = ui.element().classes("w-1/2")
            with ui.element().classes("w-1/2"):
                ui.label("Debug").classes("text-xl")
                def cleanup_all_containers():
                    #Find all containers that start with spiri-sdk
                    containers = docker_client.containers.list(all=True)
                    removal_count = 0
                    for container in containers:
                        if container.name.startswith("robot-sim-"):
                            container.remove(force=True)
                            removal_count += 1
                    ui.label(f"Removed {removal_count} containers")
                ui.button("Cleanup all containers", on_click=cleanup_all_containers)

            ui.separator()
            ui.label("Current robots").classes("text-3xl")
            robots_widget = ui.element().classes("w-full")

            #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)

# Start the NiceGUI application
ui.run(title="Spiri SDK Launcher", port=8923, dark=None)