#!/usr/bin/env python3
"""Generate rigid-link STL meshes for the Solenopsis (fire-ant) URDF.

The SolidWorks STEP export uses millimetres and this script intentionally keeps
the STL coordinates in millimetres.  The URDF applies a 0.001 mesh scale.

Rigid grouping used here:
  * base: root-level fixed parts plus each P65 fixed part and motor plate;
  * knee: arm, P65 rotating part, W190 harness and module fasteners;
  * wheel: W190 motor/wheel appearance assembly.

The grouping is explicit so it can be audited when the CAD assembly changes.
"""

from __future__ import annotations

import argparse
import json
import math
from pathlib import Path

from OCP.BRep import BRep_Builder
from OCP.BRepBuilderAPI import BRepBuilderAPI_Transform
from OCP.BRepMesh import BRepMesh_IncrementalMesh
from OCP.STEPCAFControl import STEPCAFControl_Reader
from OCP.StlAPI import StlAPI_Writer
from OCP.TCollection import TCollection_ExtendedString
from OCP.TDataStd import TDataStd_Name
from OCP.TDF import TDF_Label
from OCP.TDocStd import TDocStd_Document
from OCP.TopoDS import TopoDS_Compound
from OCP.XCAFApp import XCAFApp_Application
from OCP.XCAFDoc import XCAFDoc_DocumentTool, XCAFDoc_ShapeTool
from OCP.collections import Sequence_TDF_Label
from OCP.gp import gp_Trsf


# SolidWorks root axes -> ROS/URDF axes: (x, y, z)_urdf = (z, x, y)_sw.
SW_TO_URDF = (
    (0.0, 0.0, 1.0, 0.0),
    (1.0, 0.0, 0.0, 0.0),
    (0.0, 1.0, 0.0, 0.0),
    (0.0, 0.0, 0.0, 1.0),
)

IDENTITY = (
    (1.0, 0.0, 0.0, 0.0),
    (0.0, 1.0, 0.0, 0.0),
    (0.0, 0.0, 1.0, 0.0),
    (0.0, 0.0, 0.0, 1.0),
)

# Standardised URDF moving-link frames.  In both cases the arm points along
# link -Z at q=0 and the URDF joint axis is link -Y, matching the established
# controller convention.  The right-hand transform mirrors the CAD module in
# the physically correct way while remaining a proper rotation (determinant 1).
MODULE_TO_LEFT_LINK = (
    (0.0, 0.0, -1.0, 0.0),
    (-1.0, 0.0, 0.0, 0.0),
    (0.0, 1.0, 0.0, 0.0),
    (0.0, 0.0, 0.0, 1.0),
)
MODULE_TO_RIGHT_LINK = (
    (0.0, 0.0, 1.0, 0.0),
    (1.0, 0.0, 0.0, 0.0),
    (0.0, 1.0, 0.0, 0.0),
    (0.0, 0.0, 0.0, 1.0),
)

# The two named components are fixed to the chassis side of the P65 joint.
BASE_COMPONENT_MARKERS = (
    "Solenopsis-02-002 关节电机固定板",
    "P65MAX-S外观模型-A0-20260205-固定组件",
)
WHEEL_COMPONENT_MARKER = "W190-S 电机外观模型"


def matmul(a, b):
    return tuple(
        tuple(sum(a[i][k] * b[k][j] for k in range(4)) for j in range(4))
        for i in range(4)
    )


def transform_point(t, p):
    q = (*p, 1.0)
    return tuple(sum(t[i][k] * q[k] for k in range(4)) for i in range(3))


def translation(x, y, z):
    return (
        (1.0, 0.0, 0.0, x),
        (0.0, 1.0, 0.0, y),
        (0.0, 0.0, 1.0, z),
        (0.0, 0.0, 0.0, 1.0),
    )


def location_matrix(label):
    trsf = XCAFDoc_ShapeTool.GetLocation_s(label).Transformation()
    return tuple(
        tuple(trsf.Value(i, j) for j in range(1, 5))
        if i < 4
        else (0.0, 0.0, 0.0, 1.0)
        for i in range(1, 5)
    )


def gp_transform(matrix):
    t = gp_Trsf()
    t.SetValues(*(matrix[i][j] for i in range(3) for j in range(4)))
    return t


class CadDocument:
    def __init__(self, step_path: Path):
        self.document = TDocStd_Document(TCollection_ExtendedString("solenopsis"))
        application = XCAFApp_Application.GetApplication_s()
        application.NewDocument(TCollection_ExtendedString("MDTV-XCAF"), self.document)

        reader = STEPCAFControl_Reader()
        reader.SetNameMode(True)
        reader.SetColorMode(True)
        reader.SetPropsMode(True)
        status = reader.ReadFile(str(step_path))
        if "RetDone" not in str(status):
            raise RuntimeError(f"STEP read failed: {status}")
        if not reader.Transfer(self.document):
            raise RuntimeError("STEP transfer into XCAF failed")

        self.shape_tool = XCAFDoc_DocumentTool.ShapeTool_s(self.document.Main())
        free = Sequence_TDF_Label()
        self.shape_tool.GetFreeShapes(free)
        if free.Length() != 1:
            raise RuntimeError(f"Expected one root assembly, found {free.Length()}")
        self.root = free.Value(1)

    @staticmethod
    def label_name(label):
        attribute = TDataStd_Name()
        if label.FindAttribute(TDataStd_Name.GetID_s(), attribute):
            return attribute.Get().ToExtString()
        return ""

    @staticmethod
    def referred(label):
        target = TDF_Label()
        if XCAFDoc_ShapeTool.GetReferredShape_s(label, target):
            return target
        return label

    def referred_name(self, label):
        return self.label_name(self.referred(label))

    def components(self, label):
        target = self.referred(label)
        result = Sequence_TDF_Label()
        XCAFDoc_ShapeTool.GetComponents_s(target, result, False)
        return [result.Value(i) for i in range(1, result.Length() + 1)]

    def referred_shape(self, component):
        return XCAFDoc_ShapeTool.GetShape_s(self.referred(component))


class ShapeGroup:
    def __init__(self):
        self.builder = BRep_Builder()
        self.compound = TopoDS_Compound()
        self.builder.MakeCompound(self.compound)
        self.members = []

    def add(self, cad: CadDocument, component, output_from_parent, member_name=None):
        output_from_component = matmul(output_from_parent, location_matrix(component))
        shape = cad.referred_shape(component)
        moved = BRepBuilderAPI_Transform(
            shape, gp_transform(output_from_component), True, False
        ).Shape()
        self.builder.Add(self.compound, moved)
        self.members.append(member_name or cad.referred_name(component))

    def write_stl(self, output_path: Path, linear_deflection_mm: float):
        output_path.parent.mkdir(parents=True, exist_ok=True)
        BRepMesh_IncrementalMesh(
            self.compound, linear_deflection_mm, False, 0.35, True
        ).Perform()
        writer = StlAPI_Writer()
        # Binary STL is lossless relative to the tessellation and is typically
        # five to eight times smaller than the default ASCII representation.
        writer.ASCIIMode = False
        if not writer.Write(self.compound, str(output_path)):
            raise RuntimeError(f"Failed to write {output_path}")


def module_name_from_hip(hip_sw):
    # URDF forward = SW +Z, URDF left = SW +X.
    fore_aft = "f" if hip_sw[2] > 0.0 else "h"
    left_right = "l" if hip_sw[0] > 0.0 else "r"
    return fore_aft + left_right


def rotation_rpy(matrix):
    """Return fixed-axis URDF roll/pitch/yaw for a proper rotation matrix."""
    pitch = math.asin(max(-1.0, min(1.0, -matrix[2][0])))
    if abs(math.cos(pitch)) > 1.0e-10:
        roll = math.atan2(matrix[2][1], matrix[2][2])
        yaw = math.atan2(matrix[1][0], matrix[0][0])
    else:
        roll = math.atan2(-matrix[0][1], matrix[1][1])
        yaw = 0.0
    return roll, pitch, yaw


def generate(step_path: Path, output_dir: Path, deflection_mm: float):
    cad = CadDocument(step_path)
    base = ShapeGroup()
    knees = {}
    wheels = {}
    manifest = {
        "source_step": str(step_path),
        "stl_units": "millimetre",
        "linear_deflection_mm": deflection_mm,
        "axis_map": "URDF (x,y,z) = SolidWorks (z,x,y)",
        "modules": {},
    }

    module_components = []
    for component in cad.components(cad.root):
        name = cad.referred_name(component)
        if name == "Solenopsis-02 运动模组总成":
            module_components.append(component)
        else:
            base.add(cad, component, SW_TO_URDF, name)

    if len(module_components) != 4:
        raise RuntimeError(f"Expected four motion modules, found {len(module_components)}")

    for module in module_components:
        root_from_module = location_matrix(module)
        hip_sw = transform_point(root_from_module, (0.0, 300.0, 0.0))
        leg = module_name_from_hip(hip_sw)
        if leg in knees:
            raise RuntimeError(f"Duplicate leg classification: {leg}")

        knees[leg] = ShapeGroup()
        wheels[leg] = ShapeGroup()

        # The STEP assembly stores the four legs at four arbitrary joint
        # angles.  A URDF needs a common neutral pose, so q=0 is defined with
        # every 300 mm arm vertical and all established joint axes at -Y.
        module_to_link = (
            MODULE_TO_LEFT_LINK if leg.endswith("l") else MODULE_TO_RIGHT_LINK
        )
        knee_from_module = matmul(
            module_to_link, translation(0.0, -300.0, 0.0)
        )
        wheel_from_module = module_to_link

        for child in cad.components(module):
            name = cad.referred_name(child)
            if any(marker in name for marker in BASE_COMPONENT_MARKERS):
                base.add(cad, child, matmul(SW_TO_URDF, root_from_module), f"{leg}:{name}")
            elif WHEEL_COMPONENT_MARKER in name:
                wheels[leg].add(cad, child, wheel_from_module, name)
            else:
                knees[leg].add(cad, child, knee_from_module, name)

        urdf_from_module = matmul(SW_TO_URDF, root_from_module)
        hip_urdf_mm = transform_point(SW_TO_URDF, hip_sw)
        manifest["modules"][leg] = {
            "hip_urdf_m": [value / 1000.0 for value in hip_urdf_mm],
            "joint_rpy_rad": [0.0, 0.0, 0.0],
            "joint_axis": [0.0, -1.0, 0.0],
            "cad_assembly_rpy_rad": list(rotation_rpy(urdf_from_module)),
            "base_members": [
                name for name in base.members if name.startswith(f"{leg}:")
            ],
            "knee_members": knees[leg].members,
            "wheel_members": wheels[leg].members,
        }

    base.write_stl(output_dir / "base_link.STL", deflection_mm)
    for leg in ("fl", "fr", "hl", "hr"):
        knees[leg].write_stl(output_dir / f"{leg}_knee.STL", deflection_mm)
        wheels[leg].write_stl(output_dir / f"{leg}_wheel.STL", deflection_mm)

    (output_dir / "mesh_manifest.json").write_text(
        json.dumps(manifest, ensure_ascii=False, indent=2) + "\n", encoding="utf-8"
    )


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("step", type=Path)
    parser.add_argument("output_dir", type=Path)
    parser.add_argument(
        "--deflection-mm",
        type=float,
        default=0.2,
        help="Maximum linear tessellation deflection in millimetres (default: 0.2)",
    )
    args = parser.parse_args()
    generate(args.step.resolve(), args.output_dir.resolve(), args.deflection_mm)


if __name__ == "__main__":
    main()