#!/usr/bin/env python3
"""山猫 S10（Deep Robotics）遥控通信模块。

仅依赖标准库。实现《山猫S10 软件开发指南》第 1 章中遥控所需的部分：

* 16 字节帧头 + JSON ASDU 的封装与解包（支持 TCP 流式拆包/粘包）
* mTLS 连接（TLS 1.3，客户端证书双向认证）
* 心跳保活与后台状态接收线程
* 控制指令：使用模式、运动状态、步态、轴指令、速度指令、
  灯语、自主充电、运动 SDK、休眠
* 状态回读：基础状态 / 运控状态 / 设备状态 / 异常状态 / 导航状态

限制：只处理单包 ASDU（ASDU 长度 <= 1000 字节），所有控制指令都在此范围内；
多包拼接（总长最大 65535）未实现。

示例::

    with S10Client(cert="certs/client_chain.crt",
                   key="certs/client.key",
                   ca="certs/root_ca.crt") as robot:
        robot.stand()
        robot.wait_motion_state({MotionState.RL}, timeout=20)
        robot.move(yaw=0.3)          # 轴指令，需持续下发
        robot.stop()
        robot.lie_down()
"""

from __future__ import annotations

import json
import logging
import socket
import ssl
import struct
import threading
import time
from dataclasses import dataclass, field
from typing import Any, Iterable

__all__ = [
    "S10Client",
    "S10Error",
    "S10Timeout",
    "S10CommandError",
    "MotionState",
    "Gait",
    "Mode",
    "Charge",
    "Frame",
    "build_frame",
    "parse_frames",
]

log = logging.getLogger("s10_robot")

# ---------------------------------------------------------------------------
# 协议常量
# ---------------------------------------------------------------------------

SYNC = b"\xeb\x91\xeb\x90"
HEADER_LEN = 16
FMT_XML = 0x00
FMT_JSON = 0x01
PROTO_VERSION = 0x01
MAX_ASDU = 1000

# Type 字段
T_MOTION = 0x00100001      # 运控：状态/步态/轴指令/速度指令，也用于运控状态上报
T_MODE = 0x00100002        # 使用模式/休眠/设备状态上报
T_NAV = 0x00100003         # 导航状态
T_CHARGE = 0x00100004      # 自主充电
T_LED = 0x00100005         # 灯语 / 运动 SDK
T_HEARTBEAT = 0x00100064   # 心跳 / 基础状态上报
T_FAULT = 0x0010007F       # 异常状态上报

# Command 字段
C_HEARTBEAT = 0x00000005
C_MODE_SET = 0x00500002
C_MOTION_STATE = 0x00200002
C_GAIT = 0x00300002
C_AXIS = 0x00100002
C_VELOCITY = 0x00110002
C_CHARGE = 0x00100002
C_LED = 0x00200002
C_SDK = 0x00300002
C_SLEEP = 0x00400002
C_STATUS = 0x00F00000      # 所有主动上报的 Command

# 上报类别
_STATUS_KEYS = {
    (T_HEARTBEAT, C_STATUS): "basic",
    (T_MOTION, C_STATUS): "motion",
    (T_MODE, C_STATUS): "device",
    (T_FAULT, C_STATUS): "fault",
    (T_NAV, C_STATUS): "nav",
}


class MotionState:
    """MotionState / MotionParam 取值。"""

    SOFT_ESTOP = -2
    DEFAULT = 0        # 初始 / 趴卧
    STANDING = 1
    JOINT_DAMP = 2     # 关节阻尼 / 软急停
    LYING = 4
    CALIBRATE = 5
    RL = 17
    DAMP_LYING = 0x1001


class Gait:
    """GaitParam / Gait 取值。"""

    BASIC = 0x1001          # 常规运动模式
    HIGH_PLATFORM = 0x1002  # 高台（仅在状态上报中出现）
    FLAT_NAV = 0x3002       # 导航运动模式
    STAIR_NAV = 0x3003      # 导航运动模式


class Mode:
    """ControlUsageMode 取值。"""

    NORMAL = 0
    NAV = 1
    ASSIST = 2


class Charge:
    """Charge 动作取值。"""

    STOP = 0
    START = 1
    RESET = 2  # 强制断电并清除错误，慎用


# ---------------------------------------------------------------------------
# 异常
# ---------------------------------------------------------------------------


class S10Error(RuntimeError):
    """通信或协议层错误。"""


class S10Timeout(S10Error):
    """等待响应或状态超时。"""


class S10CommandError(S10Error):
    """机器人返回非 0 ErrorCode。"""

    def __init__(self, code: int, message: str | None, frame: "Frame | None" = None):
        self.code = code
        self.message = message
        self.frame = frame
        super().__init__(f"ErrorCode=0x{code:04X} ({code}) {message or ''}".strip())


# ---------------------------------------------------------------------------
# 帧结构
# ---------------------------------------------------------------------------


@dataclass
class Frame:
    type: int
    command: int
    time: str
    items: dict[str, Any]
    report_id: int
    fmt: int
    seq: int
    version: int
    payload: bytes = b""

    @property
    def is_status(self) -> bool:
        return self.command == C_STATUS


def _to_int(value: Any) -> int:
    if isinstance(value, bool):
        return int(value)
    if isinstance(value, int):
        return value
    if isinstance(value, float):
        return int(value)
    if isinstance(value, str):
        text = value.strip().replace(" ", "")
        return int(text, 0)
    raise ValueError(f"无法转换为整数: {value!r}")


def build_frame(asdu: bytes, report_id: int = 0, seq: int = 0, fmt: int = FMT_JSON) -> bytes:
    """把 ASDU 字节封装成 16 字节头 + 载荷。"""
    if len(asdu) > MAX_ASDU:
        raise ValueError(f"单帧 ASDU 超长: {len(asdu)} > {MAX_ASDU}")
    header = SYNC
    header += struct.pack("<H", len(asdu))
    header += struct.pack("<H", report_id & 0xFFFF)
    header += bytes([fmt, seq & 0xFF, PROTO_VERSION])
    header += b"\x00" * 5
    assert len(header) == HEADER_LEN
    return header + asdu


def build_asdu(type_: int, command: int, items: dict[str, Any] | None = None,
               time_str: str | None = None) -> bytes:
    """构造 JSON 格式的 ASDU 载荷。"""
    body = {
        "PatrolDevice": {
            "Type": type_,
            "Command": command,
            "Time": time_str or time.strftime("%Y-%m-%d %H:%M:%S"),
            "Items": items or {},
        }
    }
    return json.dumps(body, separators=(",", ":"), ensure_ascii=False).encode("utf-8")


def _parse_payload(payload: bytes) -> tuple[int, int, str, dict[str, Any]]:
    doc = json.loads(payload.decode("utf-8"))
    body = doc.get("PatrolDevice", doc)
    if not isinstance(body, dict):
        raise ValueError("ASDU 根节点格式异常")
    type_ = _to_int(body.get("Type", -1))
    command = _to_int(body.get("Command", -1))
    time_str = str(body.get("Time", ""))
    items = body.get("Items") or {}
    if not isinstance(items, dict):
        raise ValueError("Items 不是对象")
    return type_, command, time_str, items


def parse_frames(buf: bytearray) -> list[Frame]:
    """从缓冲区中提取所有完整帧，剩余的半个帧留在 buf 里。

    遇到垃圾数据时通过搜索同步字符重新对齐。
    """
    frames: list[Frame] = []
    while True:
        idx = buf.find(SYNC)
        if idx < 0:
            # 保留尾部 3 字节，避免同步字符被切开
            if len(buf) > 3:
                del buf[: len(buf) - 3]
            return frames
        if idx:
            log.debug("丢弃 %d 字节非同步数据", idx)
            del buf[:idx]
        if len(buf) < HEADER_LEN:
            return frames

        length, report_id = struct.unpack_from("<HH", buf, 4)
        if length > MAX_ASDU * 65:  # 明显异常，跳过这个同步字符重新找
            log.warning("长度字段异常 (%d)，重新同步", length)
            del buf[:4]
            continue
        total = HEADER_LEN + length
        if len(buf) < total:
            return frames

        raw = bytes(buf[:total])
        del buf[:total]
        payload = raw[HEADER_LEN:]
        try:
            type_, command, time_str, items = _parse_payload(payload)
        except Exception as exc:  # 坏包不应中断接收
            log.warning("ASDU 解析失败: %s; payload=%r", exc, payload[:200])
            continue

        frames.append(
            Frame(
                type=type_,
                command=command,
                time=time_str,
                items=items,
                report_id=report_id,
                fmt=raw[8],
                seq=raw[9],
                version=raw[10],
                payload=payload,
            )
        )


# ---------------------------------------------------------------------------
# 客户端
# ---------------------------------------------------------------------------


@dataclass
class _Pending:
    type: int
    command: int
    event: threading.Event = field(default_factory=threading.Event)
    frame: Frame | None = None


class S10Client:
    """山猫 S10 遥控客户端。

    单条 TCP/TLS 连接内完成心跳、指令下发和状态接收。轴指令/速度指令要求
    2 秒内来自同一客户端，因此不要为每条指令重新建连。
    """

    def __init__(
        self,
        host: str = "10.21.33.103",
        port: int = 30003,
        *,
        cert: str = "certs/client_chain.crt",
        key: str = "certs/client.key",
        ca: str = "certs/root_ca.crt",
        verify: bool = False,
        server_name: str = "server.local",
        heartbeat_hz: float = 2.0,
        request_timeout: float = 5.0,
        connect_timeout: float = 8.0,
    ):
        self.host = host
        self.port = port
        self.cert = cert
        self.key = key
        self.ca = ca
        self.verify = verify
        self.server_name = server_name
        self.heartbeat_hz = heartbeat_hz
        self.request_timeout = request_timeout
        self.connect_timeout = connect_timeout

        self._sock: ssl.SSLSocket | None = None
        self._running = False
        self._wlock = threading.Lock()
        self._id_lock = threading.Lock()
        self._report_id = 0
        self._pending: dict[int, _Pending] = {}
        self._pending_lock = threading.Lock()

        self._cond = threading.Condition()
        self._status: dict[str, dict[str, Any]] = {
            "basic": {}, "motion": {}, "device": {}, "fault": {}, "nav": {}
        }
        self._status_ts: dict[str, float] = {}
        self._errors: dict[int, dict[str, Any]] = {}

        self._recv_thread: threading.Thread | None = None
        self._hb_thread: threading.Thread | None = None
        self._last_rx = 0.0
        self._last_hb = 0.0

    # -- 连接管理 ---------------------------------------------------------

    def connect(self) -> "S10Client":
        if self._sock is not None:
            return self
        ctx = ssl.create_default_context(ssl.Purpose.SERVER_AUTH, cafile=self.ca)
        ctx.check_hostname = False  # 服务端证书 CN=server.local，而连接用的是 IP
        ctx.verify_mode = ssl.CERT_REQUIRED if self.verify else ssl.CERT_NONE
        ctx.load_cert_chain(self.cert, self.key)

        raw = socket.create_connection((self.host, self.port), timeout=self.connect_timeout)
        raw.settimeout(0.5)
        try:
            sock = ctx.wrap_socket(raw, server_hostname=self.server_name)
        except Exception:
            raw.close()
            raise
        self._sock = sock
        if not self.verify:
            log.warning("未严格校验服务端证书链（--verify 可开启，部分 Python 版本会因根 CA 缺 keyUsage 失败）")
        log.info("已连接 %s:%s，%s %s", self.host, self.port,
                 sock.version(), sock.cipher()[0])

        self._running = True
        self._last_rx = time.monotonic()
        self._recv_thread = threading.Thread(target=self._recv_loop, name="s10-recv", daemon=True)
        self._hb_thread = threading.Thread(target=self._hb_loop, name="s10-hb", daemon=True)
        self._recv_thread.start()
        self._hb_thread.start()
        self.heartbeat()
        return self

    def close(self) -> None:
        self._running = False
        sock, self._sock = self._sock, None
        if sock is not None:
            try:
                sock.shutdown(socket.SHUT_RDWR)
            except (OSError, AttributeError):
                pass
            try:
                sock.close()
            except (OSError, AttributeError):
                pass
        for thread in (self._recv_thread, self._hb_thread):
            if thread is not None and thread.is_alive():
                thread.join(timeout=2.0)
        with self._pending_lock:
            for pending in self._pending.values():
                pending.event.set()
            self._pending.clear()

    def __enter__(self) -> "S10Client":
        return self.connect()

    def __exit__(self, exc_type, exc, tb) -> None:
        self.close()

    @property
    def connected(self) -> bool:
        return self._running and self._sock is not None

    # -- 收发 -------------------------------------------------------------

    def _next_report_id(self) -> int:
        with self._id_lock:
            rid = self._report_id
            self._report_id = (self._report_id + 1) & 0xFFFF
            return rid

    def _send_nowait(self, type_: int, command: int, items: dict[str, Any]) -> None:
        sock = self._sock
        if sock is None:
            raise S10Error("未连接")
        packet = build_frame(build_asdu(type_, command, items), self._next_report_id())
        with self._wlock:
            try:
                sock.sendall(packet)
            except OSError as exc:
                raise S10Error(f"发送失败: {exc}") from exc

    def _request(self, type_: int, command: int, items: dict[str, Any] | None = None,
                 timeout: float | None = None) -> dict[str, Any]:
        rid = self._next_report_id()
        pending = _Pending(type_, command)
        with self._pending_lock:
            self._pending[rid] = pending
        try:
            packet = build_frame(build_asdu(type_, command, items), rid)
            sock = self._sock
            if sock is None:
                raise S10Error("未连接")
            with self._wlock:
                sock.sendall(packet)
            if not pending.event.wait(timeout or self.request_timeout):
                raise S10Timeout(
                    f"等待响应超时 Type=0x{type_:08X} Command=0x{command:08X}"
                )
            frame = pending.frame
        finally:
            with self._pending_lock:
                self._pending.pop(rid, None)

        assert frame is not None
        code = frame.items.get("ErrorCode")
        if code is not None and int(code) != 0:
            raise S10CommandError(int(code), frame.items.get("ErrorMessage"), frame)
        return frame.items

    def _recv_loop(self) -> None:
        buf = bytearray()
        while self._running:
            sock = self._sock
            if sock is None:
                break
            try:
                chunk = sock.recv(65535)
            except socket.timeout:
                continue
            except (OSError, ssl.SSLError) as exc:
                if self._running:
                    log.warning("接收失败: %s", exc)
                break
            if not chunk:
                log.warning("对端关闭连接")
                break
            self._last_rx = time.monotonic()
            buf.extend(chunk)
            for frame in parse_frames(buf):
                self._dispatch(frame)
        self._running = False

    def _dispatch(self, frame: Frame) -> None:
        if not frame.is_status:
            with self._pending_lock:
                pending = self._pending.get(frame.report_id)
            if pending and pending.type == frame.type and pending.command == frame.command:
                pending.frame = frame
                pending.event.set()
                return

        key = _STATUS_KEYS.get((frame.type, frame.command))
        if key is None:
            log.debug("忽略未识别帧 Type=0x%08X Command=0x%08X", frame.type, frame.command)
            return

        with self._cond:
            if key == "fault":
                self._apply_faults(frame.items)
            elif key == "basic":
                # Items.BasicStatus.*，部分固件可能直接平铺
                inner = frame.items.get("BasicStatus")
                self._status[key] = dict(inner if isinstance(inner, dict) else frame.items)
            else:
                self._status[key] = dict(frame.items)
            self._status_ts[key] = time.monotonic()
            if key == "basic":
                state = self._status["basic"].get("MotionState")
                hes = self._status["basic"].get("HES")
                log.debug("基础状态 MotionState=%s Gait=%s HES=%s",
                          state, self._status["basic"].get("Gait"), hes)
            self._cond.notify_all()

    def _apply_faults(self, items: dict[str, Any]) -> None:
        errors = items.get("ErrorList")
        if isinstance(errors, dict):
            errors = [errors]
        if not isinstance(errors, list):
            self._status["fault"] = dict(items)
            return
        for err in errors:
            code = int(err.get("Code", 0))
            etype = int(err.get("Type", 1))
            if etype == 2:
                self._errors.pop(code, None)
            else:
                self._errors[code] = err
        self._status["fault"] = dict(items)
        self._status["fault"]["ActiveErrors"] = list(self._errors.values())

    def _hb_loop(self) -> None:
        period = 1.0 / max(self.heartbeat_hz, 0.2)
        while self._running:
            try:
                self.heartbeat()
            except S10Error as exc:
                if self._running:
                    log.warning("心跳发送失败: %s", exc)
                break
            time.sleep(period)

    # -- 状态读取 ---------------------------------------------------------

    @property
    def basic_status(self) -> dict[str, Any]:
        with self._cond:
            return dict(self._status["basic"])

    @property
    def motion_status(self) -> dict[str, Any]:
        """运控状态的 MotionStatus 部分（速度、姿态、关节角等）。"""
        with self._cond:
            items = self._status["motion"]
            inner = items.get("MotionStatus")
            return dict(inner if isinstance(inner, dict) else items)

    @property
    def motor_status(self) -> dict[str, Any]:
        with self._cond:
            inner = self._status["motion"].get("MotorStatus")
            return dict(inner) if isinstance(inner, dict) else {}

    @property
    def device_status(self) -> dict[str, Any]:
        with self._cond:
            return dict(self._status["device"])

    @property
    def active_errors(self) -> list[dict[str, Any]]:
        with self._cond:
            return [dict(e) for e in self._errors.values()]

    @property
    def status_age(self) -> float:
        """最近一次收到基础状态距今的秒数；尚未收到时为 inf。"""
        last = self._status_ts.get("basic")
        return float("inf") if last is None else time.monotonic() - last

    def wait_basic(self, field_name: str, values: Iterable[Any], timeout: float = 10.0,
                   check_estop: bool = True) -> Any:
        """等待基础状态中某字段取到指定值之一。"""
        wanted = set(values)

        def ready() -> bool:
            basic = self._status["basic"]
            if check_estop and basic.get("HES") == 1:
                raise S10Error("硬急停已触发（HES=1）")
            return basic.get(field_name) in wanted

        with self._cond:
            ok = self._cond.wait_for(ready, timeout)
            value = self._status["basic"].get(field_name)
        if not ok:
            raise S10Timeout(f"等待 {field_name} in {sorted(wanted)} 超时，当前={value!r}")
        return value

    def wait_motion_state(self, states: Iterable[int], timeout: float = 20.0,
                          check_estop: bool = True) -> int:
        return self.wait_basic("MotionState", states, timeout, check_estop)

    def wait_gait(self, gaits: Iterable[int], timeout: float = 10.0) -> int:
        return self.wait_basic("Gait", gaits, timeout)

    def wait_mode(self, modes: Iterable[int], timeout: float = 10.0) -> int:
        return self.wait_basic("ControlUsageMode", modes, timeout)

    # -- 控制指令 ---------------------------------------------------------

    def heartbeat(self) -> None:
        """发送一次心跳。机器人只向持续发心跳的 IP:端口 上报状态。"""
        self._send_nowait(T_HEARTBEAT, C_HEARTBEAT, {})

    def set_mode(self, mode: int) -> dict[str, Any]:
        """切换使用模式（Mode.NORMAL / NAV / ASSIST）。"""
        return self._request(T_MODE, C_MODE_SET, {"Mode": int(mode)})

    def set_motion_state(self, state: int) -> dict[str, Any]:
        """转换运动状态（见 MotionState）。软急停/空闲不可下发。"""
        return self._request(T_MOTION, C_MOTION_STATE, {"MotionParam": int(state)})

    def stand(self) -> dict[str, Any]:
        return self.set_motion_state(MotionState.STANDING)

    def lie_down(self) -> dict[str, Any]:
        return self.set_motion_state(MotionState.LYING)

    def damp(self) -> dict[str, Any]:
        """进入关节阻尼状态（可作为软急停后的安全态）。"""
        return self.set_motion_state(MotionState.JOINT_DAMP)

    def set_gait(self, gait: int) -> dict[str, Any]:
        """切换步态。必须在机器人完全停止后下发。"""
        return self._request(T_MOTION, C_GAIT, {"GaitParam": int(gait)})

    def move(self, x: float = 0.0, y: float = 0.0, z: float = 0.0,
             roll: float = 0.0, pitch: float = 0.0, yaw: float = 0.0) -> None:
        """轴指令：各轴为相对最大速度的比例，取值 [-1, 1]，无响应帧。

        仅常规/辅助模式有效；基础步态只响应 X、Y、Yaw。
        """
        values = (x, y, z, roll, pitch, yaw)
        if any(v < -1.0 or v > 1.0 for v in values):
            raise ValueError(f"轴指令参数必须位于 [-1, 1]: {values}")
        self._send_nowait(T_MOTION, C_AXIS, {
            "X": float(x), "Y": float(y), "Z": float(z),
            "Roll": float(roll), "Pitch": float(pitch), "Yaw": float(yaw),
        })

    def drive_velocity(self, x: float = 0.0, y: float = 0.0, z: float = 0.0,
                       roll: float = 0.0, pitch: float = 0.0, yaw: float = 0.0) -> None:
        """速度指令：物理量（m/s、rad/s），仅导航模式生效，无响应帧。"""
        if abs(x) > 1.67 or abs(y) > 0.4 or abs(yaw) > 1.0:
            raise ValueError(f"速度指令超限: X±1.67 Y±0.4 Yaw±1.0, got {x},{y},{yaw}")
        self._send_nowait(T_MOTION, C_VELOCITY, {
            "X": float(x), "Y": float(y), "Z": float(z),
            "Roll": float(roll), "Pitch": float(pitch), "Yaw": float(yaw),
        })

    def stop(self) -> None:
        """下发全 0 轴指令，让机器人停止移动。"""
        self.move()

    def set_led(self, head: dict[str, Any] | None = None,
                tail: dict[str, Any] | None = None,
                custom: bool = True) -> dict[str, Any]:
        """设置头灯/尾灯。每项形如 {"Type": 0, "Color": [1], "Cycle": 1}。"""
        leds = []
        for spec in (head, tail):
            if spec is not None:
                leds.append({
                    "Type": int(spec.get("Type", 0)),
                    "Color": list(spec.get("Color", [1])),
                    "Cycle": int(spec.get("Cycle", 1)),
                })
        return self._request(T_LED, C_LED, {"CustomMode": bool(custom), "Led": leds})

    def charge(self, action: int) -> dict[str, Any]:
        """自主充电：0 结束 / 1 开始 / 2 清除状态（慎用）。"""
        return self._request(T_CHARGE, C_CHARGE, {"Charge": int(action)})

    def sdk_mode(self, enable: bool, frequency: int = 100) -> dict[str, Any]:
        """开启/关闭运动 SDK 模式。频率需在 [1,200] 且能整除 1000。"""
        if enable and not (1 <= frequency <= 200 and 1000 % frequency == 0):
            raise ValueError(f"Frequency 非法: {frequency}")
        return self._request(T_LED, C_SDK,
                             {"SDKEnable": bool(enable), "Frequency": int(frequency)})

    def sleep_mode(self, sleep: bool, auto: bool = True, minutes: int = 5) -> dict[str, Any]:
        """设置休眠/自动休眠。minutes 为心跳断联后的等待时间。"""
        return self._request(T_MODE, C_SLEEP,
                             {"Sleep": bool(sleep), "Auto": bool(auto), "Time": int(minutes)})