---

import asyncio
import json
from pynput import keyboard

ROBOT_IP = "10.10.10.10"
WS_URL = f"ws://{ROBOT_IP}:8765"

# Скорости
LINEAR_SPEED = 0.2    # м/с
ANGULAR_SPEED = 0.5   # рад/с

# Текущие нажатые клавиши
keys_pressed = set()
ws = None

async def send_cmd_vel(linear_x, angular_z):
    """Отправка скорости в топик /cmd_vel через WebSocket"""
    if ws is None:
        return
    
    # JSON-сообщение формата Foxglove / Rosbridge
    msg = {
        "op": "publish",
        "topic": "/cmd_vel",
        "msg": {
            "linear": {"x": float(linear_x), "y": 0.0, "z": 0.0},
            "angular": {"x": 0.0, "y": 0.0, "z": float(angular_z)}
        }
    }
    try:
        await ws.send(json.dumps(msg))
    except Exception as e:
        print(f"Ошибка отправки: {e}")

def on_press(key):
    try:
        if hasattr(key, 'char') and key.char:
            keys_pressed.add(key.char.lower())
    except AttributeError:
        pass

def on_release(key):
    try:
        if hasattr(key, 'char') and key.char:
            keys_pressed.discard(key.char.lower())
    except AttributeError:
        pass
    if key == keyboard.Key.esc:
        return False  # Остановка слушателя по клавише Esc

async def control_loop():
    global ws
    import websockets

    print(f"Подключение к роботу {WS_URL}...")
    async with websockets.connect(WS_URL, subprotocols=["foxglove.websocket.v1"]) as websocket:
        ws = websocket
        print("\n--- Робот подключен! ---")
        print("Управление WASD (зажимайте клавишу):")
        print("  W - вперед | S - назад | A - влево | D - вправо")
        print("  Отпустили клавишу -> робот сам останавливается")
        print("  Нажмите ESC для выхода\n")

        # Фоновый слушатель нажатий клавиш
        listener = keyboard.Listener(on_press=on_press, on_release=on_release)
        listener.start()

        try:
            while listener.is_alive():
                vx, wz = 0.0, 0.0

                if 'w' in keys_pressed:
                    vx += LINEAR_SPEED
                if 's' in keys_pressed:
                    vx -= LINEAR_SPEED
                if 'a' in keys_pressed:
                    wz += ANGULAR_SPEED
                if 'd' in keys_pressed:
                    wz -= ANGULAR_SPEED

                # Отправляем команду скорости 10 раз в секунду
                await send_cmd_vel(vx, wz)
                await asyncio.sleep(0.1)

        finally:
            # Безопасная остановка при выходе
            await send_cmd_vel(0.0, 0.0)
            print("Робот остановлен.")

if __name__ == "__main__":
    asyncio.run(control_loop())