""" URDF is xml description of robot shape used ny ROS internals.
    generate_urdf generates this file based on provided config.
    Also, each defined scad module is rendered into STL and used.
    """
import hashlib

from lib.config import Distance, Angle, Unit, Mass

from lxml import etree


def generate_link(xml_root, name, cfg):
    """ generates xml link element, renders .stl file from scad source if necessary """

    link = etree.SubElement(xml_root, 'link', name=name)

    #if 'mass' in cfg:
    #    inertial = etree.SubElement(link, 'inertial')
    #    etree.SubElement(inertial, 'origin', xyz="{} {} {}".format(*unit_convert(cfg['mass_center'])))
    #     mass = etree.SubElement(inertial, 'mass',value = str(unit_convert(cfg['mass'])))
    # -> TODO: inertia matrix needs to be implemented

    if 'scad' in cfg:

        visual = etree.SubElement(link, 'visual')

        geometry = etree.SubElement(visual, 'geometry')

        stl_filename = 'stl/' + name + ".stlb"

        etree.SubElement(
            geometry, 'mesh', filename="package://schpin_tote/" + stl_filename
        )

        material = etree.SubElement(
            visual, 'material', name=name + "_material"
        )

        etree.SubElement(
            material, 'color', rgba="{} {} {} 1".format(*get_color(name))
        )

    return link


def get_color(name):
    """ Generates random color based on name """
    m = hashlib.md5()
    m.update(name.encode('utf8'))
    h = m.digest()
    return [float(ord(x)) / 256 for x in h[:3]]


def unit_convert(x):

    if isinstance(x, Distance):
        return x.as_m()
    elif isinstance(x, Angle):
        return x.as_radians()
    elif isinstance(x, Mass):
        return x.as_kg()
    elif isinstance(x, Unit):
        raise Exception("don't know how to handle unit)")
    elif isinstance(x, (tuple, list)):
        return [unit_convert(y) for y in x]
    else:
        return x


def generate_urdf_joint(xml_root, name, parent, child, cfg, offset=None):
    """ Generates joint xml elemnt in xml_root """
    if 'joint' in cfg:
        joint_t = 'revolute'
    else:
        joint_t = 'fixed'
    joint = etree.SubElement(xml_root, 'joint', name=name, type=joint_t)

    if 'joint' in cfg:
        etree.SubElement(
            joint,
            'limit',
            effort="1000",
            velocity="50",
            lower=str(float(cfg['joint']['angle']['min'])),
            upper=str(float(cfg['joint']['angle']['max']))
        )
        if 'axis' in cfg['joint']:
            etree.SubElement(
                joint, 'axis', xyz="{} {} {}".format(*cfg['joint']['axis'])
            )

    etree.SubElement(joint, 'parent', link=parent.get('name'))
    etree.SubElement(joint, 'child', link=child.get('name'))

    pos = list(cfg['pos'])
    rot = list(cfg['rot'])

    if offset is not None:
        for i in range(3):
            pos[i] += offset['pos'][i]
            rot[i] += offset['rot'][i]

    etree.SubElement(
        joint,
        'origin',
        xyz="{} {} {}".format(*unit_convert(pos)),
        rpy="{} {} {}".format(*unit_convert(rot))
    )


def generate_urdf_chain(
    xml_root, prefix, segments, name, prev_link, offset=None
):
    """ generates link/name chain in xml_root, uses prefix for nameprefix, 
        picks segment based on 'name' and in case next_segment is 
        defined for segment, recursively executes again.
        
        Link is generated, and joint is than generated betwen
        prev_link and new link."""

    cfg = segments[name]

    link = generate_link(xml_root, name=prefix + "/" + name, cfg=cfg)

    generate_urdf_joint(
        xml_root,
        name=prefix + "/" + name,
        parent=prev_link,
        child=link,
        cfg=cfg,
        offset=offset
    )

    if 'next_segment' in cfg:
        generate_urdf_chain(
            xml_root, prefix, segments, cfg['next_segment'], link
        )


def generate_urdf(filename, cfg):
    """ generates urdf file for tote robot, generates full link/joint
        chain for all legs and gimbal """

    robot = etree.Element('robot', name='Tote')

    base_link = generate_link(
        robot, 'base_link', {'scad': {
            'module': 'body',
            'args': []
        }}
    )

    for leg, leg_cfg in cfg['legs'].items():

        generate_urdf_chain(
            robot,
            "leg/" + leg,  # as name prefix
            leg_cfg['segments'],
            leg_cfg['chain_root'],
            base_link,
            offset={"pos": leg_cfg['pos'],
                    "rot": leg_cfg['rot']}
        )

    generate_urdf_chain(
        robot, 'gimbal', cfg['gimbal']['segments'], cfg['gimbal']['chain_root'],
        base_link
    )

    tree = etree.ElementTree(robot)
    tree.write(filename, pretty_print=True)
