from robot.rokae.rokae_robot import RokaeRobot, RokaeXmateRobot
# Key: robot vendor; Value: robot class


rokae_classes = {
    "ROKAE_XMATE7": RokaeXmateRobot
}


def rokae_factory(robot_host, robot_type, **kwargs):
    if robot_type.startswith(("ROKAE_CR", "ROKAE_ER", "ROKAE_SR")):
        rokae_classes[robot_type] = RokaeXmateRobot
    return rokae_classes.get(robot_type, RokaeRobot)(robot_host, robot_type=robot_type, **kwargs)


robot_dict = {RokaeRobot.ROBOT_VENDOR: rokae_factory}