2

import atexit
import base64
import json
import math
import os
import queue
import shlex
import threading
import time
import uuid

try:
    import paramiko
except ImportError as exc:
    raise ImportError("缺少 paramiko，请先执行：pip install paramiko") from exc

_PROTOCOL_PREFIX = "@@XMWROBOT@@"

_REMOTE_BRIDGE_SOURCE = r'''
import glob
import json
import os
import select
import sys
import time

PREFIX = "@@XMWROBOT@@"
PERIOD = 0.01


def emit(data):
    print(PREFIX + json.dumps(data, ensure_ascii=False, separators=(",", ":")), flush=True)


def find_sdk_root():
    home = os.path.expanduser("~")
    candidates = []
    configured = os.environ.get("XMWROBOT_SDK_ROOT", "").strip()
    if configured:
        candidates.append(configured)
    candidates.extend([
        os.path.join(home, "work"),
        os.path.join(home, "noetix_sdk_bumi-main"),
        os.path.join(home, "noetix_sdk_bumi"),
        os.path.join(home, "work", "noetix_sdk_bumi-main"),
        os.path.join(home, "work", "noetix_sdk_bumi"),
    ])
    candidates.extend(glob.glob(os.path.join(home, "**", "noetix_sdk_bumi*"), recursive=True))
    seen = set()
    for root in candidates:
        root = os.path.abspath(os.path.expanduser(root))
        if root in seen:
            continue
        seen.add(root)
        if os.path.isfile(os.path.join(root, "config", "dds.xml")):
            return root
    raise RuntimeError("没有找到新版 Bumi SDK")


def load_sdk():
    root = find_sdk_root()
    os.chdir(root)
    os.environ["CYCLONEDDS_URI"] = "file://" + os.path.join(root, "config", "dds.xml")
    search_dirs = [
        root,
        os.path.join(root, "build"),
        os.path.join(root, "build", "lib"),
        os.path.join(root, "examples_py"),
    ]
    for so_path in glob.glob(os.path.join(root, "**", "highcontrol_py*.so"), recursive=True):
        search_dirs.append(os.path.dirname(so_path))
    for item in search_dirs:
        if os.path.isdir(item) and item not in sys.path:
            sys.path.insert(0, item)
    from highcontrol_py import ControlCmd, HighController
    controller = HighController.instance()
    controller.init()
    return controller, ControlCmd


def stop(controller, ControlCmd, count=8):
    for _ in range(max(1, int(count))):
        controller.publish_cmd(0.0, 0.0, 0.0, ControlCmd.DEFAULT, 0)
        time.sleep(PERIOD)


def move_forward(controller, ControlCmd, speed, duration):
    speed = max(0.1, min(0.5, float(speed)))
    duration = max(0.1, min(8.0, float(duration)))
    started = time.monotonic()
    while True:
        elapsed = time.monotonic() - started
        if elapsed >= duration:
            break
        remaining = duration - elapsed
        if elapsed < 0.15:
            factor = max(0.45, min(1.0, 0.45 + elapsed / 0.15 * 0.55))
        elif remaining < 0.22:
            factor = max(0.0, remaining / 0.22)
        else:
            factor = 1.0
        controller.publish_cmd(speed * factor, 0.0, 0.0, ControlCmd.DEFAULT, 0)
        time.sleep(PERIOD)
    stop(controller, ControlCmd, 10)


try:
    controller, ControlCmd = load_sdk()
    emit({"event": "ready", "ok": True})
except Exception as exc:
    emit({"event": "ready", "ok": False, "error": str(exc)})
    raise

pending_action = None
pending_index = 0
running = True

while running:
    started = time.monotonic()
    try:
        readable, _, _ = select.select([sys.stdin], [], [], 0)
        if readable:
            line = sys.stdin.readline()
            if line == "":
                running = False
            else:
                request_id = ""
                try:
                    message = json.loads(line)
                    request_id = str(message.get("id", ""))
                    command = str(message.get("command", "")).lower()

                    if command == "move":
                        move_forward(
                            controller,
                            ControlCmd,
                            message.get("speed", 0.42),
                            message.get("duration", 0.8),
                        )
                        emit({"id": request_id, "ok": True})

                    elif command == "action":
                        action_name = str(message.get("name", "")).upper()
                        if not hasattr(ControlCmd, action_name):
                            raise ValueError("SDK 不支持动作：" + action_name)
                        pending_action = getattr(ControlCmd, action_name)
                        pending_index = max(0, int(message.get("index", 0)))
                        emit({"id": request_id, "ok": True})

                    elif command == "pose":
                        imu = controller.get_imu_data()
                        emit({
                            "id": request_id,
                            "ok": True,
                            "result": {
                                "linear_acc": [float(imu.linear_acc[i]) for i in range(3)],
                                "angular_vel": [float(imu.angular_vel[i]) for i in range(3)],
                            },
                        })

                    elif command == "shutdown":
                        stop(controller, ControlCmd, 10)
                        emit({"id": request_id, "ok": True})
                        running = False

                    elif command == "ping":
                        emit({"id": request_id, "ok": True})

                    else:
                        raise ValueError("未知命令：" + command)

                except Exception as exc:
                    emit({"id": request_id, "ok": False, "error": str(exc)})

        if running:
            action = pending_action if pending_action is not None else ControlCmd.DEFAULT
            index = pending_index if pending_action is not None else 0
            controller.publish_cmd(0.0, 0.0, 0.0, action, index)
            pending_action = None
            pending_index = 0

    except Exception as exc:
        emit({"event": "runtime_error", "ok": False, "error": str(exc)})
        pending_action = None
        pending_index = 0
        try:
            controller.publish_cmd(0.0, 0.0, 0.0, ControlCmd.DEFAULT, 0)
        except Exception:
            pass

    elapsed = time.monotonic() - started
    if elapsed < PERIOD:
        time.sleep(PERIOD - elapsed)

try:
    stop(controller, ControlCmd, 10)
except Exception:
    pass
'''


class RobotError(RuntimeError):
    pass


class RobotConnectionError(RobotError):
    pass


class RobotCommandError(RobotError):
    pass


class _Robot:
    _ACTIONS = {
        "wave": "SWING",
        "swing": "SWING",
        "挥手": "SWING",
        "handshake": "SHAKE",
        "shake": "SHAKE",
        "握手": "SHAKE",
        "cheer": "CHEER",
        "欢呼": "CHEER",
        "tear": "TEAR",
        "wipe_tears": "TEAR",
        "擦眼泪": "TEAR",
        "dance": "DANCE",
        "dance1": "DANCE",
        "dance2": "DANCE1",
        "dance3": "DANCE2",
        "stand_up": "FALLTOSTAND",
        "falltostand": "FALLTOSTAND",
        "倒地起身": "FALLTOSTAND",
        "lie_down": "STANDTOFALL",
        "standtofall": "STANDTOFALL",
        "站立倒地": "STANDTOFALL",
    }

    _WAITS = {
        "SWING": 0.8,
        "SHAKE": 1.0,
        "CHEER": 4.0,
        "TEAR": 1.0,
        "DANCE": 1.0,
        "DANCE1": 1.0,
        "DANCE2": 1.0,
        "FALLTOSTAND": 0.5,
        "STANDTOFALL": 0.5,
    }

    def __init__(self, ip, ssh_port=None, username=None, password=None):
        self.ip = str(ip).strip()
        if not self.ip:
            raise ValueError("机器人 IP 不能为空")
        self.ssh_port = int(ssh_port if ssh_port is not None else os.getenv("XMWROBOT_SSH_PORT", "2222"))
        self.username = username if username is not None else os.getenv("XMWROBOT_SSH_USER", "noetix")
        self.password = password if password is not None else os.getenv("XMWROBOT_SSH_PASSWORD", "bumi")
        self.forward_speed = float(os.getenv("XMWROBOT_FORWARD_SPEED", "0.42"))
        self.seconds_per_step = float(os.getenv("XMWROBOT_SECONDS_PER_STEP", "0.55"))
        self._ssh = None
        self._bridge_stdin = None
        self._bridge_stdout = None
        self._bridge_stderr = None
        self._closed = False
        self._request_lock = threading.RLock()
        self._pending_lock = threading.RLock()
        self._pending = {}
        self._ready_queue = queue.Queue(maxsize=1)
        self._bridge_errors = []
        self._connect()
        self._ensure_standing()
        atexit.register(self.close)

    def _connect(self):
        client = paramiko.SSHClient()
        client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
        try:
            client.connect(
                hostname=self.ip,
                port=self.ssh_port,
                username=self.username,
                password=self.password,
                timeout=8,
                banner_timeout=8,
                auth_timeout=8,
                look_for_keys=False,
                allow_agent=False,
            )
            transport = client.get_transport()
            if transport is not None:
                transport.set_keepalive(15)
        except Exception as exc:
            client.close()
            raise RobotConnectionError(f"无法连接机器人 {self.ip}:{self.ssh_port}：{exc}") from exc
        self._ssh = client
        self._stop_conflicts()
        self._start_bridge()

    def _stop_conflicts(self):
        script = """
systemctl --user stop bumi-onboard-voice.service >/dev/null 2>&1 || true
for pattern in '(^|/)operation([[:space:]]|\\.py)' '\\.bumi_qwen_voice_highcontrol\\.py' 'Bumi_机器人本体千问语音控制'; do
  pgrep -u "$(id -u)" -f "$pattern" 2>/dev/null | while read -r pid; do
    [ "$pid" = "$$" ] || kill -TERM "$pid" >/dev/null 2>&1 || true
  done
done
sleep 0.3
"""
        try:
            _, stdout, stderr = self._ssh.exec_command("bash -lc " + shlex.quote(script), timeout=5)
            stdout.channel.recv_exit_status()
            stderr.read()
        except Exception:
            pass

    def _start_bridge(self):
        encoded = base64.b64encode(_REMOTE_BRIDGE_SOURCE.encode("utf-8")).decode("ascii")
        command = "python3 -u -c " + shlex.quote("import base64;exec(base64.b64decode('" + encoded + "').decode('utf-8'))")
        try:
            stdin, stdout, stderr = self._ssh.exec_command(command, get_pty=False)
        except Exception as exc:
            raise RobotConnectionError(f"无法启动机器人控制桥：{exc}") from exc
        self._bridge_stdin = stdin
        self._bridge_stdout = stdout
        self._bridge_stderr = stderr
        threading.Thread(target=self._bridge_reader, daemon=True).start()
        threading.Thread(target=self._bridge_stderr_reader, daemon=True).start()
        try:
            ready = self._ready_queue.get(timeout=15)
        except queue.Empty as exc:
            details = "；".join(self._bridge_errors[-3:])
            raise RobotConnectionError("机器人 SDK 初始化超时" + ("：" + details if details else "")) from exc
        if not ready.get("ok"):
            raise RobotConnectionError("机器人 SDK 初始化失败：" + str(ready.get("error", "未知错误")))

    def _bridge_reader(self):
        try:
            for raw_line in iter(self._bridge_stdout.readline, ""):
                line = raw_line.strip()
                if not line.startswith(_PROTOCOL_PREFIX):
                    continue
                try:
                    message = json.loads(line[len(_PROTOCOL_PREFIX):])
                except Exception:
                    continue
                if message.get("event") == "ready":
                    try:
                        self._ready_queue.put_nowait(message)
                    except queue.Full:
                        pass
                    continue
                request_id = str(message.get("id", ""))
                with self._pending_lock:
                    waiter = self._pending.get(request_id)
                if waiter is not None:
                    try:
                        waiter.put_nowait(message)
                    except queue.Full:
                        pass
        except Exception as exc:
            self._bridge_errors.append(str(exc))

    def _bridge_stderr_reader(self):
        try:
            for raw_line in iter(self._bridge_stderr.readline, ""):
                line = raw_line.strip()
                if line:
                    self._bridge_errors.append(line)
                    if len(self._bridge_errors) > 20:
                        del self._bridge_errors[:-20]
        except Exception:
            pass

    def _request(self, command, timeout=5.0, **payload):
        if self._closed or self._bridge_stdin is None:
            raise RobotConnectionError("机器人连接已经关闭")
        request_id = uuid.uuid4().hex
        waiter = queue.Queue(maxsize=1)
        with self._pending_lock:
            self._pending[request_id] = waiter
        message = {"id": request_id, "command": command}
        message.update(payload)
        try:
            with self._request_lock:
                self._bridge_stdin.write(json.dumps(message, ensure_ascii=False, separators=(",", ":")) + "\n")
                self._bridge_stdin.flush()
            try:
                response = waiter.get(timeout=timeout)
            except queue.Empty as exc:
                details = "；".join(self._bridge_errors[-3:])
                raise RobotCommandError(f"命令 {command} 等待超时" + ("：" + details if details else "")) from exc
        finally:
            with self._pending_lock:
                self._pending.pop(request_id, None)
        if not response.get("ok"):
            raise RobotCommandError(str(response.get("error", "机器人执行失败")))
        return response.get("result", {})

    def _pose_state(self):
        try:
            data = self._request("pose", timeout=4.0)
            linear_acc = [float(value) for value in data.get("linear_acc", [])]
            angular_vel = [float(value) for value in data.get("angular_vel", [])]
        except Exception:
            return "unknown"
        if len(linear_acc) != 3 or not all(math.isfinite(value) for value in linear_acc):
            return "unknown"
        if len(angular_vel) != 3 or not all(math.isfinite(value) for value in angular_vel):
            angular_vel = [0.0, 0.0, 0.0]
        acceleration_norm = math.sqrt(sum(value * value for value in linear_acc))
        angular_speed = math.sqrt(sum(value * value for value in angular_vel))
        if acceleration_norm < 0.1 or acceleration_norm > 100.0:
            return "unknown"
        vertical_ratio = abs(linear_acc[2]) / acceleration_norm
        if angular_speed > 1.5:
            return "unknown"
        if vertical_ratio >= 0.72:
            return "standing"
        if vertical_ratio <= 0.52:
            return "fallen"
        return "unknown"

    def _confirmed_pose_state(self, attempts=10, delay=0.3):
        standing = 0
        fallen = 0
        for _ in range(attempts):
            current = self._pose_state()
            if current == "standing":
                standing += 1
                fallen = 0
                if standing >= 2:
                    return "standing"
            elif current == "fallen":
                fallen += 1
                standing = 0
                if fallen >= 2:
                    return "fallen"
            else:
                standing = 0
                fallen = 0
            time.sleep(delay)
        return "unknown"

    def _ensure_standing(self):
        time.sleep(0.7)
        current = self._confirmed_pose_state(attempts=12, delay=0.25)
        if current == "standing":
            return
        if current == "unknown":
            current = self._confirmed_pose_state(attempts=12, delay=0.25)
            if current == "standing":
                return
        if current != "fallen":
            raise RobotError("无法确认机器人姿态，请让机器人保持静止后重新运行")
        try:
            import tkinter as tk
        except Exception as exc:
            raise RobotError("机器人处于倒地状态，当前电脑无法打开起身窗口") from exc
        root = tk.Tk()
        root.title("步数争夺战")
        root.geometry("360x180")
        root.resizable(False, False)
        status = tk.StringVar(value="机器人倒地，请点击起身")
        tk.Label(root, textvariable=status, font=("Microsoft YaHei", 12), pady=28).pack(fill="x")
        button = tk.Button(root, text="倒地起身", font=("Microsoft YaHei", 12), width=14)
        button.pack(pady=(0, 20))
        state = {"busy": False, "closed": False}

        def close_window():
            state["closed"] = True
            try:
                root.destroy()
            except Exception:
                pass

        def check_pose():
            if state["closed"] or state["busy"]:
                return
            current_state = self._confirmed_pose_state(attempts=3, delay=0.15)
            if current_state == "standing":
                close_window()
                return
            if current_state == "fallen":
                status.set("机器人倒地，请点击起身")
                button.config(state="normal")
            else:
                status.set("正在确认机器人姿态……")
                button.config(state="disabled")
            root.after(500, check_pose)

        def stand_up():
            if state["busy"]:
                return
            state["busy"] = True
            button.config(state="disabled")
            status.set("机器人正在起身……")

            def worker():
                error = None
                try:
                    self._request("action", timeout=5.0, name="FALLTOSTAND")
                    deadline = time.monotonic() + 14.0
                    standing = 0
                    while time.monotonic() < deadline and not state["closed"]:
                        time.sleep(0.5)
                        if self._pose_state() == "standing":
                            standing += 1
                            if standing >= 2:
                                root.after(0, close_window)
                                return
                        else:
                            standing = 0
                except Exception as exc:
                    error = str(exc)

                def finish():
                    state["busy"] = False
                    if state["closed"]:
                        return
                    status.set("起身失败，请重试" if error else "尚未确认站立，请重试")
                    button.config(state="normal")

                root.after(0, finish)

            threading.Thread(target=worker, daemon=True).start()

        button.config(command=stand_up, state="normal")
        root.protocol("WM_DELETE_WINDOW", lambda: None)
        root.after(300, check_pose)
        root.mainloop()

    def forward(self, steps=1):
        if isinstance(steps, bool) or not isinstance(steps, int):
            return
        if steps not in (1, 2, 3):
            return
        duration = steps * self.seconds_per_step + 0.37
        self._request(
            "move",
            timeout=duration + 4.0,
            speed=self.forward_speed,
            duration=duration,
        )

    def action(self, name):
        key = str(name).strip().lower()
        command = self._ACTIONS.get(key)
        if command is None:
            raise ValueError("不支持的动作：" + str(name))
        self._request("action", name=command)
        time.sleep(self._WAITS.get(command, 0.5))
        if command == "CHEER":
            try:
                self._request("action", name="WALK")
                time.sleep(0.8)
            except Exception:
                pass

    def close(self):
        if self._closed:
            return
        try:
            if self._bridge_stdin is not None:
                self._request("shutdown", timeout=2.0)
        except Exception:
            pass
        self._closed = True
        try:
            if self._ssh is not None:
                self._ssh.close()
        except Exception:
            pass


def connect(ip, ssh_port=None, username=None, password=None):
    return _Robot(ip, ssh_port=ssh_port, username=username, password=password)


__all__ = ["connect", "RobotError", "RobotConnectionError", "RobotCommandError"]