#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
WL-PRO 轮组机器人 遥控脚本（纯标准库，无第三方依赖，Windows/macOS/Linux 通用）

协议与网页 index.html 完全一致：ws://<host>:81/ ，发送 JSON 状态帧：
  {"dir":..., "mode":"basic", "roll":0, "height":38, "linear":0,
   "angular":0, "stable":0, "joy_x":0, "joy_y":0}
  dir 取值：stop | forward | back | left | right | jump

用法：
  python robot-ctl.py --host 192.168.4.1 forward          # 前进 0.8s 后停
  python robot-ctl.py jump                                   # 跳一下
  python robot-ctl.py --host 192.168.4.1 joy 0 100         # 摇杆 前推到底 0.8s
  python robot-ctl.py joy -60 0 --ms 500                     # 左移
  python robot-ctl.py height 60                              # 站高 60mm（32~85）
  python robot-ctl.py roll 10                                # 横滚 10°（-30~30）
  python robot-ctl.py linear 150 --ms 600                    # 线速度 150mm/s
  python robot-ctl.py seq "height 60; forward 800; stop"
  python robot-ctl.py stop                                   # 急停
  python robot-ctl.py probe                                  # 只连不动作，打印回传

参数：--host（默认 192.168.4.1） --port（默认 81） --ms（动作持续毫秒）
      --rate（推帧间隔毫秒，默认 150，与网页一致） --timeout（连接超时秒）
"""

import argparse
import base64
import hashlib
import json
import os
import socket
import struct
import sys
import time

DEFAULT_STATE = {
    "roll": 0, "height": 38, "linear": 0, "angular": 0,
    "stable": 0, "mode": "basic", "dir": "stop", "joy_x": 0, "joy_y": 0,
}
DIRS = {"forward", "back", "left", "right", "jump"}
GUID = "258EAFA5-E914-47DA-95CA-C5AB0DC85B11"


class WS:
    """极简 WebSocket 客户端（仅文本帧 + ping/pong/close）"""

    def __init__(self, host, port, timeout=8.0, verbose=True):
        self.host, self.port, self.timeout, self.verbose = host, port, timeout, verbose
        self.sock = None
        self.buf = b""

    def log(self, *a):
        if self.verbose:
            print("[robot]", *a, flush=True)

    def connect(self):
        s = socket.create_connection((self.host, self.port), timeout=self.timeout)
        s.settimeout(self.timeout)
        key = base64.b64encode(os.urandom(16)).decode()
        req = (
            "GET / HTTP/1.1\r\n"
            f"Host: {self.host}:{self.port}\r\n"
            "Upgrade: websocket\r\nConnection: Upgrade\r\n"
            f"Sec-WebSocket-Key: {key}\r\nSec-WebSocket-Version: 13\r\n"
            f"Origin: http://{self.host}:{self.port}\r\n\r\n"
        )
        s.sendall(req.encode())
        # 读响应头
        head = b""
        while b"\r\n\r\n" not in head:
            chunk = s.recv(4096)
            if not chunk:
                raise RuntimeError("握手时连接被关闭")
            head += chunk
            if len(head) > 65536:
                raise RuntimeError("响应头过大")
        header, rest = head.split(b"\r\n\r\n", 1)
        lines = header.decode("utf-8", "replace").split("\r\n")
        status = lines[0]
        if "101" not in status:
            raise RuntimeError(f"设备未接受 WebSocket 升级：{status}")
        hdrs = {}
        for ln in lines[1:]:
            if ":" in ln:
                k, v = ln.split(":", 1)
                hdrs[k.strip().lower()] = v.strip()
        expect = base64.b64encode(hashlib.sha1((key + GUID).encode()).digest()).decode()
        if hdrs.get("sec-websocket-accept") != expect:
            raise RuntimeError("Sec-WebSocket-Accept 校验失败")
        self.sock, self.buf = s, rest
        self.log("已连接", f"ws://{self.host}:{self.port}/")
        return self

    def send_text(self, text):
        if self.sock is None:
            raise RuntimeError("未连接")
        payload = text.encode()
        n = len(payload)
        header = bytearray([0x81])
        if n < 126:
            header.append(0x80 | n)
        elif n < 65536:
            header.append(0x80 | 126)
            header += struct.pack("!H", n)
        else:
            header.append(0x80 | 127)
            header += struct.pack("!Q", n)
        mask = os.urandom(4)
        header += mask
        masked = bytes(b ^ mask[i % 4] for i, b in enumerate(payload))
        self.sock.sendall(bytes(header) + masked)

    def _recv_more(self):
        chunk = self.sock.recv(4096)
        if not chunk:
            raise ConnectionError("连接已关闭")
        self.buf += chunk

    def recv_frame(self, timeout=0.2):
        """返回 (opcode, payload_bytes) 或 None（超时）"""
        self.sock.settimeout(timeout)
        try:
            while True:
                if len(self.buf) >= 2:
                    b0, b1 = self.buf[0], self.buf[1]
                    opcode = b0 & 0x0F
                    masked = b1 & 0x80
                    ln = b1 & 0x7F
                    off = 2
                    if ln == 126:
                        if len(self.buf) < 4:
                            self._recv_more(); continue
                        ln = struct.unpack("!H", self.buf[2:4])[0]; off = 4
                    elif ln == 127:
                        if len(self.buf) < 10:
                            self._recv_more(); continue
                        ln = struct.unpack("!Q", self.buf[2:10])[0]; off = 10
                    if masked:
                        if len(self.buf) < off + 4:
                            self._recv_more(); continue
                        mk = self.buf[off:off + 4]; off += 4
                    else:
                        mk = None
                    if len(self.buf) < off + ln:
                        self._recv_more(); continue
                    payload = self.buf[off:off + ln]
                    self.buf = self.buf[off + ln:]
                    if mk:
                        payload = bytes(b ^ mk[i % 4] for i, b in enumerate(payload))
                    if opcode == 9:  # ping -> pong
                        self._send_control(0x0A, payload); continue
                    if opcode == 10:
                        continue
                    return opcode, payload
                self._recv_more()
        except socket.timeout:
            return None
        except (BlockingIOError, InterruptedError):
            return None

    def _send_control(self, opcode, payload=b""):
        mask = os.urandom(4)
        masked = bytes(b ^ mask[i % 4] for i, b in enumerate(payload))
        self.sock.sendall(bytes([0x80 | opcode, 0x80 | len(payload)]) + mask + masked)

    def drain(self, seconds):
        """收一段时间的回传（设备可能会回状态/心跳）"""
        end = time.time() + seconds
        got = 0
        while time.time() < end:
            f = self.recv_frame(timeout=min(0.2, max(0.02, end - time.time())))
            if f is None:
                continue
            op, payload = f
            if op == 8:
                break
            if op == 1:
                got += 1
                try:
                    txt = payload.decode("utf-8", "replace")
                except Exception:
                    txt = repr(payload)
                self.log("收到:", txt[:400])
        return got

    def close(self):
        if self.sock:
            try:
                self._send_control(0x08, struct.pack("!H", 1000))
            except Exception:
                pass
            try:
                self.sock.close()
            except Exception:
                pass
            self.sock = None


class Robot:
    def __init__(self, ws, rate=150, base=None):
        self.ws, self.rate = ws, rate
        self.state = dict(DEFAULT_STATE)
        if base:
            # 基线：先用设备当前的真实姿态覆盖默认值，避免首个动作帧改掉姿态
            for k, v in base.items():
                if v is not None and k in self.state:
                    self.state[k] = v

    def send(self, **extra):
        self.state.update(extra)
        frame = {
            "dir": self.state["dir"], "mode": self.state["mode"],
            "roll": self.state["roll"], "height": self.state["height"],
            "linear": self.state["linear"], "angular": self.state["angular"],
            "stable": self.state["stable"],
            "joy_x": self.state["joy_x"], "joy_y": self.state["joy_y"],
        }
        self.ws.send_text(json.dumps(frame, separators=(",", ":")))
        return frame

    def stream(self, extra, ms):
        step = max(30, self.rate)
        t0 = time.time()
        self.send(**extra)
        while (time.time() - t0) * 1000.0 < ms:
            left = ms / 1000.0 - (time.time() - t0)
            if left <= 0:
                break
            time.sleep(min(step / 1000.0, left))
            if (time.time() - t0) * 1000.0 >= ms:
                break
            self.send(**extra)


def run(robot, cmd, args, opts):
    def num(v, d):
        return d if v is None or v == "" else float(v) if "." in str(v) else int(v)

    if cmd == "stop":
        robot.send(dir="stop", joy_x=0, joy_y=0)
        robot.ws.log("已停止")
        return
    if cmd == "probe":
        robot.send(dir="stop")
        got = robot.ws.drain(float(args[0]) if args else 3.0)
        robot.ws.log(f"监听结束，收到 {got} 帧")
        return
    if cmd == "raw":
        obj = json.loads(" ".join(args))
        robot.send(**obj)
        robot.ws.log("已发:", json.dumps(obj, ensure_ascii=False))
        return
    if cmd == "height":
        robot.send(height=num(args[0] if args else None, 38))
        robot.ws.log(f"站高 = {robot.state['height']}mm")
        return
    if cmd == "roll":
        robot.send(roll=num(args[0] if args else None, 0))
        robot.ws.log(f"横滚 = {robot.state['roll']}°")
        return
    if cmd in ("linear", "angular"):
        v = num(args[0] if args else None, 0)
        ms = opts.ms or (num(args[1], 800) if len(args) > 1 else 800)
        robot.stream({cmd: v}, ms)
        robot.send(**{cmd: 0})
        robot.ws.log(f"{cmd} = {v} 持续 {ms}ms 后归零")
        return
    if cmd == "stable":
        robot.send(stable=1 if num(args[0] if args else None, 0) else 0)
        robot.ws.log(f"stable = {robot.state['stable']}")
        return
    if cmd == "joy":
        x = num(args[0] if args else None, 0)
        y = num(args[1], 0) if len(args) > 1 else 0
        ms = opts.ms or (num(args[2], 800) if len(args) > 2 else 800)
        robot.stream({"joy_x": x, "joy_y": y}, ms)
        robot.send(joy_x=0, joy_y=0, dir="stop")
        robot.ws.log(f"摇杆 x={x} y={y} 持续 {ms}ms 后回中")
        return
    if cmd == "seq":
        script = " ".join(args)
        for part in [p.strip() for p in script.split(";") if p.strip()]:
            seg = part.split()
            run(robot, seg[0], seg[1:], opts)
            time.sleep(0.12)
        return
    if cmd in DIRS:
        ms = opts.ms or (300 if cmd == "jump" else 800)
        robot.stream({"dir": cmd}, ms)
        robot.send(dir="stop")
        robot.ws.log(f"{cmd} 持续 {ms}ms 后停止")
        return
    raise SystemExit(
        "未知命令: %s\n可用: stop forward back left right jump joy height roll "
        "linear angular stable seq raw probe" % cmd
    )


def main():
    ap = argparse.ArgumentParser(add_help=True, description="WL-PRO 轮组机器人遥控")
    ap.add_argument("--host", default="192.168.4.1")
    ap.add_argument("--port", type=int, default=81)
    ap.add_argument("--ms", type=int, default=0)
    ap.add_argument("--rate", type=int, default=150)
    ap.add_argument("--timeout", type=float, default=8.0)
    # 设备当前真实姿态（设备不回传状态，必须由人告知，否则动作帧会改掉姿态）
    ap.add_argument("--height", type=float, default=None, help="设备当前站高 mm")
    ap.add_argument("--roll", type=float, default=None, help="设备当前横滚 deg")
    ap.add_argument("--stable", type=int, default=None, help="设备当前 stable 0/1")
    ap.add_argument("--linear", type=float, default=None, help="基线线速度，一般留空")
    ap.add_argument("--angular", type=float, default=None, help="基线角速度，一般留空")
    ap.add_argument("cmd", nargs="?", default=None)
    ap.add_argument("args", nargs="*")
    opts = ap.parse_args()

    if not opts.cmd:
        ap.print_help()
        sys.exit(1)

    ws = WS(opts.host, opts.port, timeout=opts.timeout)
    code = 0
    try:
        ws.connect()
        base = {
            "height": opts.height, "roll": opts.roll, "stable": opts.stable,
            "linear": opts.linear, "angular": opts.angular,
        }
        robot = Robot(ws, rate=opts.rate, base=base)
        if opts.height is None and opts.cmd in ("forward", "back", "left", "right", "jump", "joy", "stop", "seq"):
            print("[robot] 注意: 未用 --height 指定设备当前站高，将按默认 38mm 发送。"
                  "若设备实际站高不是 38mm，本命令会把它改到 38mm。", file=sys.stderr)
        run(robot, opts.cmd, opts.args, opts)
        time.sleep(0.15)
    except KeyboardInterrupt:
        code = 130
        print("\n[robot] 已中断")
    except Exception as e:
        print(f"[robot] 失败: {e}", file=sys.stderr)
        code = 1
    finally:
        try:
            ws.close()
        except Exception:
            pass
    sys.exit(code)


if __name__ == "__main__":
    main()