#!/usr/bin/env python3
"""Tiny client to send sim-level control commands to a running g1-sim container.

The sim (in meshcat mode) binds a ZMQ REP socket on G1_CONTROL_PORT (default 6003;
port 6002 is reserved for the gripper channel that mirrors the real robot).
This talks to it over IP, so it works from the host or a remote machine.

Only dependency is pyzmq. Examples:

    # toggle the harness / elastic band (equivalent to pressing "9" in the viewer)
    python g1_sim_ctl.py band_toggle
    python g1_sim_ctl.py band off
    python g1_sim_ctl.py band on

    # raise / lower the harness attachment length
    python g1_sim_ctl.py band_length 0.5
    python g1_sim_ctl.py band_length_delta -0.1

    # reset the sim, or send a raw sim keybind (9/backspace/up/down/left/right)
    python g1_sim_ctl.py reset
    python g1_sim_ctl.py key 9

    # point at a remote host / custom port
    python g1_sim_ctl.py --host 192.168.1.50 --port 6003 band_toggle
"""

import argparse
import json
import sys

import zmq


def build_message(args: argparse.Namespace) -> dict:
    action = args.action
    if action == "key":
        if args.value is None:
            sys.exit("`key` requires a value, e.g. `key 9`")
        return {"key": str(args.value)}
    if action == "band":
        if args.value not in ("on", "off"):
            sys.exit("`band` requires `on` or `off`")
        return {"cmd": "band_enable", "value": args.value == "on"}
    if action in ("band_length", "band_length_delta"):
        if args.value is None:
            sys.exit(f"`{action}` requires a float value")
        return {"cmd": action, "value": float(args.value)}
    # reset, band_toggle, and any other bare command
    return {"cmd": action}


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("action", help="reset | band_toggle | band | band_length | band_length_delta | key")
    parser.add_argument("value", nargs="?", default=None, help="value for the action, when applicable")
    parser.add_argument("--host", default="127.0.0.1")
    parser.add_argument("--port", type=int, default=6003)
    parser.add_argument("--timeout", type=float, default=2.0, help="reply timeout in seconds")
    args = parser.parse_args()

    msg = build_message(args)

    ctx = zmq.Context.instance()
    sock = ctx.socket(zmq.REQ)
    sock.setsockopt(int(zmq.RCVTIMEO), int(args.timeout * 1000))
    sock.setsockopt(int(zmq.LINGER), 0)
    sock.connect(f"tcp://{args.host}:{args.port}")
    sock.send_json(msg)
    try:
        resp = sock.recv_json()
    except zmq.Again:
        sys.exit(f"No reply from {args.host}:{args.port} within {args.timeout}s (is the sim running in meshcat mode?)")
    print(json.dumps(resp, indent=2))


if __name__ == "__main__":
    main()