run_example.py

"""
Launcher for the dual-gem5 Crazyflie light-chase simulation.

Starts:
  1. Bridge helper server (this process)
  2. Webots                (one subprocess; runs leader + chaser + supervisor controllers)
  3. gem5-leader           (one subprocess; runs leader.elf)
  4. gem5-chaser           (one subprocess; runs chaser.elf)

The bridge helper matches the two Webots controllers by robot name to
the two gem5 server names:
    Crazyflie_Leader  <-> gem5-leader
    Crazyflie_Chaser  <-> gem5-chaser
"""

import argparse
import os
import subprocess
import time
from pathlib import Path

import bridge._bridge as br

SCRIPT_DIR = Path(__file__).resolve().parent
DEFAULT_GEM5_SCRIPT = SCRIPT_DIR / "gem5" / "gem5_script.py"
DEFAULT_WORLD = SCRIPT_DIR / "webots" / "worlds" / "crazyflie_dual_gem5.wbt"
DEFAULT_LEADER_ELF = SCRIPT_DIR / "gem5" / "firmware" / "build" / "leader.elf"
DEFAULT_CHASER_ELF = SCRIPT_DIR / "gem5" / "firmware" / "build" / "chaser.elf"

parser = argparse.ArgumentParser(
    description="Launch the dual-gem5 Crazyflie light-chase simulation."
)
parser.add_argument("--gem5-path",       required=True,
                    help="Path to gem5 executable (e.g. .../build/ARM/gem5.opt)")
parser.add_argument("--webots-path",     required=True,
                    help="Path to the Webots executable")
parser.add_argument("--gem5-script",     default=str(DEFAULT_GEM5_SCRIPT),
                    help="Path to the gem5 simulation script")
parser.add_argument("--gem5-leader-binary", default=str(DEFAULT_LEADER_ELF),
                    help="Leader firmware ELF")
parser.add_argument("--gem5-chaser-binary", default=str(DEFAULT_CHASER_ELF),
                    help="Chaser firmware ELF")
parser.add_argument("--webots-world",    default=str(DEFAULT_WORLD),
                    help="Path to the Webots .wbt world file")
parser.add_argument("--output-dir",      default=str(SCRIPT_DIR / "output"),
                    help="Directory for gem5 m5out trees")

# -- Webots display / speed control --
parser.add_argument(
    "--webots-mode", choices=["pause", "realtime", "fast"], default="realtime",
    help="Webots simulation mode (default: realtime). 'fast' runs as fast "
         "as the host can compute and implies no rendering.",
)
parser.add_argument(
    "--no-rendering", action="store_true",
    help="Disable Webots 3D rendering (physics still runs). Big speed win. "
         "Implied by --webots-mode=fast.",
)
parser.add_argument(
    "--no-minimize", action="store_true",
    help="Don't start the Webots window minimized (default: minimized).",
)
parser.add_argument(
    "--batch", action="store_true",
    help="Run Webots in batch mode (no save dialogs / quit prompts).",
)
parser.add_argument(
    "--headless", action="store_true",
    help="Shortcut for: --webots-mode=fast --no-rendering --batch. "
         "Webots still needs an X display — use xvfb-run for true headless.",
)


def start_executable(path, args, friendly_name):
    p = Path(path)
    if not p.exists():
        raise FileNotFoundError(f"{friendly_name} not found at {path}")
    if not os.access(str(p), os.X_OK):
        raise PermissionError(f"{friendly_name} at {path} is not executable")
    print(f"[launcher] Starting {friendly_name}: {path} {' '.join(args)}")
    return subprocess.Popen([str(path)] + args)


def main():
    args = parser.parse_args()

    # Robot name (Webots side) -> gem5 server name
    client_to_server = {
        "Crazyflie_Leader": "gem5-leader",
        "Crazyflie_Chaser": "gem5-chaser",
    }

    output_dir = Path(args.output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)

    for label, path in [
        ("gem5",            args.gem5_path),
        ("gem5 script",     args.gem5_script),
        ("leader firmware", args.gem5_leader_binary),
        ("chaser firmware", args.gem5_chaser_binary),
        ("webots",          args.webots_path),
        ("world",           args.webots_world),
    ]:
        if not Path(path).exists():
            raise FileNotFoundError(f"{label} not found at {path}")

    print("[launcher] Starting bridge helper server...")
    listen_fd = br.bridge_setup_helper_server_socket()
    print(f"[launcher] Helper listening on fd {listen_fd}")

    # Resolve --headless into the underlying Webots flags
    mode        = "fast" if args.headless else args.webots_mode
    no_render   = args.no_rendering or args.headless or mode == "fast"
    batch       = args.batch or args.headless

    webots_args = []
    if not args.no_minimize:
        webots_args.append("--minimize")
    webots_args.append(f"--mode={mode}")
    if no_render:
        webots_args.append("--no-rendering")
    if batch:
        webots_args.append("--batch")
    webots_args += [args.webots_world, "--stdout", "--stderr"]

    webots_proc = start_executable(args.webots_path, webots_args, "webots")
    time.sleep(0.5)

    gem5_procs = []
    for server_name, binary in [
        ("gem5-leader", args.gem5_leader_binary),
        ("gem5-chaser", args.gem5_chaser_binary),
    ]:
        m5out = output_dir / f"{server_name}-m5out"
        gem5_args = [
            "-re",
            "-d", str(m5out),
            args.gem5_script,
            "--binary", binary,
            "--server-name", server_name,
        ]
        proc = start_executable(args.gem5_path, gem5_args, server_name)
        gem5_procs.append((server_name, proc))

    print("[launcher] All processes started. Entering bridge helper loop...")
    print("[launcher] Press Ctrl+C to stop.")
    try:
        br.bridge_helper_server_loop(listen_fd, client_to_server)
    except KeyboardInterrupt:
        print("\n[launcher] Interrupted, shutting down...")
    finally:
        br.bridge_close_helper_server_socket(listen_fd)

    print("[launcher] Bridge matchmaking done. Keeping launcher alive...")
    try:
        while True:
            time.sleep(1)
    except KeyboardInterrupt:
        print("\n[launcher] Interrupted, shutting down...")

    procs = [("webots", webots_proc)] + gem5_procs
    for name, proc in procs:
        if proc is None:
            continue
        try:
            if proc.poll() is None:
                proc.terminate()
                try:
                    proc.wait(timeout=5)
                except subprocess.TimeoutExpired:
                    proc.kill()
                    proc.wait()
            print(f"[launcher] {name} exited with code {proc.returncode}")
        except Exception as e:
            print(f"[launcher] Error terminating {name}: {e}")

    print("[launcher] Done.")


if __name__ == "__main__":
    main()