#!/usr/bin/env python3
"""Visualize URDF collision geometry and optional visual meshes in MeshCat."""

from __future__ import annotations

import argparse
from dataclasses import dataclass
import math
import os
from pathlib import Path
import subprocess
import sys
import time
import traceback
import webbrowser
import xml.etree.ElementTree as ET

import meshcat
import meshcat.geometry as geometry
import numpy as np


@dataclass(frozen=True)
class Joint:
    name: str
    kind: str
    parent: str
    child: str
    origin: np.ndarray
    axis: np.ndarray
    mimic: tuple[str, float, float] | None


def parse_vector(
    value: str | None, size: int, default: tuple[float, ...]
) -> np.ndarray:
    if value is None:
        return np.asarray(default, dtype=np.float64)
    parts = value.split()
    if len(parts) != size:
        raise ValueError(f"expected {size} values, got {value!r}")
    result = np.asarray([float(part) for part in parts], dtype=np.float64)
    if not np.isfinite(result).all():
        raise ValueError(f"values contain NaN or infinity: {value!r}")
    return result


def rpy_rotation(rpy: np.ndarray) -> np.ndarray:
    roll, pitch, yaw = rpy
    cr, sr = math.cos(roll), math.sin(roll)
    cp, sp = math.cos(pitch), math.sin(pitch)
    cy, sy = math.cos(yaw), math.sin(yaw)
    return np.array(
        [
            [cy * cp, cy * sp * sr - sy * cr, cy * sp * cr + sy * sr],
            [sy * cp, sy * sp * sr + cy * cr, sy * sp * cr - cy * sr],
            [-sp, cp * sr, cp * cr],
        ],
        dtype=np.float64,
    )


def origin_transform(origin: ET.Element | None) -> np.ndarray:
    transform = np.eye(4)
    if origin is None:
        return transform
    transform[:3, :3] = rpy_rotation(
        parse_vector(origin.get("rpy"), 3, (0.0, 0.0, 0.0))
    )
    transform[:3, 3] = parse_vector(
        origin.get("xyz"), 3, (0.0, 0.0, 0.0)
    )
    return transform


def axis_angle_transform(axis: np.ndarray, angle: float) -> np.ndarray:
    norm = float(np.linalg.norm(axis))
    if norm < 1e-12:
        raise ValueError("joint axis must not be zero")
    x, y, z = axis / norm
    c, s = math.cos(angle), math.sin(angle)
    one_minus_c = 1.0 - c
    transform = np.eye(4)
    transform[:3, :3] = np.array(
        [
            [c + x * x * one_minus_c, x * y * one_minus_c - z * s, x * z * one_minus_c + y * s],
            [y * x * one_minus_c + z * s, c + y * y * one_minus_c, y * z * one_minus_c - x * s],
            [z * x * one_minus_c - y * s, z * y * one_minus_c + x * s, c + z * z * one_minus_c],
        ],
        dtype=np.float64,
    )
    return transform


def translation_transform(offset: np.ndarray) -> np.ndarray:
    transform = np.eye(4)
    transform[:3, 3] = offset
    return transform


def parse_joint_values(values: list[str]) -> dict[str, float]:
    result: dict[str, float] = {}
    for assignment in values:
        name, separator, raw_value = assignment.partition("=")
        if not separator or not name or not raw_value:
            raise ValueError(
                f"invalid --joint value {assignment!r}; expected NAME=VALUE"
            )
        if name in result:
            raise ValueError(f"joint value specified more than once: {name}")
        value = float(raw_value)
        if not math.isfinite(value):
            raise ValueError(f"joint value must be finite: {assignment!r}")
        result[name] = value
    return result


def parse_joints(root: ET.Element) -> dict[str, Joint]:
    joints: dict[str, Joint] = {}
    children: set[str] = set()
    for element in root.findall("joint"):
        name = element.get("name")
        kind = element.get("type")
        parent_element = element.find("parent")
        child_element = element.find("child")
        if not name or not kind or parent_element is None or child_element is None:
            raise ValueError("every joint needs name, type, parent, and child")
        parent = parent_element.get("link")
        child = child_element.get("link")
        if not parent or not child:
            raise ValueError(f"joint {name!r} has an empty parent or child")
        if name in joints:
            raise ValueError(f"duplicate joint name: {name}")
        if child in children:
            raise ValueError(f"link {child!r} has more than one parent joint")
        axis_element = element.find("axis")
        axis = parse_vector(
            axis_element.get("xyz") if axis_element is not None else None,
            3,
            (1.0, 0.0, 0.0),
        )
        mimic_element = element.find("mimic")
        mimic = None
        if mimic_element is not None:
            source = mimic_element.get("joint")
            if not source:
                raise ValueError(f"mimic joint {name!r} has no source joint")
            mimic = (
                source,
                float(mimic_element.get("multiplier", "1")),
                float(mimic_element.get("offset", "0")),
            )
        joints[name] = Joint(
            name=name,
            kind=kind,
            parent=parent,
            child=child,
            origin=origin_transform(element.find("origin")),
            axis=axis,
            mimic=mimic,
        )
        children.add(child)
    return joints


def resolve_joint_values(
    joints: dict[str, Joint], requested: dict[str, float]
) -> dict[str, float]:
    unknown = set(requested) - set(joints)
    if unknown:
        raise ValueError(f"unknown joints: {sorted(unknown)}")
    fixed = [name for name in requested if joints[name].kind == "fixed"]
    if fixed:
        raise ValueError(f"fixed joints cannot be assigned: {sorted(fixed)}")

    resolved: dict[str, float] = {}

    def resolve(name: str, stack: set[str]) -> float:
        if name in resolved:
            return resolved[name]
        if name in stack:
            raise ValueError(f"mimic joint cycle contains {name!r}")
        joint = joints[name]
        if name in requested:
            value = requested[name]
        elif joint.mimic is not None:
            source, multiplier, offset = joint.mimic
            if source not in joints:
                raise ValueError(
                    f"mimic joint {name!r} references unknown joint {source!r}"
                )
            value = multiplier * resolve(source, stack | {name}) + offset
        else:
            value = 0.0
        resolved[name] = value
        return value

    for joint_name in joints:
        resolve(joint_name, set())
    return resolved


def joint_motion(joint: Joint, value: float) -> np.ndarray:
    if joint.kind == "fixed":
        return np.eye(4)
    if joint.kind in ("revolute", "continuous"):
        return axis_angle_transform(joint.axis, value)
    if joint.kind == "prismatic":
        norm = float(np.linalg.norm(joint.axis))
        if norm < 1e-12:
            raise ValueError(f"joint {joint.name!r} axis must not be zero")
        axis = joint.axis / norm
        return translation_transform(axis * value)
    raise ValueError(
        f"joint {joint.name!r} uses unsupported type {joint.kind!r}; "
        "supported types are fixed, revolute, continuous, and prismatic"
    )


def compute_link_transforms(
    root: ET.Element, joints: dict[str, Joint], values: dict[str, float]
) -> dict[str, np.ndarray]:
    links = {element.get("name") for element in root.findall("link")}
    if None in links:
        raise ValueError("every link needs a name")
    children = {joint.child for joint in joints.values()}
    roots = links - children
    if not roots:
        raise ValueError("URDF has no root link")

    transforms = {name: np.eye(4) for name in roots}
    pending = list(joints.values())
    while pending:
        unresolved: list[Joint] = []
        for joint in pending:
            if joint.parent not in links or joint.child not in links:
                raise ValueError(
                    f"joint {joint.name!r} references a missing parent or child link"
                )
            if joint.parent not in transforms:
                unresolved.append(joint)
                continue
            transforms[joint.child] = (
                transforms[joint.parent]
                @ joint.origin
                @ joint_motion(joint, values[joint.name])
            )
        if len(unresolved) == len(pending):
            names = [joint.name for joint in unresolved]
            raise ValueError(f"joint graph is cyclic or disconnected: {names}")
        pending = unresolved
    return transforms


def resolve_mesh_path(filename: str, urdf_path: Path) -> Path:
    if filename.startswith("file://"):
        path = Path(filename.removeprefix("file://"))
    elif filename.startswith("package://"):
        package_path = Path(filename.removeprefix("package://"))
        if len(package_path.parts) < 2:
            raise ValueError(f"invalid package URI: {filename}")
        package_name, relative_parts = package_path.parts[0], package_path.parts[1:]
        candidates = [
            parent / package_name / Path(*relative_parts)
            for parent in (urdf_path.parent, *urdf_path.parents)
        ]
        candidates.extend(
            parent / Path(*relative_parts)
            for parent in urdf_path.parents
            if parent.name == package_name
        )
        for candidate in candidates:
            if candidate.is_file():
                return candidate.resolve()
        raise FileNotFoundError(
            f"could not resolve {filename!r} relative to {urdf_path}"
        )
    else:
        path = Path(filename)
    if not path.is_absolute():
        path = urdf_path.parent / path
    path = path.resolve()
    if not path.is_file():
        raise FileNotFoundError(path)
    return path


def load_mesh(path: Path) -> geometry.Geometry:
    suffix = path.suffix.lower()
    if suffix == ".stl":
        return geometry.StlMeshGeometry.from_file(str(path))
    if suffix == ".obj":
        return geometry.ObjMeshGeometry.from_file(str(path))
    if suffix == ".dae":
        return geometry.DaeMeshGeometry.from_file(str(path))
    raise ValueError(
        f"unsupported mesh format {path.suffix!r}: {path}; "
        "MeshCat viewer supports STL, OBJ, and DAE"
    )


def cylinder_dimensions(cylinder: ET.Element) -> tuple[float, float]:
    radius = float(cylinder.get("radius", "nan"))
    length = float(cylinder.get("length", "nan"))
    if (
        not math.isfinite(radius)
        or not math.isfinite(length)
        or radius <= 0.0
        or length <= 0.0
    ):
        raise ValueError(
            f"cylinder radius and length must be positive: {radius}, {length}"
        )
    return radius, length


def cylinder_correction() -> np.ndarray:
    correction = np.eye(4)
    # Three.js cylinders use local Y; URDF cylinders use local Z.
    correction[:3, :3] = rpy_rotation(np.array([math.pi / 2.0, 0.0, 0.0]))
    return correction


def cylinder_wireframe(
    radius: float,
    length: float,
    generator_count: int,
    color: int,
    opacity_value: float,
    ring_segments: int = 64,
) -> geometry.LineSegments:
    vertices: list[tuple[float, float, float]] = []
    half_length = length / 2.0

    # Smooth top and bottom rings, without cap triangulation spokes.
    for y in (-half_length, half_length):
        for index in range(ring_segments):
            first = 2.0 * math.pi * index / ring_segments
            second = 2.0 * math.pi * (index + 1) / ring_segments
            vertices.extend(
                [
                    (radius * math.cos(first), y, radius * math.sin(first)),
                    (radius * math.cos(second), y, radius * math.sin(second)),
                ]
            )

    # Sparse axial generator lines; six means one line every 60 degrees.
    for index in range(generator_count):
        angle = 2.0 * math.pi * index / generator_count
        x = radius * math.cos(angle)
        z = radius * math.sin(angle)
        vertices.extend([(x, -half_length, z), (x, half_length, z)])

    points = np.asarray(vertices, dtype=np.float32).T
    material = geometry.LineBasicMaterial(
        color=color,
        transparent=opacity_value < 1.0,
        opacity=opacity_value,
    )
    return geometry.LineSegments(geometry.PointsGeometry(points), material)


def geometry_object(
    geometry_element: ET.Element, urdf_path: Path
) -> tuple[geometry.Geometry, np.ndarray]:
    box = geometry_element.find("box")
    sphere = geometry_element.find("sphere")
    cylinder = geometry_element.find("cylinder")
    mesh = geometry_element.find("mesh")
    correction = np.eye(4)

    if box is not None:
        size = parse_vector(box.get("size"), 3, ())
        if (size <= 0.0).any():
            raise ValueError(f"box size must be positive: {size}")
        return geometry.Box(size), correction
    if sphere is not None:
        radius = float(sphere.get("radius", "nan"))
        if not math.isfinite(radius) or radius <= 0.0:
            raise ValueError(f"sphere radius must be positive: {radius}")
        return geometry.Sphere(radius), correction
    if cylinder is not None:
        radius, length = cylinder_dimensions(cylinder)
        return geometry.Cylinder(length, radius), cylinder_correction()
    if mesh is not None:
        filename = mesh.get("filename")
        if not filename:
            raise ValueError("mesh geometry has no filename")
        scale = parse_vector(mesh.get("scale"), 3, (1.0, 1.0, 1.0))
        correction[:3, :3] = np.diag(scale)
        return load_mesh(resolve_mesh_path(filename, urdf_path)), correction
    raise ValueError("geometry must contain box, sphere, cylinder, or mesh")


def parse_rgba(value: str) -> tuple[float, float, float, float]:
    rgba = parse_vector(value, 4, ())
    if ((rgba < 0.0) | (rgba > 1.0)).any():
        raise ValueError(f"RGBA values must be between 0 and 1: {value!r}")
    return tuple(float(component) for component in rgba)


def rgb_integer(rgb: tuple[float, float, float]) -> int:
    red, green, blue = (round(component * 255.0) for component in rgb)
    return (red << 16) | (green << 8) | blue


def visual_rgba(
    visual: ET.Element,
    named_materials: dict[str, tuple[float, float, float, float]],
) -> tuple[float, float, float, float]:
    material = visual.find("material")
    if material is None:
        return (0.65, 0.68, 0.72, 1.0)
    color = material.find("color")
    if color is not None and color.get("rgba"):
        return parse_rgba(color.get("rgba", ""))
    name = material.get("name")
    if name and name in named_materials:
        return named_materials[name]
    return (0.65, 0.68, 0.72, 1.0)


def named_materials(
    root: ET.Element,
) -> dict[str, tuple[float, float, float, float]]:
    result: dict[str, tuple[float, float, float, float]] = {}
    for material in root.findall("material"):
        name = material.get("name")
        color = material.find("color")
        if name and color is not None and color.get("rgba"):
            result[name] = parse_rgba(color.get("rgba", ""))
    return result


def parse_color(value: str) -> int:
    normalized = value.removeprefix("#").removeprefix("0x")
    if len(normalized) != 6:
        raise argparse.ArgumentTypeError("color must use RRGGBB format")
    try:
        result = int(normalized, 16)
    except ValueError as error:
        raise argparse.ArgumentTypeError("color must use RRGGBB format") from error
    return result


def opacity(value: str) -> float:
    result = float(value)
    if not 0.0 <= result <= 1.0:
        raise argparse.ArgumentTypeError("opacity must be between 0 and 1")
    return result


def cylinder_lines(value: str) -> int:
    result = int(value)
    if result < 3:
        raise argparse.ArgumentTypeError("cylinder line count must be at least 3")
    return result


def safe_name(value: str) -> str:
    return value.replace("/", "_")


def render_urdf(
    viewer: meshcat.Visualizer,
    root: ET.Element,
    urdf_path: Path,
    link_transforms: dict[str, np.ndarray],
    *,
    collision_only: bool,
    visual_opacity: float,
    collision_opacity: float,
    collision_color: int,
    wireframe: bool,
    collision_cylinder_lines: int,
) -> tuple[int, int]:
    viewer.delete()
    materials = named_materials(root)
    visual_count = 0
    collision_count = 0
    collision_material = geometry.MeshPhongMaterial(
        color=collision_color,
        transparent=collision_opacity < 1.0,
        opacity=collision_opacity,
        wireframe=wireframe,
    )

    for link in root.findall("link"):
        link_name = link.get("name", "")
        link_transform = link_transforms[link_name]
        if not collision_only:
            for index, visual in enumerate(link.findall("visual")):
                geometry_element = visual.find("geometry")
                if geometry_element is None:
                    raise ValueError(f"visual geometry missing on link {link_name!r}")
                shape, correction = geometry_object(geometry_element, urdf_path)
                rgba = visual_rgba(visual, materials)
                alpha = visual_opacity * rgba[3]
                material = geometry.MeshPhongMaterial(
                    color=rgb_integer(rgba[:3]),
                    transparent=alpha < 1.0,
                    opacity=alpha,
                )
                node = viewer[
                    f"robot/visual/{safe_name(link_name)}/visual_{index}"
                ]
                node.set_object(shape, material)
                node.set_transform(
                    link_transform
                    @ origin_transform(visual.find("origin"))
                    @ correction
                )
                visual_count += 1

        for index, collision in enumerate(link.findall("collision")):
            geometry_element = collision.find("geometry")
            if geometry_element is None:
                raise ValueError(f"collision geometry missing on link {link_name!r}")
            cylinder = geometry_element.find("cylinder")
            if wireframe and cylinder is not None:
                radius, length = cylinder_dimensions(cylinder)
                shape = cylinder_wireframe(
                    radius,
                    length,
                    collision_cylinder_lines,
                    collision_color,
                    collision_opacity,
                )
                correction = cylinder_correction()
                custom_line_object = True
            else:
                shape, correction = geometry_object(geometry_element, urdf_path)
                custom_line_object = False
            node = viewer[
                f"robot/collision/{safe_name(link_name)}/collision_{index}"
            ]
            if custom_line_object:
                node.set_object(shape)
            else:
                node.set_object(shape, collision_material)
            node.set_transform(
                link_transform
                @ origin_transform(collision.find("origin"))
                @ correction
            )
            collision_count += 1
    return visual_count, collision_count


def close_viewer(viewer: meshcat.Visualizer) -> None:
    """Close MeshCat 0.3.x without relying on its broken Visualizer.close()."""
    window = viewer.window
    window.zmq_socket.close(linger=0)
    server_process = window.server_proc
    if server_process is None or server_process.poll() is not None:
        return
    server_process.terminate()
    try:
        server_process.wait(timeout=3.0)
    except subprocess.TimeoutExpired:
        server_process.kill()
        server_process.wait(timeout=3.0)


def export_static_html(viewer: meshcat.Visualizer, output: Path) -> None:
    output = output.resolve()
    if output.suffix.lower() != ".html":
        raise ValueError(f"MeshCat snapshot must use an .html extension: {output}")
    output.parent.mkdir(parents=True, exist_ok=True)
    temporary = output.with_name(f".{output.stem}.tmp{output.suffix}")
    temporary.write_text(viewer.static_html(), encoding="utf-8")
    os.replace(temporary, output)


def parse_args(argv: list[str]) -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--input", required=True, type=Path, help="URDF to display")
    parser.add_argument(
        "--joint",
        action="append",
        default=[],
        metavar="NAME=VALUE",
        help="joint position in radians, or meters for prismatic joints",
    )
    parser.add_argument(
        "--collision-only",
        action="store_true",
        help="hide visual geometry and display only collisions",
    )
    parser.add_argument(
        "--visual-opacity",
        type=opacity,
        default=1.0,
        help="visual geometry opacity (default: 1.0)",
    )
    parser.add_argument(
        "--collision-opacity",
        type=opacity,
        default=1.0,
        help="collision geometry opacity (default: 1.0)",
    )
    parser.add_argument(
        "--collision-color",
        type=parse_color,
        default=parse_color("00ff00"),
        metavar="RRGGBB",
        help="collision color in hexadecimal (default: 00ff00)",
    )
    parser.add_argument(
        "--collision-cylinder-lines",
        "--collision-cylinder-segments",
        dest="collision_cylinder_lines",
        type=cylinder_lines,
        default=12,
        metavar="COUNT",
        help="axial lines on collision cylinders (default: 12, every 30 degrees)",
    )
    collision_style = parser.add_mutually_exclusive_group()
    collision_style.add_argument(
        "--wireframe",
        dest="wireframe",
        action="store_true",
        default=True,
        help="draw collision geometry as wireframe (default)",
    )
    collision_style.add_argument(
        "--solid-collisions",
        dest="wireframe",
        action="store_false",
        help="draw collision geometry as translucent solids",
    )
    parser.add_argument(
        "--no-browser", action="store_true", help="do not automatically open a browser"
    )
    parser.add_argument(
        "--export-html",
        type=Path,
        help="write a standalone MeshCat HTML snapshot and exit",
    )
    parser.add_argument(
        "--exit-after-load",
        action="store_true",
        help=argparse.SUPPRESS,
    )
    parser.add_argument(
        "--zmq-url", help="connect to an existing MeshCat ZMQ server"
    )
    return parser.parse_args(argv)


def run(args: argparse.Namespace) -> int:
    urdf_path = args.input.resolve()
    if not urdf_path.is_file():
        raise FileNotFoundError(urdf_path)
    if urdf_path.suffix.lower() != ".urdf":
        raise ValueError(f"input must be a .urdf file: {urdf_path}")
    root = ET.parse(urdf_path).getroot()
    if root.tag != "robot":
        raise ValueError(f"URDF root must be <robot>, found <{root.tag}>")

    joints = parse_joints(root)
    requested = parse_joint_values(args.joint)
    values = resolve_joint_values(joints, requested)
    link_transforms = compute_link_transforms(root, joints, values)

    viewer = meshcat.Visualizer(zmq_url=args.zmq_url)
    visual_count, collision_count = render_urdf(
        viewer,
        root,
        urdf_path,
        link_transforms,
        collision_only=args.collision_only,
        visual_opacity=args.visual_opacity,
        collision_opacity=args.collision_opacity,
        collision_color=args.collision_color,
        wireframe=args.wireframe,
        collision_cylinder_lines=args.collision_cylinder_lines,
    )
    if collision_count == 0:
        raise RuntimeError(f"URDF contains no <collision> elements: {urdf_path}")

    url = viewer.url()
    print(f"Loaded: {urdf_path}")
    print(f"Visual geometry: {visual_count}")
    print(f"Collision geometry: {collision_count}")
    print(f"MeshCat URL: {url}", flush=True)
    if args.export_html is not None:
        export_static_html(viewer, args.export_html)
        snapshot_uri = args.export_html.resolve().as_uri()
        print(f"Standalone snapshot: {snapshot_uri}", flush=True)
        if not args.no_browser:
            webbrowser.open(snapshot_uri, new=2)
        close_viewer(viewer)
        return 0
    if not args.no_browser:
        viewer.open()
    if args.exit_after_load:
        close_viewer(viewer)
        return 0

    print("Press Ctrl+C to stop the viewer.", flush=True)
    try:
        while True:
            time.sleep(1.0)
    except KeyboardInterrupt:
        print("Stopping MeshCat viewer.")
    finally:
        close_viewer(viewer)
    return 0


def main(argv: list[str] | None = None) -> int:
    try:
        return run(parse_args(sys.argv[1:] if argv is None else argv))
    except Exception:
        traceback.print_exc()
        return 1


if __name__ == "__main__":
    raise SystemExit(main())