import logging
from functools import cached_property
import time
from typing import Any
from lerobot.cameras.utils import make_cameras_from_configs
from .motors_bus import RealmanMotorsBus
from .config import RealmanRobotConfig
from lerobot.robots.robot import Robot
from lerobot.robots.utils import ensure_safe_goal_position
from lerobot.utils.errors import DeviceAlreadyConnectedError, DeviceNotConnectedError

logger = logging.getLogger(__name__)
class RealmanRobot(Robot):
    config_class = RealmanRobotConfig
    name = "realman"
    def __init__(self, config: RealmanRobotConfig):
        super().__init__(config)
        self.config = config
        self.bus = RealmanMotorsBus(
            port=self.config.port,
            motors=config.motors,
            joint=config.joint,
            mock=self.config.mock,
            gripper_range=config.gripper_range,
            calibration=self.calibration,
        )
        self.cameras = make_cameras_from_configs(config.cameras)
    @property
    def _motors_ft(self) -> dict[str, type]:
        return {f"{motor}.pos": float for motor in self.bus.motors}

    @property
    def _cameras_ft(self) -> dict[str, tuple]:
        return {
            cam: (self.config.cameras[cam].height, self.config.cameras[cam].width, 3) for cam in self.cameras
        }

    @cached_property
    def observation_features(self) -> dict[str, type | tuple]:
        return {**self._motors_ft, **self._cameras_ft}

    @cached_property
    def action_features(self) -> dict[str, type]:
        if self.config.mode==0:
            return self._motors_ft
        else:
            return {"x.vel":float,"y.vel":float,"z.vel":float,"rx.vel":float,"ry.vel":float,"rz.vel":float,"gripper.vel":float}

    @property
    def is_connected(self) -> bool:
        return self.bus.is_connected and all(cam.is_connected for cam in self.cameras.values())
    
    def connect(self, calibrate: bool = True) -> None:
        """
        We assume that at connection time, arm is in a rest position,
        and torque can be safely disabled to run calibration.
        """
        if self.is_connected:
            raise DeviceAlreadyConnectedError(f"{self} already connected")

        self.bus.connect()

        for cam in self.cameras.values():
            cam.connect()

        if not self.is_calibrated and calibrate:
            logger.info(
                "Mismatch between calibration values in the motor and the calibration file or no calibration file found"
            )
            self.calibrate()
        # self.configure()
        logger.info(f"{self} connected.")
    
    @property
    def is_calibrated(self) -> bool:
        return self.bus.is_calibrated
    
    def calibrate(self) -> None:
        """move piper to the home position"""
        if not self.is_connected:
            raise ConnectionError()
        
        self.bus.apply_calibration()
    
    def configure(self) -> None:
        with self.bus.torque_disabled():
            self.bus.configure_motors()
    
    def setup_motors(self) -> None:
        pass

    def get_observation(self) -> dict[str, Any]:
        if not self.is_connected:
            raise DeviceNotConnectedError(f"{self} is not connected.")

        # Read arm position
        start = time.perf_counter()
        obs_dict = self.bus.sync_read("Present_Position")
        obs_dict = {f"{motor}.pos": val for motor, val in obs_dict.items()}
        dt_ms = (time.perf_counter() - start) * 1e3
        logger.debug(f"{self} read state: {dt_ms:.1f}ms")

        # Capture images from cameras
        for cam_key, cam in self.cameras.items():
            start = time.perf_counter()
            obs_dict[cam_key] = cam.async_read()
            dt_ms = (time.perf_counter() - start) * 1e3
            logger.debug(f"{self} read {cam_key}: {dt_ms:.1f}ms")

        return obs_dict
    
    def send_action(self, action: dict[str, Any]) -> dict[str, Any]:
        """Command arm to move to a target joint configuration.

        The relative action magnitude may be clipped depending on the configuration parameter
        `max_relative_target`. In this case, the action sent differs from original action.
        Thus, this function always returns the action actually sent.

        Raises:
            RobotDeviceNotConnectedError: if robot is not connected.

        Returns:
            the action sent to the motors, potentially clipped.
        """
        if not self.is_connected:
            raise DeviceNotConnectedError(f"{self} is not connected.")
        if self.config.mode==0:

            goal_pos = {key.removesuffix(".pos"): val for key, val in action.items() if key.endswith(".pos")}

            if len(goal_pos)>0:
                # Cap goal position when too far away from present position.
                # /!\ Slower fps expected due to reading from the follower.
                if self.config.max_relative_target is not None:
                    present_pos = self.bus.sync_read("Present_Position")
                    goal_present_pos = {key: (g_pos, present_pos[key]) for key, g_pos in goal_pos.items()}
                    goal_pos = ensure_safe_goal_position(goal_present_pos, self.config.max_relative_target)

                # Send goal position to the arm
                self.bus.sync_write("Goal_Position", goal_pos)
            return {f"{motor}.pos": val for motor, val in goal_pos.items()}
        else:
            valAction = {key.removesuffix(".vel"): val for key, val in action.items() if key.endswith(".vel")}
            #笛卡尔速度透传
            self.bus.sync_write("Goal_Velocity", valAction)
            return action

    def disconnect(self):
        if not self.is_connected:
            raise DeviceNotConnectedError(f"{self} is not connected.")

        self.bus.disconnect(self.config.disable_torque_on_disconnect)
        for cam in self.cameras.values():
            cam.disconnect()

        logger.info(f"{self} disconnected.")