import time

import can
import threading
from typing import Dict, List, Tuple, Optional
from canbus import RxCache4Ch, BroadcastSenderThread, SingleSenderThread
from usb_can import resolve_canalyst_device_index
from can_msg_builder import *


class Exoskeleton:
    def __init__(self):
        self.bitrate = 1000000
        self.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],
        }
        self.joint_map = {
            "sender-dev0-ch0": ["LJ1", "LJ2", "RJ1", "RJ2"],
            "sender-dev0-ch1": ["LJ5", "LJ3", "RJ5", "RJ3"],
            "sender-dev1-ch0": ["LJ7", "LJ6", "LJ4"],
            "sender-dev1-ch1": ["RJ7", "RJ6", "RJ4"],
        }

        idx0 = resolve_canalyst_device_index("/dev/canalystii_1")
        idx1 = resolve_canalyst_device_index("/dev/canalystii_0")
        self.bus0 = can.Bus(interface="canalystii", channel=[0, 1], bitrate=self.bitrate, device=idx0)
        self.bus1 = can.Bus(interface="canalystii", channel=[0, 1], bitrate=self.bitrate, device=idx1)
        
        # ========= 接收缓存：每个设备一个 Notifier，一个 cache（按 channel 分流） =========
        self.cache0 = RxCache4Ch({0: ["LJ1", "LJ2", "RJ1", "RJ2"], 1: ["LJ5", "LJ3", "RJ5", "RJ3"]})
        self.cache1 = RxCache4Ch({0: ["LJ7", "LJ6", "LJ4"], 1: ["RJ7", "RJ6", "RJ4"]})
        self.notifier0 = can.Notifier(self.bus0, [self.cache0], timeout=0.01)
        self.notifier1 = can.Notifier(self.bus1, [self.cache1], timeout=0.01)

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

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

        sender_dev0_ch0 = BroadcastSenderThread(
            "sender-dev0-ch0", self.bus0, self.bus0_lock, 0, self.cmd_event, self.barrier, self.stop_event)
        sender_dev0_ch1 = SingleSenderThread(
            "sender-dev0-ch1", self.bus0, self.bus0_lock, 1, self.cmd_event, self.barrier, self.stop_event, [1, 3])
        sender_dev1_ch0 = BroadcastSenderThread(
            "sender-dev1-ch0", self.bus1, self.bus1_lock, 0, self.cmd_event, self.barrier, self.stop_event)
        sender_dev1_ch1 = BroadcastSenderThread(
            "sender-dev1-ch1", self.bus1, self.bus1_lock, 1, self.cmd_event, self.barrier, self.stop_event)
        self.senders = [sender_dev0_ch0, sender_dev0_ch1, sender_dev1_ch0, sender_dev1_ch1]
        for sender in self.senders:
            sender.start()
    
    def turn_on(self):
        for sender in self.senders:
            sender.turn_on()
        self.cmd_event.set()
        try:
            self.barrier.wait(timeout=0.05)
        except threading.BrokenBarrierError:
            pass
        self.cmd_event.clear()

    def turn_off(self):
        for sender in self.senders:
            sender.turn_off()
        self.cmd_event.set()
        try:
            self.barrier.wait(timeout=0.05)
        except threading.BrokenBarrierError:
            pass
        self.cmd_event.clear()

    def clear_error(self):
        for sender in self.senders:
            sender.clear_error()
        self.cmd_event.set()
        try:
            self.barrier.wait(timeout=0.05)
        except threading.BrokenBarrierError:
            pass
        self.cmd_event.clear()

    def stop_all(self):
        for sender in self.senders:
            sender.stop()
        self.cmd_event.set()
        try:
            self.barrier.wait(timeout=0.05)
        except threading.BrokenBarrierError:
            pass
        self.cmd_event.clear()

    def set_zero(self):
        for sender in self.senders:
            sender.set_zero()
        self.cmd_event.set()
        try:
            self.barrier.wait(timeout=0.05)
        except threading.BrokenBarrierError:
            pass
        self.cmd_event.clear()

    def set_torq(self, torq: Dict[str, float]):
        torq_cmd = {}
        for sender in self.senders:
            torq_cmd[sender.name] = [torq[x] for x in self.joint_map[sender.name]]
        for sender in self.senders:
            sender.set_torq(torq_cmd[sender.name])
        self.cmd_event.set()
        try:
            self.barrier.wait(timeout=0.05)
        except threading.BrokenBarrierError:
            pass
        self.cmd_event.clear()

    def set_pos(self, pos: Dict[str, float]):
        pos_cmd = {}
        for sender in self.senders:
            pos_cmd[sender.name] = [pos[x] for x in self.joint_map[sender.name]]
        for sender in self.senders:
            sender.set_pos(pos_cmd[sender.name])
        self.cmd_event.set()
        try:
            self.barrier.wait(timeout=0.05)
        except threading.BrokenBarrierError:
            pass
        self.cmd_event.clear()

    def get_states(self):
        c0 = self.cache0.snapshot()
        c1 = self.cache1.snapshot()
        print({"c0": c0, "c1": c1})

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


if __name__ == "__main__":
    exoskeleton = Exoskeleton()
    exoskeleton.turn_on()
    count = 0
    try:
        while True:
            time.sleep(0.02)
            exoskeleton.set_torq({
                "LJ1": 30, "LJ2": 30, "LJ3": 30, "LJ4": 30, "LJ5": 100, "LJ6": 100, "LJ7": 100,
                "RJ1": 30, "RJ2": 30, "RJ3": 30, "RJ4": 30, "RJ5": 100, "RJ6": 100, "RJ7": 100,
            })
            count += 1
            if count % 10 == 0:
                exoskeleton.get_states()
    except KeyboardInterrupt:
        exoskeleton.set_torq({
            "LJ1": 0, "LJ2": 0, "LJ3": 0, "LJ4": 0, "LJ5": 0, "LJ6": 0, "LJ7": 0,
            "RJ1": 0, "RJ2": 0, "RJ3": 0, "RJ4": 0, "RJ5": 0, "RJ6": 0, "RJ7": 0,
        })
        exoskeleton.stop_all()
        exoskeleton.turn_off()
        exoskeleton.stop_event.set()
        exoskeleton.cmd_event.set()
        for th in exoskeleton.senders:
            th.join(timeout=0.5)
        exoskeleton.notifier0.stop()
        exoskeleton.notifier1.stop()
        exoskeleton.bus0.shutdown()
        exoskeleton.bus1.shutdown()