from feathersdk.comms.socketcan_tcp import SocketResult
from feathersdk.comms.system import S_uint16le, S_uint16be, S_uint64be, S_uint32le, S_uint32be
from feathersdk.robot.motors.robstride import RobstrideMotorMode, RobstrideMotorError, RobstrideMotorMsg, \
    get_param_info, RobstrideFeedbackMessageBounds
from feathersdk.utils.feathertypes import TYPE_CHECKING, NamedTuple
from feathersdk.utils.common import try_log_exception_peacefully

if TYPE_CHECKING:
    from feathersdk.robot.motors.motors_manager import MotorsManager


_SKIPPED_MESSAGES = [RobstrideMotorMsg.BeginFlash, RobstrideMotorMsg.FlashData, RobstrideMotorMsg.EndFlash,
                     RobstrideMotorMsg.FlashFileInfo]

def _handle_feedback_message(manager: 'MotorsManager', m_id: int, data: bytes, e_byte: int, _msg_op: RobstrideMotorMsg):
    if m_id not in manager.motors_by_id:
        return
    
    motor = manager.motors_by_id[m_id]
    old_mode = motor.mode[0]

    # Update the motor's feedback values
    motor.update_feedback(
        angle=S_uint16be.unpack_and_map(data[0:2], *RobstrideFeedbackMessageBounds.ANGLE_MAP[motor.version]), 
        velocity=S_uint16be.unpack_and_map(data[2:4], *RobstrideFeedbackMessageBounds.VELOCITY_MAP[motor.version]), 
        torque=S_uint16be.unpack_and_map(data[4:6], *RobstrideFeedbackMessageBounds.TORQUE_MAP[motor.version]), 
        temp=S_uint16be.unpack(data[6:8]) / RobstrideFeedbackMessageBounds.TEMP_DIVISOR, 
        errors=RobstrideMotorError.parse_motor_feedback_errors(e_byte & 0x3F),  # Remove motor mode bits
        mode=RobstrideMotorMode.parse_mode(e_byte)
    )

    # Trigger compliance mode if necessary
    if motor.should_trigger_compliance_mode() and old_mode != RobstrideMotorMode.Reset:
        manager.trigger_compliance_mode(motor.family_name, motor.motor_name)


def _handle_read_param_message(manager: 'MotorsManager', motor_id: int, data: bytes):
    """Parse the value and update the motor's property."""
    p_info = get_param_info(S_uint16le.unpack(data[:2]))
    ver = manager.motors_by_id[motor_id].version
    val = p_info.dtype.unpack(data[4:])
    if not p_info.is_within_range(val, ver):
        raise ValueError(f"Read param value {val} out of range for {p_info.name}, version {ver}: {p_info.range[ver]}")
    manager.motors_by_id[motor_id].update_property(p_info.name, val)


def _handle_get_device_id_message(manager: 'MotorsManager', motor_id: int, data: bytes):
    """Parse the device id and update the motor's property."""
    manager.motors_by_id[motor_id].update_property('device_id', S_uint64be.unpack(data))


def _handle_fault_feedback_message(manager: 'MotorsManager', motor_id: int, data: bytes):
    """Parse the fault feedback and update the motor's property."""
    fault_bytes = S_uint32le.unpack(data) & 0x00FF_FFFF  # Only 3 bytes are used for fault bits
    parsed = RobstrideMotorError.parse_fault_feedback_errors(fault_bytes)
    motor = manager.motors_by_id[motor_id]
    motor._log_robstride_error_changes(parsed, "fault_feedback")
    motor.errors.update(parsed)


@try_log_exception_peacefully
def handle_robstride_motor_message(manager: 'MotorsManager', result: SocketResult) -> None:
    """Handle an incoming Robstride motor message."""
    msg_op, motor_id, host_id, e_byte = parse_robstride_can_id(result.can_id, motor_to_host=True)
    
    # Skip if motor is not in our map (might be from an unknown motor)
    if motor_id not in manager.motors_by_id:
        return

    if msg_op == RobstrideMotorMsg.GetDeviceId:
        _handle_get_device_id_message(manager, motor_id, result.data)
    elif msg_op in [RobstrideMotorMsg.MotorFeedback, RobstrideMotorMsg.SetMotorActiveReporting]:
        _handle_feedback_message(manager, motor_id, result.data, e_byte, msg_op)
    elif msg_op == RobstrideMotorMsg.ReadParam:
        _handle_read_param_message(manager, motor_id, result.data)
    elif msg_op == RobstrideMotorMsg.FaultFeedback:
        _handle_fault_feedback_message(manager, motor_id, result.data)
    elif msg_op in _SKIPPED_MESSAGES:
        pass
    else:
        raise NotImplementedError(f"Unhandled Robstride motor message type: {msg_op}")


class RobstrideParsedCanId(NamedTuple):
    msg_op: RobstrideMotorMsg
    motor_id: int
    host_id: int
    extra_byte: int


def parse_robstride_can_id(can_id: int, motor_to_host: bool = True) -> RobstrideParsedCanId:
    """Parses the CAN ID into named tuple of (msg_op, motor_id, host_id, extra_byte)
    
    CAN_ID format: [CAN_FLAGS (3 bits)] [MSG_OP (4 bits)] [EXTRA_BYTE (8 bits)] [ID_1 (8 bits)] [ID_2 (8 bits)]
                       31-29                28-24             23-16                 15-8            7-0

    The extra_byte value is only used in some robstride messages, such as OperationControl or MotorFeedback frames.

    Args:
        can_id: The CAN ID to parse.
        motor_to_host: If True, then it is assumed the message is sent from the motor to the host, in which case ID_1 
            is the motor ID and ID_2 is the host ID. If False, then the reverse is true and ID_1 is the host ID
            and ID_2 is the motor ID.
    
    Returns:
        RobstrideParsedCanId: (msg_op, motor_id, host_id, extra_byte)
    """
    msg_op = RobstrideMotorMsg((can_id & 0x1F00_0000) >> 24)
    extra_byte = (can_id & 0x00FF_0000) >> 16
    motor_id = (can_id & 0x0000_FF00) >> 8
    host_id = can_id & 0x0000_00FF

    if not motor_to_host:
        motor_id, host_id = host_id, motor_id
    
    return RobstrideParsedCanId(msg_op=msg_op, motor_id=motor_id, host_id=host_id, extra_byte=extra_byte)
