import time
import traceback
from feathersdk.utils.feathertypes import *
from .system import is_can_interface, enable_can_interface, is_can_enabled, get_iface_uid
from .socketcan_tcp import _create_callback, _cansend, _tcpsend_modbus, SocketResult, \
    _clear_all_sockets, _add_endpoint, _remove_endpoint, _polling_step
from feathersdk.utils.common import FeatherThreadLock, exponential_timeout, make_exception_pickleable, currtime, \
    timediff, timestamped_class, dont_allow_on_fork, FeatherThread
import math
from collections import namedtuple
from feathersdk.utils.logger import warning, debug, error, info
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


def _DANGEROUS_set_can_bus_decay_limit(limit: float) -> None:
    """WARNING: Only use this for testing!"""
    global CAN_BUS_DECAY_LIMIT
    CAN_BUS_DECAY_LIMIT = limit


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


@dataclass
@timestamped_class
class SendCanMessage:
    """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
@timestamped_class
class SendTCPMessage:
    """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, *args, **kwargs) -> Self:
        with cls._lock:
            if cls._instance is None:
                cls._instance = super(CommsManager, cls).__new__(cls)
                cls._instance._init(start_polling=True)
        
        if kwargs.get("reset", False):
            cls._instance.reset()

        return cls._instance
    
    def _init(self, start_polling: bool = True) -> None:
        """Initialize the comms manager. Assumes you have the lock already."""
        if hasattr(self, "thread") and start_polling:
            raise ValueError("Cannot re-initialize CommsManager with start_polling=True")
        if hasattr(self, "is_forked") and self.is_forked:
            raise ValueError("Cannot re-initialize CommsManager after fork")

        self.is_forked = False
        self.is_dry = False
        self._callbacks: List[Callable[[SocketResult], None]] = []
        self.endpoints: Dict[str, int] = {}  # {endpoint: iface_uid}
        self._can_loads: Dict[str, LoadAndTime] = {}
        self.__overload_check = True
        self._restarting_thread_id: Optional[int] = None
        self._aborted = False
        self._pending_msgs: Dict[str, Set[MotorUIDType]] = {}
        self._single_msg_checkers: Dict[str, Callable[[SocketResult], MotorUIDType]] = {}
        self._single_msg_mode: bool = False

        # Time to spend polling for messages in the thread function before returning to check for aborting.
        self.polling_func_time_ms: int = 30  

        self._send_callbacks: List[Callable[[SendCallbackParamType], None]] = []
        
        if start_polling:
            self.thread = FeatherThread(target=self._thread_func, name="comms_manager", daemon=True)
            self._polling_callback = _create_callback(self.__main_callback)
            self._thread_error = None
            self._thread_traceback = None
            self._thread_started = False
            self.thread.start()

            try:
                exponential_timeout(0.1, lambda: True if self.is_running() else None)
            except TimeoutError:
                raise CommsManagerNotRunningError()
    
    @dont_allow_on_fork
    def reset(self) -> None:
        """Reset the comms manager to empty initial state."""
        with self._lock:
            _clear_all_sockets()
            self._init(start_polling=False)

    @dont_allow_on_fork
    def _DANGEROUS_disable_overload_check(self) -> None:
        """WARNING: Only use this for testing!"""
        with self._lock:
            self.__overload_check = False
    
    @dont_allow_on_fork
    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

    @dont_allow_on_fork
    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 show info about them
        """
        with self._lock:
            self.is_dry = is_dry
    
    @dont_allow_on_fork
    def add_endpoint(
        self, 
        endpoint: Union[str, List[str]], 
        enable_cans: bool = True, 
        allow_no_enable_can: bool = False, 
        skip_duplicates: bool = False
    ) -> None:
        """Add an endpoint to the comms manager.
        
        Args:
            endpoint: The endpoint to add. Can be a string or a list of strings.
            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.
            skip_duplicates: If True, will not raise an error if the endpoint is already added.
        """
        if isinstance(endpoint, str):
            endpoint = [endpoint]
        
        with self._lock:
            for ep in endpoint:
                if ep in self.endpoints:
                    if skip_duplicates:
                        continue
                    raise ValueError(f"Endpoint {ep} already added")
                
                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`")

                self.endpoints[ep] = get_iface_uid(ep)
                self._can_loads[ep] = LoadAndTime(0.0, 0.0)
                _add_endpoint(ep)

    @dont_allow_on_fork
    def remove_endpoint(self, endpoint: Union[str, List[str]]) -> None:
        """Remove an endpoint from the comms manager.
        
        Args:
            endpoint: The endpoint to remove. Can be a string or a list of strings.
        """
        if isinstance(endpoint, str):
            endpoint = [endpoint]

        with self._lock:
            for ep in endpoint:
                if ep not in self.endpoints:
                    raise ValueError(f"Endpoint {ep} not found")
                del self.endpoints[ep]
                del self._can_loads[ep]
                _remove_endpoint(ep)
    
    @dont_allow_on_fork
    def clear_endpoints(self) -> None:
        """Clear all endpoints from the comms manager."""
        with self._lock:
            self.endpoints.clear()
            self._can_loads.clear()
            self._pending_msgs.clear()
            _clear_all_sockets()
    
    def _thread_func(self) -> None:
        """Thread function for the comms manager."""
        try:
            self._thread_started = True
            if self._aborted:
                error(f"Comms manager thread started but aborted. Aborted: {self._aborted}")
                return
            while not self._aborted:
                _polling_step(self._polling_callback, self.polling_func_time_ms)
            warning(f"Comms manager thread exited. Aborted: {self._aborted}")
        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}")
        finally:
            info(f"Final comms manager thread exiting")
    
    @dont_allow_on_fork
    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)
    
    @dont_allow_on_fork
    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)
    
    @dont_allow_on_fork
    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)
    
    @dont_allow_on_fork
    def register_single_msg_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
    
    @dont_allow_on_fork
    def __main_callback(self, result: SocketResult) -> None:
        """Called by libsocketcantcp library to notify us of a new message.
        
        NOTE: the sockets lock is already acquired by the polling thread, so we don't need to acquire it here.
        """
        try:
            if self._aborted:
                return
            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)
        except Exception as e:
            error(f"Error in comms manager callback: {e}\n{traceback.format_exc()}")
            self._thread_error = e
            self._thread_traceback = traceback.format_exc()
    
    def is_running(self) -> bool:
        """Check if the background polling thread is currently running."""
        return hasattr(self, "thread") and self.thread is not None and self.thread.is_alive() and self._thread_started \
            and not self._aborted

    @dont_allow_on_fork
    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.
        """
        if self._aborted:
            raise RuntimeError("CommsManager is aborted")
        
        with self._lock:
            for callback in self._send_callbacks:
                callback(SendCanMessage(iface=interface, extended=extended, can_id=can_id, data=data, motor_uid=motor_uid))
                
            # 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()) 
        
            if self.is_dry:
                debug("Dry cansend: ", interface, can_id, data)
            else:
                self._run_and_check(_cansend, interface, extended, can_id, data)

    @dont_allow_on_fork
    def tcpsend_modbus(self, ip: str, tid: int, uid: int, fcode: int, reg_addr: int, reg_val: int) -> None:
        """Send a Modbus TCP message."""
        with self._lock:
            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)
            else:
                self._run_and_check(_tcpsend_modbus, ip, tid, uid, fcode, reg_addr, reg_val)
    
    @dont_allow_on_fork
    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, show a warning.
        
        Assumes you have the lock already.
        """
        if not self.is_running():
            raise CommsManagerNotRunningError()
        func(*args)

    @dont_allow_on_fork
    def on_abort(self) -> None:
        """Called when the comms manager is aborted."""
        if self.is_forked or self._aborted:
            return
        
        with self._lock:
            self._aborted = True

            if hasattr(self, "thread") and self.thread is not None:
                self.thread.join(timeout=0.5)
                if self.thread.is_alive():
                    error(f"Comms manager thread did not exit in time")
                self.thread = None
    
    def on_fork(self) -> None:
        """Called when the process forks."""
        self.is_forked = True
        self.is_dry = False
        self._callbacks: List[Callable[[SocketResult], None]] = []
        self.endpoints: Dict[str, int] = {}  # {endpoint: iface_uid}
        self._can_loads: Dict[str, LoadAndTime] = {}
        self.__overload_check = True
        self._restarting_thread_id: Optional[int] = None
        self._aborted = True
        self._pending_msgs: Dict[str, Set[MotorUIDType]] = {}
        self._single_msg_checkers: Dict[str, Callable[[SocketResult], MotorUIDType]] = {}
        self._single_msg_mode: bool = False
        self._send_callbacks: List[Callable[[SendCallbackParamType], None]] = []
        self.thread = None
        self._thread_error = None
        self._thread_traceback = None
        self._thread_started = False
