"""Detachable fiducial mount v2 (tool-free, drop-on, FLAT bottom, joint at the mount).

Joint sits right at the base-mount / connector interface, so:
  Part A (stays on arm) = holder + base-mount + a short SHELF stub (at the connector plane, ON the table).
  Part B (detachable)   = the whole connector + fiducial board, whose END LAPS OVER Part A's shelf.
The lap and pins are ABOVE the connector plane (bump on top, near the mount), so EVERYTHING's bottom stays
at the table plane (nothing pokes below). Part B drops straight down: its tray lands on the table and its
lap-end drops onto Part A's upward pins. Lift to detach. Servo frame, mm.
"""
import os
import numpy as np
import trimesh
import gen_two_servo as g

GLB = "two_servo_from_blender.glb"
PIN_R = 1.5
CLEAR = 0.35
PIN_PROUD = 3.5
PIN_ROOT = 2.0
HOLE_DEPTH = 3.8
SHELF_LEN = 24.0     # lap length from the mount along the connector
HOLDER = "holder_fillet.stl"


def _box(x0, x1, y0, y1, z0, z1):
    b = trimesh.creation.box(extents=[x1 - x0, y1 - y0, z1 - z0])
    b.apply_translation([(x0 + x1) / 2, (y0 + y1) / 2, (z0 + z1) / 2])
    return b


def _boxb(b):
    return _box(b[0][0], b[1][0], b[0][1], b[1][1], b[0][2], b[1][2])


def _cylY(x, z, r, y0, y1):
    c = trimesh.creation.cylinder(radius=r, height=y1 - y0, sections=40)
    c.apply_transform(trimesh.transformations.rotation_matrix(np.pi / 2, [1, 0, 0]))
    c.apply_translation([x, (y0 + y1) / 2, z])
    return c


def _U(parts):
    return trimesh.boolean.union(parts, engine="manifold")


# --- classify base pieces ---
scene = trimesh.load(GLB)
base_inv = np.linalg.inv(np.array(scene.graph["servo_0"][0], float))
P = {}
for node in sorted(scene.graph.nodes_geometry):
    if node.startswith("base_cube"):
        T, gn = scene.graph[node]
        m = scene.geometry[gn].copy()
        m.apply_transform(base_inv @ np.array(T, float))
        m.apply_scale(1000.0 / g.GLB_SCALE)
        e = sorted(m.extents, reverse=True)
        kind = "board" if e[1] > 100 else ("connector" if e[0] > 100 else "mount")
        P[kind] = m.bounds
assert set(P) == {"board", "connector", "mount"}, f"classify failed: {set(P)}"

(cx0, cy0, cz0), (cx1, cy1, cz1) = P["connector"]          # connector: Y[cy0,cy1] (thin), Z long
wy = cy1 - cy0
mount_zc = P["mount"].mean(0)[2]
zbase, zfid = (cz1, cz0) if abs(cz1 - mount_zc) < abs(cz0 - mount_zc) else (cz0, cz1)
sgn = float(np.sign(zfid - zbase))
z_shelf = zbase + sgn * SHELF_LEN                          # shelf end (away from mount)
z_step = z_shelf + sgn * 4.0                               # small step block just past the shelf
lap_z0, lap_z1 = sorted([z_shelf, zbase - sgn * 2.0])      # lap over the shelf, stop 2mm short of the mount face
pin_z = [zbase + sgn * SHELF_LEN * f for f in (0.30, 0.72)]
pin_x = [cx0 + (cx1 - cx0) * f for f in (0.25, 0.75)]
pins = [(x, z) for z in pin_z for x in pin_x]

# --- Part A: holder + mount + shelf (on the table) + upward pins ---
shelfA = _box(cx0, cx1, cy0, cy1, min(zbase, z_shelf), max(zbase, z_shelf))
partA = _U([trimesh.load(HOLDER, force="mesh"), _boxb(P["mount"]), shelfA])
for (x, z) in pins:
    body = _cylY(x, z, PIN_R, cy1 - PIN_ROOT, cy1 + PIN_PROUD - PIN_R)     # pins rise from the shelf TOP (cy1)
    dome = trimesh.creation.icosphere(radius=PIN_R, subdivisions=3)
    dome.apply_translation([x, cy1 + PIN_PROUD - PIN_R, z])
    partA = _U([partA, body, dome])
partA.export("base_holder_detach.stl")

# --- Part B: board + connector (on the table) + step + lap-over-end (above), minus holes ---
connB = _box(cx0, cx1, cy0, cy1, min(zfid, z_step), max(zfid, z_step))
step = _box(cx0, cx1, cy0, cy1 + wy, min(z_step, z_shelf), max(z_step, z_shelf))   # welds tray to the raised lap
lapB = _box(cx0, cx1, cy1, cy1 + wy, lap_z0, lap_z1)                               # laps OVER Part A's shelf
partB = _U([_boxb(P["board"]), connB, step, lapB])
for (x, z) in pins:
    partB = trimesh.boolean.difference([partB, _cylY(x, z, PIN_R + CLEAR, cy1 - 0.1, cy1 + HOLE_DEPTH)], engine="manifold")
partB.export("fiducial_mount.stl")

print(f"joint at mount; shelf Z[{min(zbase,z_shelf):.0f},{max(zbase,z_shelf):.0f}] lap Y[{cy1:.0f},{cy1+wy:.0f}] (above plane)")
print(f"lowest Y: partA={partA.bounds[0][1]:.1f} partB={partB.bounds[0][1]:.1f}  (connector plane cy0={cy0:.0f})")
print(f"Part A: wt={partA.is_watertight} bodies={len(partA.split(only_watertight=False))} vol={partA.volume/1000:.1f}cm3")
print(f"Part B: wt={partB.is_watertight} bodies={len(partB.split(only_watertight=False))} vol={partB.volume/1000:.1f}cm3")

# --- URDF: beauty arm; servo_0 base = Part A (blue) + Part B (orange), mated ---
import shutil
import xml.etree.ElementTree as ET
trimesh.util.concatenate([partA, partB]).export("base_combined.stl")
g.LAYOUT = g.read_layout(GLB)
g.PARTS[1] = ("holder_fillet.stl",) + tuple(g.PARTS[1][1:])
g.PARTS[2] = ("yoke_fillet.stl",) + tuple(g.PARTS[2][1:])
g.CONNECTOR_FMT = "connector_beauty_{}.stl"
shutil.copy("base_holder.stl", "_bh_bak.stl")
shutil.copy("base_combined.stl", "base_holder.stl")
try:
    g.gen_urdf("two_servo_detach.urdf")
finally:
    shutil.copy("_bh_bak.stl", "base_holder.stl")
    os.remove("_bh_bak.stl")
t = ET.parse("two_servo_detach.urdf")
for link in t.getroot().findall("link"):
    if link.get("name") == "servo_0":
        for v in link.findall("visual"):
            mesh = v.find("geometry/mesh")
            if mesh is not None and mesh.get("filename") == "base_holder.stl":
                mesh.set("filename", "base_holder_detach.stl")
        v2 = ET.SubElement(link, "visual")
        ET.SubElement(v2, "origin", {"xyz": "0 0 0", "rpy": "0 0 0"})
        ET.SubElement(ET.SubElement(v2, "geometry"), "mesh", {"filename": "fiducial_mount.stl", "scale": "0.001 0.001 0.001"})
        ET.SubElement(ET.SubElement(v2, "material", {"name": "fidmount"}), "color", {"rgba": "0.93 0.55 0.16 1"})
        break
ET.indent(t, "  ")
t.write("two_servo_detach.urdf", encoding="unicode", xml_declaration=True)
print("wrote two_servo_detach.urdf")