from feathersdk.utils.feathertypes import Optional, Callable, Any, TYPE_CHECKING
from feathersdk.comms import SocketResult, MotorUIDType
from .params import RobstrideMotorMsg as MotorMsg, RobstrideRunMode, _ROBSTRIDE_PARAMS
from .message_handler import parse_robstride_can_id
from feathersdk.comms.system import maybe_enable_all_can_interfaces, S_uint16be, S_uint32le, S_uint64be
from feathersdk.utils import progressbar
from feathersdk.utils.common import timediff, currtime
import threading
import time
from feathersdk.utils.feathertypes import *

if TYPE_CHECKING:
    from .robstride_motor import RobstrideMotor


# The location in firmware file where the version number is stored.
BIN_VERSION_LOC = 2062


def prepare_robstride_admin_command(func: Callable) -> Callable:
    """Decorator to make robstride admin function get their self lock before running, and init/clean comms"""
    global _WITHIN_PREPARE

    def wrapper(self: 'RobstrideAdminControl', *args: Any, **kwargs: Any) -> Any:
        with self.lock:
            self.comms.add_endpoint(maybe_enable_all_can_interfaces(), skip_duplicates=True)

            if self.motor.manager is None:  # Create a motor manager if not present
                from feathersdk.robot.motors.motors_manager import MotorsManager
                mm = MotorsManager()
                mm.add_motors({self.motor.motor_name: self.motor}, "motors", {})
                self.motor.manager = mm
            
            if self.motor.iface is None:  # Find the motor's interface if not present
                self.motor.manager.find_motors([self.motor.motor_name])
                self.iface = self.motor.iface
            
            self.motor.default_host_id = self.main_host_id

            self.comms.register_single_msg_checker(self.iface, self._msg_checker)
            self.comms.set_single_message_mode(True)
            result = func(self, *args, **kwargs)
            self.comms.set_single_message_mode(False)

            return result
    return wrapper


class RobstrideAdminControl:
    """Used to perform administrative tasks on a Robstride motor.

    For example: setting motor's ID, or flashing the motor's firmware.
    """

    def __init__(self, motor: 'RobstrideMotor'):
        self.motor = motor
        self.iface = motor.iface
        self.comms = motor.comms

        self.main_host_id = 0xFD
        self.last_response = None
        self.lock = threading.Lock()

        self.checking_op = MotorMsg.ReadParam  # Empty default
        self.checking_motor_id = None

    def _msg_checker(self, result: SocketResult) -> Optional[MotorUIDType]:
        """Check the incomming message to see if it is a response to one of our commands"""
        if not result.is_can() or not result.is_extended_can():
            return None
        
        msg_op, motor_id, host_id, extra_byte = parse_robstride_can_id(result.can_id)

        if msg_op == self.checking_op and motor_id == self.checking_motor_id and host_id == self.exp_host:
            self.last_response = result
            return MotorUIDType(self.checking_motor_id)
        
        return None

    def _send_await(self, exp_msg: MotorMsg, func: Callable, *args: Any, exp_id: Optional[int] = None, 
                    timeout: float = 1.0, exp_byte: int = 0, **kwargs: Any) -> SocketResult:
        """Send a command (by calling func(*args, **kwargs)) and await a response."""
        if self.last_response is not None:
            raise RuntimeError(f"Called _send_await() while already processing a response: {self.last_response}")
        
        self.exp_host = self.main_host_id if exp_msg not in [MotorMsg.SetID, MotorMsg.GetDeviceId] else 0xFE
        self.checking_motor_id = exp_id if exp_id is not None else self.motor.motor_id
        self.checking_op = exp_msg

        func(*args, **kwargs)

        start_time = currtime()
        while self.last_response is None:
            time.sleep(0.001)
            if timediff(start_time) > timeout:
                raise TimeoutError(f"Waiting for {exp_msg} response after {timediff(start_time)} seconds")

        response, self.last_response = self.last_response, None

        assert len(self.comms._pending_msgs[self.iface]) == 0, \
            f"Expected no pending messages on interface {self.iface}, got: {self.comms._pending_msgs[self.iface]}"

        # Check we got a good response from the motor.
        msg_op, motor_id, host_id, e_byte = parse_robstride_can_id(response.can_id)
        assert msg_op == exp_msg and motor_id == self.checking_motor_id and host_id == self.exp_host, \
            f"Got bad {exp_msg} response: msg_op={msg_op} motor_id={motor_id} host_id={host_id}," \
            f" expected: {exp_msg} {self.checking_motor_id} {self.exp_host}"
        assert e_byte == exp_byte, f"Expected {exp_msg} extra byte {exp_byte}, got: {e_byte}. Data: {response.data}"

        return response
    
    @prepare_robstride_admin_command
    def set_id(self, new_id: int, timeout: float = 1.0) -> None:
        """Set a new ID on the motor. Make sure it doesn't clash with any other motor's ID!"""
        if new_id < 0 or new_id >= 0xFE:
            raise ValueError("Motor ID must be in range [0, 0xFD]")

        m_id = self.motor.motor_id
        can_id = self.motor._msg_can_id(MotorMsg.SetID, m_id, host_id=self.main_host_id, e_byte=new_id)
        self._send_await(MotorMsg.GetDeviceId, self.motor.comms.cansend, self.motor.iface, True, can_id,
                            bytes([0] * 8), MotorUIDType(new_id), exp_id=new_id, timeout=timeout)
        self.motor.motor_id = new_id
    
    @prepare_robstride_admin_command
    def flash_firmware(self, firmware_path: str, progress: bool = False) -> None:
        """Flash firmware from a .bin file to the motor."""
        m_uid = MotorUIDType(self.motor.uid)

        # Check the motor's current firmware version
        version_vals = tuple(map(int, self._read_firmware_version().split(".")))
        print(f"Current firmware version: {'.'.join(map(str, version_vals))}")

        # Read in the firmware file and make sure it matches the motor version
        with open(firmware_path, 'rb') as f:
            firmware_data = f.read()
        bin_version = tuple(firmware_data[BIN_VERSION_LOC:BIN_VERSION_LOC + 8])

        assert bin_version[:3] == (0, 0, ord('V')) and bin_version[-1] == 7, \
            f"Invalid firmware file version format in location {(BIN_VERSION_LOC, BIN_VERSION_LOC + 8)}: {bin_version}"
        assert version_vals[0:2] == bin_version[3:5], \
            f"Firmware file version \"{'.'.join(map(str, bin_version[3:5]))}\"" \
            f" does not match motor version \"{'.'.join(map(str, version_vals[0:2]))}\""
        
        print(f"Flashing new firmware version: {'.'.join(map(str, bin_version[3:-1]))}")

        # Check the file size and calculate the number of packets we will send
        file_size = len(firmware_data)
        
        if file_size == 0:
            raise ValueError("Firmware file is empty")
        if file_size >= 0x80001:  # 512KB + 1
            raise ValueError(f"Firmware file too large: {file_size} bytes (max 512KB)")
        
        num_packets = (file_size + 7) // 8  # 8 bytes per packet
        print(f"Read firmware file: {firmware_path} - size: {file_size} bytes, will send {num_packets} CAN packets")

        # Get the device ID of the motor
        self._send_await(MotorMsg.GetDeviceId, self.motor.get_device_id)
        device_id = self.motor.get_property_value('device_id')
        assert device_id is not None, "Failed to get device ID"
        print(f"Got motor Device ID: {device_id}")

        # Send the BeginFlash command.
        print(f"Sending BeginFlash command...")
        can_id = self.motor._msg_can_id(MotorMsg.BeginFlash, self.motor.motor_id, self.main_host_id)
        print(f"BeginFlash command CAN ID: {hex(can_id)}")
        data = S_uint64be.pack(device_id)
        self._send_await(MotorMsg.BeginFlash, self.comms.cansend, self.iface, True, can_id, data, m_uid, timeout=10.0)

        # Send the FileInfo command.
        print(f"BeginFlash command sent successfully, sending file info...")
        can_id = self.motor._msg_can_id(MotorMsg.FlashFileInfo, self.motor.motor_id, self.main_host_id)
        data = S_uint32le.pack(file_size) + S_uint32le.pack(num_packets)
        self._send_await(MotorMsg.FlashFileInfo, self.comms.cansend, self.iface, True, can_id, data, m_uid)

        # Send the file data packets.
        print(f"Sending {num_packets} file data packets...", flush=True)
        for i in progressbar(range(num_packets), progress=progress):
            # Get the next 8 bytes of the firmware data, and pad with 0xff if necessary.
            data = firmware_data[i * 8:(i + 1) * 8]
            if len(data) != 8:
                assert i == num_packets - 1, f"bad num bytes: {i} != {num_packets - 1}, {len(data)} != 8"
                data += bytes([0xff] * (8 - len(data)))
            
            idx_bytes = S_uint16be.pack(i)  # Pack as big endian, but host/extra bytes are also reversed below
            motor_id = self.motor.motor_id
            can_id = self.motor._msg_can_id(MotorMsg.FlashData, motor_id, host_id=idx_bytes[1:], e_byte=idx_bytes[0])

            self._send_await(MotorMsg.FlashData, self.comms.cansend, self.iface, True, can_id, data, m_uid)

        # Send the EndFlash command.
        print(f"File data packets sent successfully, sending EndFlash command...")
        can_id = self.motor._msg_can_id(MotorMsg.EndFlash, self.motor.motor_id, host_id=0x00)
        data = S_uint32le.pack(num_packets) + bytes([0] * 4)
        self._send_await(MotorMsg.EndFlash, self.comms.cansend, self.iface, True, can_id, data, m_uid)

        print(f"Firmware flash completed!")
    
    @prepare_robstride_admin_command
    def get_limits(self, print_info: bool = True) -> Dict[str, Tuple[float, float]]:
        """Get a dictionary of the (min, max) limits for each robstride param for this motor."""
        max_val = float('inf')
        params = {name: [None, None] for name, p in _ROBSTRIDE_PARAMS.as_dict.items() if p.mode == "rw" and p.dtype.is_float}

        # Put into operation mode so we don't accidentally move the motor
        self._send_await(MotorMsg.MotorFeedback, self.motor.disable)
        self._send_await(MotorMsg.MotorFeedback, self.motor.set_run_mode, RobstrideRunMode.Operation)
        self._send_await(MotorMsg.ReadParam, self.motor.read_param, "run_mode")
        assert self.motor.run_mode == RobstrideRunMode.Operation

        sleep_time = 0.1
        for name in params.keys():
            if print_info: 
                print(f"Getting limits for {name}")
            for max_mult, idx in [(1, 1), (-1, 0)]:
                val = max_mult * max_val

                # Get bounded value. Use _write_param_checked() to bypass the safety checks for out of range values.
                self.motor.update_property(name, None)
                self._send_await(MotorMsg.MotorFeedback, self.motor._write_param_checked, name, val, True)
                self._send_await(MotorMsg.ReadParam, self.motor.read_param, name)
                time.sleep(sleep_time)

                if self.motor.properties[name].value is not None:
                    params[name][idx] = self.motor.properties[name].value
                else:
                    print(f"Failed to get {'max' if idx == 1 else 'min'} value for {name}")
                
            # Reset to default
            default = _ROBSTRIDE_PARAMS[name].default[self.motor.version]
            self._send_await(MotorMsg.MotorFeedback, self.motor._write_param_checked, name, default, True)
        
        return params
    
    @prepare_robstride_admin_command
    def read_current_firmware_version(self) -> str:
        """Read the current firmware version of the motor."""
        return self._read_firmware_version()
    
    def _read_firmware_version(self) -> str:
        """Helper method since we also call this for flashing"""
        m_uid = MotorUIDType(self.motor.uid)
        can_id = self.motor._msg_can_id(MotorMsg.Disable, self.motor.motor_id, host_id=self.main_host_id)
        data = b'\x00\xC4\x00\x00\x00\x00\x00\x00'
        result = self._send_await(MotorMsg.MotorFeedback, self.comms.cansend, self.iface, True, can_id, data, m_uid)

        assert result.data[:3] == b'\x00\xC4V', f"Expected first three bytes b'\\x00\\xC4V', got: {result.data[:3]}"

        # The last byte looks like it's just always '7', also the case in firmware binaries. Guessing this command just
        # reads 8 consecutive bytes and that '7' is used for something else
        return f"{result.data[3]}.{result.data[4]}.{result.data[5]}.{result.data[6]}"
    
    @prepare_robstride_admin_command
    def write_all_param_defaults(self) -> None:
        """Write default values for all parameters to the motor. Will disable the motor and set operation mode first."""
        # Disable the motor and set operation mode to avoid accidentally moving it
        self._send_await(MotorMsg.MotorFeedback, self.motor.disable)
        self._send_await(MotorMsg.MotorFeedback, self.motor.set_run_mode, RobstrideRunMode.Operation)
        self._send_await(MotorMsg.ReadParam, self.motor.read_param, "run_mode")
        assert self.motor.run_mode == RobstrideRunMode.Operation

        for name, p in _ROBSTRIDE_PARAMS.as_dict.items():
            if p.writable:
                self._send_await(MotorMsg.MotorFeedback, self.motor.write_param, name, p.default[self.motor.version])
    