#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import time
import struct
import threading
from dataclasses import dataclass, field
from typing import Dict, List, Tuple
import can


TORQUE_CMD_ID = 0x280


def _clamp_int16(x: int) -> int:
    return max(-32768, min(32767, int(x)))


def _pack_i16_le(x: int) -> Tuple[int, int]:
    x = _clamp_int16(x)
    u = x & 0xFFFF
    return (u & 0xFF, (u >> 8) & 0xFF)


def build_torque_broadcast(t1: int, t2: int, t3: int, t4: int) -> can.Message:
    d0, d1 = _pack_i16_le(t1)
    d2, d3 = _pack_i16_le(t2)
    d4, d5 = _pack_i16_le(t3)
    d6, d7 = _pack_i16_le(t4)
    return can.Message(arbitration_id=TORQUE_CMD_ID, is_extended_id=False,
                       data=[d0, d1, d2, d3, d4, d5, d6, d7])


def _arb_id_single(motor_id: int) -> int:
    return 0x140 + int(motor_id)


def dump_frame(mid: int, fr):
    if fr.t <= 0:
        print(f"M{mid}: (no reply yet)")
        return

    d = fr.data
    arb = fr.arb
    cmd = d[0]

    if cmd == 0xA1:
        temperature = struct.unpack("<b", d[1:2])[0]
        iq_or_power = struct.unpack("<h", d[2:4])[0]
        speed_dps = struct.unpack("<h", d[4:6])[0]
        encoder = struct.unpack("<H", d[6:8])[0]
        print(f"M{mid}: arb=0x{arb:X}  T={temperature}C  iq/p={iq_or_power}  spd={speed_dps}dps  enc={encoder}")

    else:
        # 未知回包：先原样打印，后续按V2.36对应命令回包格式解析
        print(f"M{mid}: arb=0x{arb:X} cmd=0x{cmd:02X} raw={[hex(x) for x in d]}")


@dataclass
class RxFrame:
    t: float = 0.0
    arb: int = 0
    data: bytes = b""

class RxCache(can.Listener):
    def __init__(self, motor_ids: List[int]):
        super().__init__()
        self.motor_ids = set(motor_ids)
        self._lock = threading.Lock()
        self.latest: Dict[int, RxFrame] = {i: RxFrame() for i in motor_ids}

    def on_message_received(self, msg: can.Message):
        arb = msg.arbitration_id
        if arb < 0x141 or arb > 0x160:
            return
        mid = arb - 0x140
        if mid not in self.motor_ids:
            return
        with self._lock:
            self.latest[mid] = RxFrame(time.time(), arb, bytes(msg.data))

    def snapshot(self) -> Dict[int, RxFrame]:
        with self._lock:
            return {k: v for k, v in self.latest.items()}


def main():
    motor_ids = [1,2,3,4]
    bus = can.Bus(interface="canalystii", channel=0, bitrate=1000000)

    # 只收 1~4 的回包（0x141~0x144）
    try:
        bus.set_filters([{"can_id": _arb_id_single(i), "can_mask": 0x7FF, "extended": False} for i in motor_ids])
    except Exception:
        pass

    cache = RxCache(motor_ids)
    notifier = can.Notifier(bus, [cache], timeout=0.01)

    hz = 500
    dt = 1.0 / hz
    k = 0
    next_tick = time.perf_counter()

    try:
        while True:
            # 500Hz 广播力矩（示例：全 0，你可替换成控制输出）
            bus.send(build_torque_broadcast(30, 30, 30, 30))

            # 每 100ms 打印一次收到的自动回包
            if k % int(hz * 0.1) == 0:
                snap = cache.snapshot()
                print("---- RX snapshot ----")
                for mid in motor_ids:
                    fr = snap[mid]
                    if fr.t <= 0:
                        print(f"M{mid}: (no reply yet)")
                    else:
                        dump_frame(mid, fr)

            k += 1

            # 固定周期
            next_tick += dt
            now = time.perf_counter()
            if next_tick > now:
                time.sleep(next_tick - now)
            else:
                next_tick = now

    except KeyboardInterrupt:
        pass
    finally:
        notifier.stop()
        bus.send(build_torque_broadcast(0, 0, 0, 0))
        bus.shutdown()

if __name__ == "__main__":
    main()