#!/usr/bin/env python3
"""
Launch Isaac Sim + retargeting bridge, and print the Quest 3 connection URL.

Starts everything needed in one command:
  1. Isaac Sim (G1 with Inspire hands, DDS teleoperation)
  2. Retargeting bridge (WebSocket server + body retargeting + DDS publisher)

Usage:
    python launch_bridge.py
    python launch_bridge.py --task Isaac-PickPlace-RedBlock-G129-Inspire-Joint
    python launch_bridge.py --no-sim          # bridge only (if sim is already running)
    python launch_bridge.py --device cuda     # use GPU for simulation
"""

import socket
import subprocess
import sys
import os
import signal
import time
import threading


SIM_DIR = os.path.expanduser("~/git/unitree_sim_isaaclab")
BRIDGE_DIR = os.path.dirname(os.path.abspath(__file__))

# Default Isaac Sim arguments
DEFAULT_TASK = "Isaac-PickPlace-Cylinder-G129-Inspire-Joint"
DEFAULT_DEVICE = "cpu"

# Message printed by sim when DDS is ready
SIM_READY_MARKER = "start controller success"


def get_local_ip():
    """Get local network IP via UDP socket (no traffic sent)."""
    s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    try:
        s.connect(("8.8.8.8", 80))
        return s.getsockname()[0]
    finally:
        s.close()


def stream_and_watch(proc, marker, event):
    """Read process stdout line-by-line, print it, and set event when marker seen."""
    for line in iter(proc.stdout.readline, ""):
        sys.stdout.write(f"  [sim] {line}")
        sys.stdout.flush()
        if marker in line:
            event.set()
    proc.stdout.close()


def stream_stderr(proc):
    """Forward process stderr."""
    for line in iter(proc.stderr.readline, ""):
        sys.stderr.write(f"  [sim] {line}")
        sys.stderr.flush()
    proc.stderr.close()


def main():
    # --- Parse our args (separate from bridge args) ---
    port = 8765
    task = DEFAULT_TASK
    device = DEFAULT_DEVICE
    no_sim = False
    bridge_args = []

    # Simple arg parser that pulls out our flags and passes the rest to the bridge
    args = sys.argv[1:]
    i = 0
    while i < len(args):
        if args[i] == "--port" and i + 1 < len(args):
            port = int(args[i + 1])
            bridge_args.extend(args[i:i+2])
            i += 2
        elif args[i] == "--task" and i + 1 < len(args):
            task = args[i + 1]
            i += 2
        elif args[i] == "--device" and i + 1 < len(args):
            device = args[i + 1]
            i += 2
        elif args[i] == "--no-sim":
            no_sim = True
            i += 1
        else:
            bridge_args.append(args[i])
            i += 1

    ip = get_local_ip()

    # --- Banner ---
    print()
    print("=" * 60)
    print("  G1 Teleop -> Isaac Sim Retargeting Bridge")
    print("=" * 60)
    print()
    print(f"  Quest 3 connection:")
    print(f"    IP:   {ip}")
    print(f"    Port: {port}")
    print()
    if not no_sim:
        print(f"  Isaac Sim:")
        print(f"    Task:   {task}")
        print(f"    Device: {device}")
        print()
    print("=" * 60)
    print()
    sys.stdout.flush()

    sim_proc = None
    bridge_proc = None

    def cleanup(sig=None, frame=None):
        """Shut down both processes on exit."""
        print("\nShutting down...")
        if bridge_proc and bridge_proc.poll() is None:
            bridge_proc.send_signal(signal.SIGINT)
            bridge_proc.wait(timeout=5)
        if sim_proc and sim_proc.poll() is None:
            sim_proc.send_signal(signal.SIGINT)
            sim_proc.wait(timeout=10)
        sys.exit(0)

    signal.signal(signal.SIGINT, cleanup)
    signal.signal(signal.SIGTERM, cleanup)

    # --- Launch Isaac Sim ---
    if not no_sim:
        sim_cmd = [
            sys.executable, os.path.join(SIM_DIR, "sim_main.py"),
            "--device", device,
            "--enable_cameras",
            "--task", task,
            "--enable_inspire_dds",
            "--robot_type", "g129",
        ]
        print(f"[1/2] Starting Isaac Sim...")
        print(f"       {' '.join(sim_cmd)}")
        print()
        sys.stdout.flush()

        sim_proc = subprocess.Popen(
            sim_cmd,
            cwd=SIM_DIR,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            text=True,
            bufsize=1,
        )

        # Stream sim output and watch for ready marker
        sim_ready = threading.Event()
        stdout_thread = threading.Thread(
            target=stream_and_watch, args=(sim_proc, SIM_READY_MARKER, sim_ready),
            daemon=True)
        stderr_thread = threading.Thread(
            target=stream_stderr, args=(sim_proc,),
            daemon=True)
        stdout_thread.start()
        stderr_thread.start()

        # Wait for sim to be ready (timeout after 120s)
        print("  Waiting for Isaac Sim to initialize (this may take a minute)...")
        sys.stdout.flush()
        if not sim_ready.wait(timeout=120):
            if sim_proc.poll() is not None:
                print("\n  ERROR: Isaac Sim exited unexpectedly.")
                sys.exit(1)
            print("\n  WARNING: Timed out waiting for sim ready marker.")
            print("           Starting bridge anyway (sim may still be loading).")
        else:
            print("  Isaac Sim is ready!")
        print()
        sys.stdout.flush()

    # --- Launch bridge ---
    bridge_cmd = [
        sys.executable, os.path.join(BRIDGE_DIR, "retarget_bridge.py"),
        "--port", str(port),
    ] + bridge_args

    step = "2/2" if not no_sim else "1/1"
    print(f"[{step}] Starting retargeting bridge...")
    print()
    print("=" * 60)
    print(f"  READY -- connect Quest 3 to:  {ip}:{port}")
    print("=" * 60)
    print()
    sys.stdout.flush()

    bridge_proc = subprocess.Popen(bridge_cmd, cwd=BRIDGE_DIR)

    try:
        # Wait for bridge to exit (or Ctrl-C)
        bridge_proc.wait()
    except KeyboardInterrupt:
        pass
    finally:
        cleanup()


if __name__ == "__main__":
    main()