import threading
import traceback
from feathersdk.utils.feathertypes import Optional, Callable, Any, Union, Self
from .system import is_can_interface, enable_can_interface, is_can_enabled
from .socketcan_tcp import _subscribe_multi, _cansend, _tcpsend_modbus, SocketResult, _stop_polling, \
    _clear_all_sockets, _get_err_str
from feathersdk.utils.common import make_exception_pickleable, FeatherThreadLock, exponential_timeout, currtime, \
    timediff, TimestampedMixin
import math
from collections import namedtuple
from feathersdk.utils.logger import warning, debug
from dataclasses import dataclass


MotorUIDType = str


class SocketCANLibError(Exception):
    """Error raised when the socketcan library returns a non-zero error code."""
    pass

class CanOverloadError(Exception):
    """Error raised when the CAN bus is overloaded."""
    pass

@make_exception_pickleable
class UnknownInterfaceError(Exception):
    """Error raised when the interface is not an interface, or is not being tracked by CommsManager."""
    def __init__(self, *, interface: str):
        super().__init__(f"Interface \"{interface}\" is not an interface, or is not being tracked by CommsManager")

@make_exception_pickleable
class CommsManagerNotRunningError(Exception):
    """Error raised when the comms manager is not running."""
    def __init__(self):
        super().__init__("CommsManager is not running")

@make_exception_pickleable
class CanNotEnabledError(Exception):
    """Error raised when a CAN interface is not enabled."""
    def __init__(self, *, interface: str):
        super().__init__(f"CAN interface \"{interface}\" is not enabled")

@make_exception_pickleable
class MissingMotorUIDError(Exception):
    """Error raised when a motor UID is required but not provided."""
    def __init__(self):
        super().__init__(f"`motor_uid` is required in single message mode")

@make_exception_pickleable
class PendingReceiveMessageError(Exception):
    """Error raised when a message is to be sent to a motor when we are awaiting a response in single message mode."""
    def __init__(self, *, motor_uid: MotorUIDType):
        super().__init__(f"Message to motor \"{motor_uid}\" is pending, cannot send more until response is received")


# We calculate the maximum number of messages that can be sent over a CAN bus per second using math similar to:
# https://electronics.stackexchange.com/questions/121329/whats-the-maximum-can-bus-frame-message-rate-at-125-kbit-s
# 
# Including all the extra overhead, CRC, bit stuffing, etc, we get a theoretical 144 bits per frame. At 1Mbs, this
# would be a theoretical max of 6944 frames per second. Since most all messages we send will have 1 reply from the
# device, we divide by 2 to get 3472 messages-sent per second. To give us ~10% leeway, we limit to 3000 messages-sent
# per second on average.
# 
# Using our exponential decay formula of N_{x+1} = N_x * e^(-t * decay_rate) + decay_scale, and a decay rate of 0.01,
# experiments (see /playground/justin_exp_decay_math.py) show that at a rate of 3000 messages-sent per second, the
# decayed-count maxes out at ~300. So, we use this limit here.
CAN_BUS_DECAY_LIMIT: float = 300.0;
CAN_BUS_DECAY_RATE: float = 0.01;
CAN_BUS_DECAY_SCALE: float = 10.0;  # Controls how quickly we reach our limit. Higher=more quicker.
                                    # Value of 10.0 = ~30 messages before limit at full speed, ~50 at 100us/message


LoadAndTime = namedtuple("LoadAndTime", ["load", "last_message_time"])


@dataclass
class SendCanMessage(TimestampedMixin):
    """Data class representing a CAN message that is about to be sent. Only used for WireHawk"""
    iface: str
    extended: bool
    can_id: int
    data: bytes
    motor_uid: Optional[MotorUIDType] = None

    def __hash__(self) -> int:
        return hash((self.iface, self.extended, self.can_id, self.data, self.motor_uid))

@dataclass
class SendTCPMessage(TimestampedMixin):
    """Data class representing a TCP message that is about to be sent. Only used for WireHawk"""
    ip: str
    tid: int
    uid: int
    fcode: int
    reg_addr: int
    reg_val: int

    def __hash__(self) -> int:
        return hash((self.ip, self.tid, self.uid, self.fcode, self.reg_addr, self.reg_val))


SendCallbackParamType = Union[SendCanMessage, SendTCPMessage]


class CommsManager:
    """Handles sending and receiving messages over CAN and TCP.

    Provides a single point of entry for sending and receiving messages over CAN and TCP. Any new instances of this
    class will point to the same instance.

    The manager also checks for message overload on the CAN bus and will raise an error if the message rate is too high.

    You can 
    """

    _instance: Optional[Self] = None
    """The singleton instance of the CommsManager. Do not modify at runtime!!!"""

    _lock: FeatherThreadLock = FeatherThreadLock("CommsManager")
    
    def __new__(cls) -> Self:
        with cls._lock:
            if cls._instance is None:
                cls._instance = super(CommsManager, cls).__new__(cls)
                cls._instance.__reset()
            return cls._instance
    
    def __reset(self) -> None:
        """Reset the comms manager to a clean state. Assumes you have the lock already."""
        if hasattr(self, "thread") and self.is_running():
            raise ValueError("Cannot reset CommsManager while it is running")
        
        self.thread = None
        self._thread_error: Optional[Exception] = None
        self._thread_traceback: Optional[str] = None

        self.is_dry = False
        self._callbacks: list[Callable[[SocketResult], None]] = []
        self.endpoints: list[str] = []
        self._can_loads: dict[str, LoadAndTime] = {}
        self.__overload_check = True

        self._pending_msgs: dict[str, set[MotorUIDType]] = {}
        self._single_msg_checkers: dict[str, Callable[[SocketResult], MotorUIDType]] = {}
        self._single_msg_mode: bool = False

        # Do not override any existing send callbacks since they're used by WireHawk
        send_callbacks = [] if not hasattr(self, "_send_callbacks") else self._send_callbacks
        self._send_callbacks: list[Callable[[SendCallbackParamType], None]] = send_callbacks
        
    def _DANGEROUS_disable_overload_check(self) -> None:
        """WARNING: Only use this for testing!"""
        with self._lock:
            self.__overload_check = False
    
    def set_single_message_mode(self, single_message_mode: bool) -> None:
        """Set the comms manager to single message mode.
        
        In single message mode, the comms manager will only send one message at a time to any one motor.
        """
        with self._lock:
            self._single_msg_mode = single_message_mode

    def set_is_dry(self, is_dry: bool) -> None:
        """Enable/disable 'dry run' mode.

        In dry run mode, the comms manager will not send any messages, and instead just print info about them
        """
        with self._lock:
            self.is_dry = is_dry

    def start(self, endpoints: list[str], enable_cans: bool = True, allow_no_enable_can: bool = False) -> None:
        """Start the comms manager
        
        This will start a background thread that will poll all the endpoints and call the callback for each message.

        Will call `enable_can_interface` for each endpoint that is a CAN interface with default bitrate of 1Mbs. Will
        create vCAN interfaces for each endpoint that is a vCAN interface.

        Args:
            endpoints: list of endpoints to subscribe to, eg: ["can0", "vcan1", "192.168.11.210"], etc.
            enable_cans: If True, will attempt to enable any CAN interface that is not already enabled.
            allow_no_enable_can: If True, will not raise an error if any of the CAN interfaces are not enabled.
        """
        with self._lock:
            if self.is_running():
                raise ValueError("CommsManager is already running")
            
            for ep in endpoints:
                if is_can_interface(ep):
                    try:
                        if not is_can_enabled(ep, bitrate=1_000_000):
                            if not enable_cans:
                                raise CanNotEnabledError(interface=ep)
                            enable_can_interface(ep, bitrate=1_000_000)
                    except Exception as e:
                        if not allow_no_enable_can:
                            raise e
                        else:
                            warning(f"Could not enable CAN interface {ep}: {e}, for the robot to work properly, "
                                    f"run `sudo ip link set {ep} up`")
            
            _get_err_str()  # Just to make sure library is compiled and ready
            
            self.endpoints = endpoints
            self._can_loads = {ep: LoadAndTime(0.0, 0.0) for ep in endpoints if is_can_interface(ep)}

            self.thread = threading.Thread(target=self._thread_func, daemon=True)
            self._thread_error = None
            self._thread_traceback = None
            self.thread.start()
    
    def _thread_func(self) -> None:
        """Thread function for the comms manager."""
        try:
            _subscribe_multi(self.endpoints, self.__main_callback)
        except Exception as e:
            self._thread_error = e
            self._thread_traceback = traceback.format_exc()
            warning(f"Error in comms manager thread:\n{self._thread_traceback}")
    
    def restart_with_added_endpoints(self, endpoints: Union[str, list[str]], start_ok: bool = False) -> None:
        """Add endpoints to the (running) comms manager and restart it
        
        Note: if the endpoint is already in the comms manager, this will do nothing.
        
        Args:
            endpoints: The endpoint(s) to add. Can be string or multiple strings.
            start_ok: By default, an error will be raised if the comms manager is not running. If True, will start the 
                comms manager if it is not running.
        """
        if isinstance(endpoints, str):
            endpoints = [endpoints]
        
        if not self.is_running():
            if start_ok:
                return self.start(endpoints)
            raise CommsManagerNotRunningError()
        
        with self._lock:
            new_endpoints = list(set(endpoints).union(set(self.endpoints)))
        
        if len(new_endpoints) > 0:
            self.close()
            self.start(new_endpoints)
    
    def add_callback(self, callback: Callable[[SocketResult], None]) -> None:
        """Add a callback to be called for each message."""
        with self._lock:
            self._callbacks.append(callback)
    
    def _add_send_callback(self, callback: Callable[[SendCallbackParamType], None]) -> None:
        """Add a callback to be called for each message that is about to be sent."""
        with self._lock:
            self._send_callbacks.append(callback)
    
    def remove_callback(self, callback: Callable[[SocketResult], None]) -> None:
        """Remove a callback from being called for each message."""
        with self._lock:
            self._callbacks.remove(callback)
    
    def close(self, timeout: float = 0.1) -> None:
        """Close the comms manager.
        
        This will stop the background thread and close all the endpoints.

        Args:
            timeout: The timeout in seconds to wait for the comms manager to ensure closed.
        """
        with self._lock:
            if not self.is_running():
                raise CommsManagerNotRunningError()
        
            _stop_polling()
        
        def _clear_all_sockets_and_reset():
            """Clear all sockets and reset the comms manager. 
            
            Used within exponential_timeout. Checks to make sure the thread is actually stopped, then clears and 
            resets the comms manager. We need to do it this way because we can get into deadlocks if we keep the lock
            the entire time and a message comes in to the handler.
            """
            with self._lock:
                if self.is_running():
                    return None
                
                _clear_all_sockets()
                self.__reset()

                return True
            
        # Make sure the thread is actually stopped
        try:
            exponential_timeout(timeout, _clear_all_sockets_and_reset)
        except TimeoutError:
            raise SystemError("CommsManager is still running after closing")
    
    def register_single_message_checker(self, iface: str, func: Callable[[SocketResult], Optional[MotorUIDType]]) -> None:
        """Register a function to check for expected responses in single message mode.
        
        Args:
            func: A function that takes a SocketResult and returns a string unique to the motor that sent the message,
                or None if the message is not expected and should be ignored.
        """
        with self._lock:
            self._single_msg_checkers[iface] = func
    
    def __main_callback(self, result: SocketResult) -> None:
        """Called by libsocketcantcp library to notify us of a new message."""
        if result.is_error():
            warning(f"Error in comms manager callback: {result.err_type} - {result.err_msg}")
            return
        
        # Check for expected responses in single message mode
        with self._lock:
            if self._single_msg_mode and result.socket_name in self._pending_msgs:
                if (res := self._single_msg_checkers[result.socket_name](result)) is not None:
                    if not isinstance(res, MotorUIDType):
                        raise TypeError(f"Expected MotorUIDType, got {type(res)}")
                    self._pending_msgs[result.socket_name].discard(res)

        for callback in self._callbacks:
            callback(result)
    
    def is_running(self) -> bool:
        """Check if the background polling thread is currently running."""
        return self.thread is not None and self.thread.is_alive()

    def cansend(
        self, 
        interface: str, 
        extended: bool, 
        can_id: int, 
        data: bytes, 
        motor_uid: Optional[MotorUIDType] = None
    ) -> None:
        """Send a CAN message.
        
        Args:
            interface: The interface to send the message on.
            extended: Whether the message is an extended (29-bit) CAN message, or a standard (11-bit) CAN message.
            can_id: The CAN ID of the message. Must be in the range [0-0x1FFFFFFF] for extended messages, or 
                [0-0x7FF] for standard messages.
            data: The data to send in the message. Must be 8 bytes long.
            motor_uid: The unique identifier of the motor that sent the message. Only needed in single message mode.
        """
        for callback in self._send_callbacks:
            callback(SendCanMessage(iface=interface, extended=extended, can_id=can_id, data=data, motor_uid=motor_uid))

        if self.is_dry:
            debug("Dry cansend: ", interface, can_id, data)
            return
        
        with self._lock:
            # Make sure we are not in single message mode and have a pending message for this interface
            if self._single_msg_mode:
                if interface not in self._pending_msgs:
                    self._pending_msgs[interface] = set()
                if motor_uid is None:
                    raise MissingMotorUIDError()
                if not isinstance(motor_uid, MotorUIDType):
                    raise TypeError(f"Expected MotorUIDType, got {type(motor_uid)}")
                if motor_uid in self._pending_msgs[interface]:
                    raise PendingReceiveMessageError(motor_uid=motor_uid)
                
                # We are good to send this message, mark it as pending
                self._pending_msgs[interface].add(motor_uid)
            
            # Make sure we are not overloading the CAN bus
            if interface not in self._can_loads:
                raise UnknownInterfaceError(interface=interface)
            
            if self.__overload_check:
                dt_ms = timediff(self._can_loads[interface].last_message_time) * 1000.0 * CAN_BUS_DECAY_SCALE
                new_load = self._can_loads[interface].load * math.exp(-dt_ms * CAN_BUS_DECAY_RATE) + CAN_BUS_DECAY_SCALE
                if new_load > CAN_BUS_DECAY_LIMIT:
                    raise CanOverloadError(f"CAN bus {interface} is overloaded")
            
                self._can_loads[interface] = LoadAndTime(load=new_load, last_message_time=currtime()) 
        
        self._run_and_check(_cansend, interface, extended, can_id, data)

    def tcpsend_modbus(self, ip: str, tid: int, uid: int, fcode: int, reg_addr: int, reg_val: int) -> None:
        """Send a Modbus TCP message."""
        for callback in self._send_callbacks:
            callback(SendTCPMessage(ip=ip, tid=tid, uid=uid, fcode=fcode, reg_addr=reg_addr, reg_val=reg_val))

        if self.is_dry:
            debug("dry tcpsend", ip, tid, uid, fcode, reg_addr, reg_val)
            return
        
        self._run_and_check(_tcpsend_modbus, ip, tid, uid, fcode, reg_addr, reg_val)
    
    def _run_and_check(self, func: Callable[[Any], None], *args: Any) -> None:
        """Run a library function and check the result. If the result is non-zero, print a warning."""
        if not self.is_running():
            raise CommsManagerNotRunningError()
        func(*args)
