from __future__ import annotations

import xml.etree.ElementTree as ET
from pathlib import Path
from xml.dom import minidom

from app.models import RobotDraft


def _fmt(nums: list[float]) -> str:
    return " ".join(f"{float(x):.6g}" for x in nums)


def build_urdf_xml(draft: RobotDraft, mesh_dir_uri: str = "package://robot_description/meshes") -> str:
    robot = ET.Element("robot", name=draft.name)

    for link in draft.links:
        link_el = ET.SubElement(robot, "link", name=link.name)
        visual = ET.SubElement(link_el, "visual")
        ET.SubElement(visual, "origin", xyz="0 0 0", rpy="0 0 0")
        geom = ET.SubElement(visual, "geometry")
        mesh_name = link.mesh or f"{link.name}.stl"
        ET.SubElement(geom, "mesh", filename=f"{mesh_dir_uri}/{mesh_name}", scale="0.001 0.001 0.001")
        coll = ET.SubElement(link_el, "collision")
        ET.SubElement(coll, "origin", xyz="0 0 0", rpy="0 0 0")
        cgeom = ET.SubElement(coll, "geometry")
        ET.SubElement(cgeom, "mesh", filename=f"{mesh_dir_uri}/{mesh_name}", scale="0.001 0.001 0.001")
        inertial = ET.SubElement(link_el, "inertial")
        ET.SubElement(inertial, "origin", xyz="0 0 0", rpy="0 0 0")
        ET.SubElement(inertial, "mass", value="1.0")
        ET.SubElement(
            inertial,
            "inertia",
            ixx="0.001",
            ixy="0",
            ixz="0",
            iyy="0.001",
            iyz="0",
            izz="0.001",
        )

    for joint in draft.joints:
        j_el = ET.SubElement(robot, "joint", name=joint.name, type=joint.joint_type)
        ET.SubElement(j_el, "parent", link=joint.parent)
        ET.SubElement(j_el, "child", link=joint.child)
        ET.SubElement(j_el, "origin", xyz=_fmt(joint.origin_xyz), rpy=_fmt(joint.origin_rpy))
        if joint.joint_type in {"revolute", "prismatic", "continuous"}:
            ET.SubElement(j_el, "axis", xyz=_fmt(joint.axis))
        if joint.joint_type in {"revolute", "prismatic"}:
            ET.SubElement(
                j_el,
                "limit",
                lower=str(joint.lower),
                upper=str(joint.upper),
                effort=str(joint.effort),
                velocity=str(joint.velocity),
            )

    rough = ET.tostring(robot, encoding="utf-8")
    pretty = minidom.parseString(rough).toprettyxml(indent="  ")
    # drop extra XML declaration newline quirks
    lines = [ln for ln in pretty.splitlines() if ln.strip()]
    return "\n".join(lines) + "\n"


def write_urdf_package(draft: RobotDraft, out_dir: Path) -> dict[str, str]:
    out_dir.mkdir(parents=True, exist_ok=True)
    meshes = out_dir / "meshes"
    meshes.mkdir(exist_ok=True)
    urdf_path = out_dir / f"{draft.name}.urdf"
    urdf_xml = build_urdf_xml(draft, mesh_dir_uri="meshes")
    urdf_path.write_text(urdf_xml, encoding="utf-8")

    # Placeholder note for missing meshes
    readme = out_dir / "README.txt"
    readme.write_text(
        "\n".join(
            [
                f"Robot: {draft.name}",
                f"Profile: {draft.profile}",
                f"Links: {len(draft.links)}",
                f"Joints: {len(draft.joints)}",
                "",
                "Mesh files are expected under ./meshes as <link>.stl.",
                "If meshes were not generated, URDF still validates structurally;",
                "install OCP/cadquery or cascadio later for STEP tessellation.",
                "",
                "Notes:",
                *[f"- {n}" for n in draft.notes],
            ]
        )
        + "\n",
        encoding="utf-8",
    )
    draft_json = out_dir / "draft.json"
    draft_json.write_text(draft.model_dump_json(indent=2), encoding="utf-8")
    return {
        "urdf": str(urdf_path),
        "readme": str(readme),
        "draft_json": str(draft_json),
        "meshes_dir": str(meshes),
    }