# json_to_urdf.py
import json
import math
import xml.etree.ElementTree as ET
from pathlib import Path
from xml.dom import minidom

# ============================================================
# CONFIG
# ============================================================

ROBOT_JSON = "../../data/robot/robot.json"
OUTPUT_URDF = "robot.urdf"

# ============================================================
# HELPERS
# ============================================================


def mm_to_m(v):
    return float(v) / 1000.0


def vec_mm_to_m(v):
    return [mm_to_m(x) for x in v]


def deg_to_rad(v):
    return math.radians(float(v))


def vec_deg_to_rad(v):
    return [deg_to_rad(x) for x in v]


def xyz_str(v):
    return f"{v[0]:.6f} {v[1]:.6f} {v[2]:.6f}"


# ============================================================
# LOAD JSON
# ============================================================

with open(ROBOT_JSON, "r", encoding="utf-8") as f:
    robot_json = json.load(f)

links = robot_json.get("links", {})

# ============================================================
# ROOT
# ============================================================

robot_name = Path(ROBOT_JSON).stem
robot = ET.Element("robot")
robot.set("name", robot_name)

# ============================================================
# CREATE LINKS
# ============================================================

for link_name, link_info in links.items():

    urdf_link = ET.SubElement(robot, "link")
    urdf_link.set("name", link_name)

    # --------------------------------------------------------
    # VISUALS
    # --------------------------------------------------------

    models = link_info.get("model", [])

    for model in models:

        stl_file = model.get("stlFile")
        if not stl_file:
            continue

        visual = ET.SubElement(urdf_link, "visual")

        origin = ET.SubElement(visual, "origin")

        xyz = vec_mm_to_m(model.get("originOfModel", [0, 0, 0]))
        rpy = vec_deg_to_rad(model.get("rotationOfModelDegree", [0, 0, 0]))

        origin.set("xyz", xyz_str(xyz))
        origin.set("rpy", xyz_str(rpy))

        geometry = ET.SubElement(visual, "geometry")

        mesh = ET.SubElement(geometry, "mesh")
        mesh.set("filename", stl_file)
        mesh.set("scale", "0.001 0.001 0.001")

    # --------------------------------------------------------
    # SIMPLE COLLISION FROM SKELETON
    # --------------------------------------------------------

    skeleton = link_info.get("skeleton")

    if isinstance(skeleton, dict):

        p1 = skeleton.get("from", [0, 0, 0])
        p2 = skeleton.get("to", [0, 0, 0])

        dx = p2[0] - p1[0]
        dy = p2[1] - p1[1]
        dz = p2[2] - p1[2]

        length_mm = math.sqrt(dx * dx + dy * dy + dz * dz)
        length_m = mm_to_m(length_mm)

        radius_mm = skeleton.get("radius", 4)
        radius_m = mm_to_m(radius_mm)

        cx = (p1[0] + p2[0]) * 0.5
        cy = (p1[1] + p2[1]) * 0.5
        cz = (p1[2] + p2[2]) * 0.5

        collision = ET.SubElement(urdf_link, "collision")

        collision_origin = ET.SubElement(collision, "origin")
        collision_origin.set(
            "xyz",
            xyz_str(vec_mm_to_m([cx, cy, cz]))
        )

        # ----------------------------------------------------
        # CYLINDER ORIENTATION
        # URDF cylinder points along Z
        # ----------------------------------------------------

        roll = 0.0
        pitch = 0.0
        yaw = 0.0

        if abs(dx) >= abs(dy) and abs(dx) >= abs(dz):
            pitch = math.radians(90)
            yaw = 0.0

        elif abs(dy) >= abs(dx) and abs(dy) >= abs(dz):
            pitch = math.radians(90)
            yaw = math.radians(90)

        collision_origin.set(
            "rpy",
            xyz_str([roll, pitch, yaw])
        )

        geometry = ET.SubElement(collision, "geometry")

        cylinder = ET.SubElement(geometry, "cylinder")
        cylinder.set("radius", f"{radius_m:.6f}")
        cylinder.set("length", f"{length_m:.6f}")

# ============================================================
# CREATE JOINTS
# ============================================================

for link_name, link_info in links.items():

    parent_name = link_info.get("parent")

    if parent_name is None:
        continue

    joint_info = link_info.get("jointToParent", {})

    joint_name = joint_info.get("name", f"joint_{parent_name}_{link_name}")
    joint_type = joint_info.get("type", "fixed")

    if joint_type == "linear":
        urdf_joint_type = "prismatic"
    elif joint_type == "revolute":
        urdf_joint_type = "revolute"
    else:
        urdf_joint_type = "fixed"

    joint = ET.SubElement(robot, "joint")
    joint.set("name", joint_name)
    joint.set("type", urdf_joint_type)

    # --------------------------------------------------------
    # PARENT / CHILD
    # --------------------------------------------------------

    parent = ET.SubElement(joint, "parent")
    parent.set("link", parent_name)

    child = ET.SubElement(joint, "child")
    child.set("link", link_name)

    # --------------------------------------------------------
    # ORIGIN
    # --------------------------------------------------------

    origin_xyz = joint_info.get("origin", [0, 0, 0])
    origin_rpy = joint_info.get("rotation", [0, 0, 0])

    origin = ET.SubElement(joint, "origin")
    origin.set(
        "xyz",
        xyz_str(vec_mm_to_m(origin_xyz))
    )
    origin.set(
        "rpy",
        xyz_str(vec_deg_to_rad(origin_rpy))
    )

    # --------------------------------------------------------
    # AXIS
    # --------------------------------------------------------

    axis = joint_info.get("axis", [1, 0, 0])

    axis_tag = ET.SubElement(joint, "axis")
    axis_tag.set("xyz", xyz_str(axis))

    # --------------------------------------------------------
    # DEFAULT LIMITS
    # --------------------------------------------------------

    if urdf_joint_type in ["revolute", "prismatic"]:

        limit = ET.SubElement(joint, "limit")

        if urdf_joint_type == "revolute":
            limit.set("lower", f"{-math.pi:.6f}")
            limit.set("upper", f"{math.pi:.6f}")
        else:
            limit.set("lower", "-1.0")
            limit.set("upper", "1.0")

        limit.set("effort", "100")
        limit.set("velocity", "10")

# ============================================================
# PRETTY XML
# ============================================================

xml_string = ET.tostring(robot, encoding="utf-8")
pretty = minidom.parseString(xml_string).toprettyxml(indent="  ")

with open(OUTPUT_URDF, "w", encoding="utf-8") as f:
    f.write(pretty)

print("URDF written:", OUTPUT_URDF)