import queue

import can
import usb
import time
import struct
import threading
from dataclasses import dataclass
from typing import Dict, List, Tuple, Optional
from queue import Queue

from can_msg_builder import *
from can_msg_parser import *


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,
        idle_hz: float = 500.0,
        idle_timeout_s: float = 0.01,  # 多久没控制算“空闲”
    ):
        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.queue = Queue()
        self.idle_period = 1.0 / float(idle_hz)
        self.idle_timeout_s = float(idle_timeout_s)
        self._last_cmd_ts = time.perf_counter()
        self._next_idle_ts = time.perf_counter() + self.idle_period

    def turn_on(self):
        pass

    def turn_off(self):
        pass

    def stop(self):
        pass

    def clear_error(self):
        pass

    def set_zero(self):
        pass

    def set_torq(self, torque: List[float]):
        pass

    def set_vel(self, vel: List[float]):
        pass

    def set_pos(self, pos: List[float]):
        pass

    def _mark_cmd(self):
        self._last_cmd_ts = time.perf_counter()

    def _drain_msgs(self):
        msgs = []
        while True:
            try:
                msgs.append(self.queue.get_nowait())
            except queue.Empty:
                break
        return msgs

    def _send_msgs(self, msgs):
        for msg in msgs:
            msg.channel = self.channel_index
            with self.bus_lock:
                self.bus.send(msg)

    def _thread_idle(self):
        pass

    def run(self):
        while not self.stop_event.is_set():
            now = time.perf_counter()

            # 1) 优先处理控制 tick（有控制指令）
            if self.tick_event.is_set():
                self._mark_cmd()
                try:
                    self.barrier.wait(timeout=0.2)  # 0.05 太激进，建议放宽
                except threading.BrokenBarrierError:
                    # barrier broken 时也别死循环，稍退避
                    time.sleep(0.001)
                    continue

                try:
                    msgs = self._drain_msgs()
                    if msgs:
                        with self.bus_lock:
                            for msg in msgs:
                                msg.channel = self.channel_index
                                self.bus.send(msg)
                except usb.core.USBTimeoutError:
                    # USB 写超时：退避，避免连环打爆
                    time.sleep(0.005)
                except can.CanError:
                    pass

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

                # 控制刚发完，顺手把 idle 定时器往后推，避免“控制后马上 query”
                self._next_idle_ts = time.perf_counter() + self.idle_period
                continue

            # 2) 没控制 tick：进入空闲轮询（满足空闲时间阈值）
            if (now - self._last_cmd_ts) >= self.idle_timeout_s:
                if now >= self._next_idle_ts:
                    try:
                        self._thread_idle()
                    except usb.core.USBTimeoutError:
                        time.sleep(0.005)
                    except can.CanError:
                        pass
                    # 维持均匀节拍
                    self._next_idle_ts += self.idle_period
                else:
                    # 睡到下一次 idle
                    time.sleep(min(0.001, self._next_idle_ts - now))
            else:
                # 刚停止控制不久：短睡，等进入 idle
                time.sleep(0.001)



class BroadcastSenderThread(SenderThread):
    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, bus, bus_lock, channel_index, tick_event, barrier, stop_event)

    def turn_on(self):
        self.queue.put(build_turn_on_broadcast())

    def turn_off(self):
        self.queue.put(build_turn_off_broadcast())

    def stop(self):
        self.queue.put(build_stop_broadcast())

    def clear_error(self):
        self.queue.put(build_clear_error_broadcast())

    def set_zero(self):
        pass

    def set_torq(self, torques):
        assert len(torques) <= 4
        self.queue.put(build_torque_broadcast(torques))

    def set_vel(self, vel: List[float]):
        assert len(vel) <= 4
        self.queue.put(build_velocity_broadcast(vel))

    def set_pos(self, pos):
        assert len(pos) <= 4
        self.queue.put(build_position_broadcast(pos))

    def _thread_idle(self):
        msg = build_get_state_broadcast()
        msg.channel = self.channel_index
        try:
            with self.bus_lock:
                self.bus.send(msg)
        except can.CanError:
            pass

class SingleSenderThread(SenderThread):
    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,
            ms_index: List[int]
    ):
        super().__init__(name, bus, bus_lock, channel_index, tick_event, barrier, stop_event)
        self.ms_index = ms_index

    def turn_on(self):
        for i in range(1, 5):
            self.queue.put(build_turn_on_single(i))

    def turn_off(self):
        for i in range(1, 5):
            self.queue.put(build_turn_off_single(i))

    def stop(self):
        for i in range(1, 5):
            self.queue.put(build_stop_single(i))

    def clear_error(self):
        for i in range(1, 5):
            self.queue.put(build_clear_error_single(i))

    def set_zero(self):
        pass

    def set_torq(self, torques):
        for i in range(1, 5):
            if i in self.ms_index:
                self.queue.put(build_torque_single(i, torques[i-1], True))
            else:
                self.queue.put(build_torque_single(i, torques[i-1], False))

    def set_vel(self, vel: List[float]):
        for i in range(1, 5):
            self.queue.put(build_velocity_single(i, vel[i - 1]))

    def set_pos(self, pos):
        for i in range(1, 5):
            self.queue.put(build_set_multiturn_position_single(i, pos[i - 1]))

    def _thread_idle(self):
        for i in range(1, 5):
            msg = build_get_multiturn_position_single(i)
            msg.channel = self.channel_index
            try:
                with self.bus_lock:
                    self.bus.send(msg)
            except can.CanError:
                pass
            except usb.core.USBTimeoutError:
                time.sleep(0.001)



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


@dataclass
class MotorState:
    stamp: float = 0.0
    temp: float = 0.0
    single_turn_rad: float = 0.0
    multi_turn_rad: float = 0.0
    vel: float = 0.0
    power_raw: float = 0.0


class RxCache4Ch(can.Listener):
    """
    一个 Notifier 监听一个 “设备Bus(包含两个channel)”，按 msg.channel 分流到 ch0/ch1 的缓存。
    """
    def __init__(self, joint_map: Dict[int, List[str]]) -> None:
        super().__init__()
        self._lock = threading.Lock()
        self.joint_state = {}
        self.ch_mid_idx = {}
        for ch, joint_names in joint_map.items():
            for i, joint_name in enumerate(joint_names):
                self.ch_mid_idx[(ch, i+1)] = joint_name
                self.joint_state[joint_name] = MotorState()

    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

        arb = msg.arbitration_id
        if arb < 0x141 or arb > 0x160:
            return
        mid = arb - 0x140

        joint_name = self.ch_mid_idx[(ch, mid)]
        if joint_name is None:
            return

        if msg.data[0] in [0xA0, 0xA1, 0x9C]:
            res = parse_single_turn_state(msg.data)
            # print(res)
            with self._lock:
                self.joint_state[joint_name].stamp = time.time()
                self.joint_state[joint_name].temp = res["temperature_C"]
                self.joint_state[joint_name].power_raw = res["iq_or_power_raw"]
                self.joint_state[joint_name].vel = res["speed_dps"]
                self.joint_state[joint_name].single_turn_rad = res["rad"]
        elif msg.data[0] in [0x92]:
            res = parse_multi_turn_state(msg.data)
            with self._lock:
                self.joint_state[joint_name].multi_turn_rad =res["rad"]
        else:
            print(msg.data[0])

    def snapshot(self):
        with self._lock:
            return self.joint_state