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

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

import can

TORQUE_CMD_ID = 0x280
CMD_A1 = 0xA1


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 parse_a1(data8: bytes) -> Dict:
    # A1 回包（你实测）
    # [0]=0xA1 [1]=temp(int8) [2..3]=iq/power(int16 LE) [4..5]=speed(int16 LE, dps) [6..7]=encoder(uint16 LE)
    if len(data8) != 8 or data8[0] != CMD_A1:
        raise ValueError("not A1")
    temperature = struct.unpack("<b", data8[1:2])[0]
    iq_or_power = struct.unpack("<h", data8[2:4])[0]
    speed_dps = struct.unpack("<h", data8[4:6])[0]
    encoder = struct.unpack("<H", data8[6:8])[0]
    return dict(
        temperature_C=temperature,
        iq_or_power_raw=iq_or_power,
        speed_dps=speed_dps,
        encoder=encoder,
    )


def dump_frame(tag: str, mid: int, fr) -> str:
    if fr.t <= 0:
        return f"{tag} M{mid}: (no reply yet)"
    d = fr.data
    arb = fr.arb
    cmd = d[0]
    if cmd == CMD_A1:
        s = parse_a1(d)
        return (f"{tag} M{mid}: arb=0x{arb:X}  "
                f"T={s['temperature_C']}C  iq/p={s['iq_or_power_raw']}  "
                f"spd={s['speed_dps']}dps  enc={s['encoder']}")
    return f"{tag} 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 RxCache4Ch(can.Listener):
    """
    一个 Notifier 监听一个 “设备Bus(包含两个channel)”，按 msg.channel 分流到 ch0/ch1 的缓存。
    """
    def __init__(self, ch_to_motor_ids: Dict[int, List[int]]):
        super().__init__()
        self.ch_to_motor_ids = {int(ch): set(mids) for ch, mids in ch_to_motor_ids.items()}
        self._lock = threading.Lock()
        self.latest: Dict[Tuple[int, int], RxFrame] = {}
        for ch, mids in self.ch_to_motor_ids.items():
            for mid in mids:
                self.latest[(ch, mid)] = RxFrame()

    def on_message_received(self, msg: can.Message):
        # python-can Message.channel：对 canalystii 多通道会填 0/1
        ch = getattr(msg, "channel", None)
        if ch is None:
            # 没有 channel 信息就无法分流，直接忽略或当作 ch0
            ch = 0
        try:
            ch = int(ch)
        except Exception:
            ch = 0

        if ch not in self.ch_to_motor_ids:
            return

        arb = msg.arbitration_id
        if arb < 0x141 or arb > 0x160:
            return
        mid = arb - 0x140
        if mid not in self.ch_to_motor_ids[ch]:
            return

        with self._lock:
            self.latest[(ch, mid)] = RxFrame(time.time(), arb, bytes(msg.data))

    def snapshot(self) -> Dict[Tuple[int, int], RxFrame]:
        with self._lock:
            return dict(self.latest)


class SenderThread(threading.Thread):
    """
    每个“逻辑channel(共4个)”一个发送线程：
    - 等主线程 tick_event
    - barrier 同步放行
    - 加 bus_lock 后 send（同一个设备的两个channel共享一个bus对象，必须加锁）
    """
    def __init__(
        self,
        name: str,
        bus: can.BusABC,
        bus_lock: threading.Lock,
        channel_index: int,         # 0 or 1
        tick_event: threading.Event,
        barrier: threading.Barrier,
        stop_event: threading.Event,
    ):
        super().__init__(name=name)
        self.bus = bus
        self.bus_lock = bus_lock
        self.channel_index = int(channel_index)
        self.tick_event = tick_event
        self.barrier = barrier
        self.stop_event = stop_event
        self.torque = (0, 0, 0, 0)

    def run(self):
        while not self.stop_event.is_set():
            if not self.tick_event.wait(timeout=0.5):
                continue
            if self.stop_event.is_set():
                break

            try:
                self.barrier.wait(timeout=0.05)
            except threading.BrokenBarrierError:
                continue

            t1, t2, t3, t4 = self.torque
            msg = build_torque_broadcast(t1, t2, t3, t4)
            msg.channel = self.channel_index  # ★ 指定发到该设备的哪个channel

            try:
                with self.bus_lock:
                    self.bus.send(msg)
            except can.CanError:
                pass

            # 等主线程清 tick_event，避免同tick重复发
            while self.tick_event.is_set() and not self.stop_event.is_set():
                time.sleep(0)


def open_device_bus(device_idx: int, bitrate: int) -> can.BusABC:
    """
    每块 CANalyst-II 只打开一次 Bus，channel=[0,1]
    """
    # canalystii 后端一般支持 device=0/1；如果你环境不支持会 TypeError，这里做兼容
    try:
        return can.Bus(interface="canalystii", channel=[0, 1], bitrate=bitrate, device=device_idx)
    except TypeError:
        # 某些版本 canalystii 不暴露 device 参数（只能看到一块设备）
        return can.Bus(interface="canalystii", channel=[0, 1], bitrate=bitrate)


def main():
    bitrate = 1000000
    hz = 500
    dt = 1.0 / hz

    # ========= 你需要按实际接线修改：每个“设备-channel”上有哪些电机ID =========
    # 逻辑4路：dev0-ch0, dev0-ch1, dev1-ch0, dev1-ch1
    motor_map = {
        ("dev0", 0): [1, 2, 3, 4],
        ("dev0", 1): [1, 2, 3, 4],
        ("dev1", 0): [1, 2, 3, 4],
        ("dev1", 1): [1, 2, 3, 4],
    }

    # ========= 打开两块设备（每块一次性打开2路channel） =========
    bus0 = open_device_bus(device_idx=0, bitrate=bitrate)
    bus1 = open_device_bus(device_idx=1, bitrate=bitrate)

    # 过滤：只收我们关心的 arb_id（两设备相同过滤即可）
    def set_filters_for(bus: can.BusABC, all_motor_ids: List[int]):
        try:
            bus.set_filters([{"can_id": _arb_id_single(i), "can_mask": 0x7FF, "extended": False} for i in all_motor_ids])
        except Exception:
            pass

    all_ids_dev0 = sorted(set(motor_map[("dev0", 0)] + motor_map[("dev0", 1)]))
    all_ids_dev1 = sorted(set(motor_map[("dev1", 0)] + motor_map[("dev1", 1)]))
    set_filters_for(bus0, all_ids_dev0)
    set_filters_for(bus1, all_ids_dev1)

    # ========= 接收缓存：每个设备一个 Notifier，一个 cache（按 channel 分流） =========
    cache0 = RxCache4Ch({0: motor_map[("dev0", 0)], 1: motor_map[("dev0", 1)]})
    cache1 = RxCache4Ch({0: motor_map[("dev1", 0)], 1: motor_map[("dev1", 1)]})
    notifier0 = can.Notifier(bus0, [cache0], timeout=0.01)
    notifier1 = can.Notifier(bus1, [cache1], timeout=0.01)

    # ========= 发送线程：4个逻辑channel =========
    stop_event = threading.Event()
    tick_event = threading.Event()
    barrier = threading.Barrier(4 + 1)  # 4个sender + 主线程

    bus0_lock = threading.Lock()
    bus1_lock = threading.Lock()

    sender_dev0_ch0 = SenderThread("sender-dev0-ch0", bus0, bus0_lock, 0, tick_event, barrier, stop_event)
    sender_dev0_ch1 = SenderThread("sender-dev0-ch1", bus0, bus0_lock, 1, tick_event, barrier, stop_event)
    sender_dev1_ch0 = SenderThread("sender-dev1-ch0", bus1, bus1_lock, 0, tick_event, barrier, stop_event)
    sender_dev1_ch1 = SenderThread("sender-dev1-ch1", bus1, bus1_lock, 1, tick_event, barrier, stop_event)

    senders = [sender_dev0_ch0, sender_dev0_ch1, sender_dev1_ch0, sender_dev1_ch1]
    for th in senders:
        th.start()

    # ========= 你的控制输出：这里演示固定 torque =========
    def torque_for(tag: str, ch: int) -> Tuple[int, int, int, int]:
        # TODO：替换为你的控制算法输出
        return (30, 30, 30, 30)

    k = 0
    next_tick = time.perf_counter()

    try:
        while True:
            # 1) 写入4路 torque（主线程写共享变量）
            sender_dev0_ch0.torque = torque_for("dev0", 0)
            sender_dev0_ch1.torque = torque_for("dev0", 1)
            sender_dev1_ch0.torque = torque_for("dev1", 0)
            sender_dev1_ch1.torque = torque_for("dev1", 1)

            # 2) 同步tick发布 + barrier放行（尽量同时发送）
            tick_event.set()
            try:
                barrier.wait(timeout=0.05)
            except threading.BrokenBarrierError:
                pass
            tick_event.clear()

            # 3) 打印（每100ms一次）
            if k % int(hz * 0.1) == 0:
                print(f"\n==== tick {k}  t={time.time():.3f} ====")

                snap0 = cache0.snapshot()
                print("[dev0-ch0]")
                for mid in motor_map[("dev0", 0)]:
                    print("  " + dump_frame("[dev0-ch0]", mid, snap0[(0, mid)]))
                print("[dev0-ch1]")
                for mid in motor_map[("dev0", 1)]:
                    print("  " + dump_frame("[dev0-ch1]", mid, snap0[(1, mid)]))

                snap1 = cache1.snapshot()
                print("[dev1-ch0]")
                for mid in motor_map[("dev1", 0)]:
                    print("  " + dump_frame("[dev1-ch0]", mid, snap1[(0, mid)]))
                print("[dev1-ch1]")
                for mid in motor_map[("dev1", 1)]:
                    print("  " + dump_frame("[dev1-ch1]", mid, snap1[(1, mid)]))

            k += 1

            # 4) 500Hz 固定周期（绝对时间推进）
            next_tick += dt
            now = time.perf_counter()
            if next_tick > now:
                time.sleep(next_tick - now)
            else:
                next_tick = now

    except KeyboardInterrupt:
        pass
    finally:
        # 停止发送线程
        stop_event.set()
        tick_event.set()
        for th in senders:
            th.join(timeout=0.5)

        # 停止 Notifier（避免你之前那种 interpreter shutdown 崩溃）
        try:
            notifier0.stop()
            notifier1.stop()
        except Exception:
            pass

        # 安全清零力矩（给每个设备两个channel都发一次）
        try:
            for ch in [0, 1]:
                m = build_torque_broadcast(0, 0, 0, 0)
                m.channel = ch
                with bus0_lock:
                    bus0.send(m)
                with bus1_lock:
                    bus1.send(m)
        except Exception:
            pass

        try:
            bus0.shutdown()
            bus1.shutdown()
        except Exception:
            pass


if __name__ == "__main__":
    main()