#!/usr/bin/env python3
"""
Generate prototype STL files for a clamp-on dual monitor arm/cradle set.

Design brief from Sott:
- Existing pole diameter: 35 mm.
- Clamp to existing pole using four M4 nuts/bolts.
- Lateral pivots left/right from clamp using M4 nuts/bolts.
- Two first-stage arms: 180 mm long x 20 mm high x 20 mm breadth.
- Two second-stage arms same size.
- Final monitor support is a non-VESA open cradle: 180 mm long x 25 mm high x 15 mm internal width.
- Print in PLA.

Assumptions used for this prototype:
- M4 clearance holes = 4.5 mm diameter.
- Pole clearance = 35.6 mm diameter.
- Arm pivot holes run through the 20 mm height (Z axis).
- Clamp is a two-half split collar with four transverse M4 clamp bolt holes running across the clamp halves.
- Cradle has 15 mm internal slot width, 4 mm wall thickness, open along both ends and top, with one M4 pivot boss.
"""
from __future__ import annotations
from pathlib import Path
import math, struct, zipfile, json
from typing import Iterable, List, Sequence, Tuple

from shapely.geometry import Point, Polygon, box
from shapely.ops import unary_union, triangulate

OUT = Path(__file__).resolve().parent

Vec = Tuple[float, float, float]
Tri = Tuple[Vec, Vec, Vec]

# ---------------- STL utilities ----------------

def sub(a: Vec, b: Vec) -> Vec:
    return (a[0]-b[0], a[1]-b[1], a[2]-b[2])

def cross(a: Vec, b: Vec) -> Vec:
    return (a[1]*b[2]-a[2]*b[1], a[2]*b[0]-a[0]*b[2], a[0]*b[1]-a[1]*b[0])

def norm(v: Vec) -> Vec:
    l = math.sqrt(v[0]*v[0]+v[1]*v[1]+v[2]*v[2])
    if l == 0:
        return (0.0, 0.0, 0.0)
    return (v[0]/l, v[1]/l, v[2]/l)

def tri_normal(t: Tri) -> Vec:
    return norm(cross(sub(t[1], t[0]), sub(t[2], t[0])))

def write_binary_stl(path: Path, triangles: Sequence[Tri], name: str = "mesh") -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open('wb') as f:
        header = (name[:76] + ' ' * 80).encode('ascii', errors='ignore')[:80]
        f.write(header)
        f.write(struct.pack('<I', len(triangles)))
        for tri in triangles:
            n = tri_normal(tri)
            f.write(struct.pack('<3f', *n))
            for v in tri:
                f.write(struct.pack('<3f', *v))
            f.write(struct.pack('<H', 0))

def area_xy(tri: Tri) -> float:
    (x1,y1,_),(x2,y2,_),(x3,y3,_) = tri
    return 0.5*((x2-x1)*(y3-y1)-(x3-x1)*(y2-y1))

def orient(tri: Tri, want_positive_xy: bool) -> Tri:
    pos = area_xy(tri) > 0
    if pos != want_positive_xy:
        return (tri[0], tri[2], tri[1])
    return tri

# ---------------- 2D extrusion utilities ----------------

def ring_coords(poly: Polygon) -> List[Tuple[float, float]]:
    coords = list(poly.exterior.coords)
    if coords[0] == coords[-1]:
        coords = coords[:-1]
    return [(float(x), float(y)) for x, y in coords]

def add_wall(tris: List[Tri], coords: Sequence[Tuple[float,float]], z0: float, z1: float, outward_ccw: bool=True) -> None:
    n = len(coords)
    for i in range(n):
        x1,y1 = coords[i]
        x2,y2 = coords[(i+1)%n]
        if outward_ccw:
            tris.append(((x1,y1,z0),(x2,y2,z0),(x2,y2,z1)))
            tris.append(((x1,y1,z0),(x2,y2,z1),(x1,y1,z1)))
        else:
            tris.append(((x1,y1,z0),(x2,y2,z1),(x2,y2,z0)))
            tris.append(((x1,y1,z0),(x1,y1,z1),(x2,y2,z1)))

def extrude_polygon(geom, z0: float, z1: float) -> List[Tri]:
    """Extrude a Shapely Polygon/MultiPolygon with holes into a manifold-ish STL mesh."""
    tris: List[Tri] = []
    polygons = list(geom.geoms) if hasattr(geom, 'geoms') else [geom]
    for poly in polygons:
        if poly.is_empty:
            continue
        # Top/bottom triangulation, filtered to polygon interior including holes.
        for t in triangulate(poly):
            rp = t.representative_point()
            if not poly.contains(rp):
                continue
            coords = list(t.exterior.coords)[:3]
            top = tuple((float(x), float(y), z1) for x,y in coords)  # type: ignore
            bot = tuple((float(x), float(y), z0) for x,y in coords)  # type: ignore
            tris.append(orient(top, True))
            tris.append(orient(bot, False))
        # Exterior wall: shapely exterior usually CCW.
        add_wall(tris, ring_coords(poly), z0, z1, outward_ccw=True)
        # Hole walls: interiors usually CW, reverse orientation handling.
        for interior in poly.interiors:
            coords = list(interior.coords)
            if coords[0] == coords[-1]:
                coords = coords[:-1]
            add_wall(tris, [(float(x),float(y)) for x,y in coords], z0, z1, outward_ccw=False)
    return tris

# ---------------- Part generators ----------------

M4_CLEAR = 4.5
POLE_CLEAR = 35.6


def capsule(length: float, width: float, resolution: int=32) -> Polygon:
    r = width / 2
    rect = box(-length/2 + r, -r, length/2 - r, r)
    return unary_union([
        rect,
        Point(-length/2 + r, 0).buffer(r, resolution=resolution),
        Point(length/2 - r, 0).buffer(r, resolution=resolution),
    ])

PIVOT_TAB_THICKNESS = 8.0


def arm_mesh(length=180.0, width=20.0, height=20.0, hole_d=M4_CLEAR, pivot_thickness=PIVOT_TAB_THICKNESS) -> List[Tri]:
    """Arm with 20 mm high centre body and 8 mm pivot tabs to match cradle pivot thickness."""
    r_outer = width / 2
    r_hole = hole_d / 2
    left_c = -length/2 + r_outer
    right_c = length/2 - r_outer
    body = box(left_c, -width/2, right_c, width/2)
    left_tab = Point(left_c, 0).buffer(r_outer, resolution=32).difference(Point(left_c, 0).buffer(r_hole, resolution=24))
    right_tab = Point(right_c, 0).buffer(r_outer, resolution=32).difference(Point(right_c, 0).buffer(r_hole, resolution=24))
    tris: List[Tri] = []
    # 20 mm high centre beam, preserving original arm strength/height.
    tris += extrude_polygon(body, 0, height)
    # Pivot ends/tabs reduced to the same 8 mm thickness as the cradle pivot tab.
    tris += extrude_polygon(left_tab, 0, pivot_thickness)
    tris += extrude_polygon(right_tab, 0, pivot_thickness)
    return tris

def pivot_spacer_mesh(outer_d=12.0, inner_d=M4_CLEAR, height=2.0) -> List[Tri]:
    """Thin optional M4 pivot washer/spacer, sized for the 8 mm pivot tabs."""
    shape = Point(0,0).buffer(outer_d/2, resolution=48).difference(Point(0,0).buffer(inner_d/2, resolution=24))
    return extrude_polygon(shape, 0, height)

def cradle_mesh(length=180.0, internal_width=15.0, wall=4.0, height=25.0, pivot_tab_len=25.0, hole_d=M4_CLEAR) -> List[Tri]:
    """Open U-channel along X, open top. Pivot tab is a flat boss at x=-length/2."""
    tris: List[Tri] = []
    outer_w = internal_width + 2*wall
    # Bottom plate: x length, y outer width, z wall.
    bottom = box(-length/2, -outer_w/2, length/2, outer_w/2)
    tris += extrude_polygon(bottom, 0, wall)
    # Side walls.
    left_wall = box(-length/2, -outer_w/2, length/2, -internal_width/2)
    right_wall = box(-length/2, internal_width/2, length/2, outer_w/2)
    tris += extrude_polygon(left_wall, wall, height)
    tris += extrude_polygon(right_wall, wall, height)
    # Pivot tab/boss: flat horizontal plate with M4 vertical hole, attached to left end of cradle.
    tab = capsule(pivot_tab_len + 20, 20, 24)
    # Move capsule so it projects from left end; centre at -length/2 - pivot_tab_len/2.
    tab = shapely_translate(tab, xoff=-length/2 - pivot_tab_len/2 + 10, yoff=0)
    tab = tab.difference(Point(-length/2 - pivot_tab_len + 10, 0).buffer(hole_d/2, resolution=24))
    tris += extrude_polygon(tab, 0, 8.0)
    return tris

def shapely_translate(geom, xoff=0.0, yoff=0.0):
    from shapely import affinity
    return affinity.translate(geom, xoff=xoff, yoff=yoff)

def clamp_shapes(side='left', inner_d=POLE_CLEAR, outer_d=68.0, split_gap=2.0, hole_d=M4_CLEAR):
    """Return collar footprint and separate lateral pivot boss footprint."""
    ro = outer_d/2
    ri = inner_d/2
    ring = Point(0,0).buffer(ro, resolution=64).difference(Point(0,0).buffer(ri, resolution=64))
    if side == 'left':
        collar = ring.intersection(box(-ro-1, -ro-1, -split_gap/2, ro+1))
        boss = capsule(38, 22, 24)
        boss = shapely_translate(boss, xoff=-ro-18, yoff=0)
        boss = boss.difference(Point(-ro-28, 0).buffer(hole_d/2, resolution=24))
    else:
        collar = ring.intersection(box(split_gap/2, -ro-1, ro+1, ro+1))
        boss = capsule(38, 22, 24)
        boss = shapely_translate(boss, xoff=ro+18, yoff=0)
        boss = boss.difference(Point(ro+28, 0).buffer(hole_d/2, resolution=24))
    return collar, boss


def clamp_half_mesh(side='left', height=50.0, inner_d=POLE_CLEAR, outer_d=68.0, split_gap=2.0, hole_d=M4_CLEAR) -> List[Tri]:
    """Two split-collar halves around a vertical 35 mm pole.

    Amendment requested by Sott:
    - The four M4 clamp bolts are transverse, not longitudinal.
    - In this coordinate system the pole is vertical on Z, and the clamp bolts run horizontally
      across the left/right clamp halves on the X axis.
    - Each half has four actual transverse M4 clearance channels: two front/back positions and
      two heights. The channels are cut into the relevant Z bands so slicers see real openings.
    - The lateral arm pivot remains a vertical M4 pivot hole through the left/right boss.
    """
    ro = outer_d/2
    collar_shape, boss_shape = clamp_shapes(side, inner_d, outer_d, split_gap, hole_d)
    y_positions = (-22.0, 22.0)
    z_centres = (15.0, 35.0)
    channel_h = hole_d + 0.4          # vertical clearance band around M4 shaft
    channel_w_y = hole_d + 1.0        # front/back clearance
    # Through-channel along X for each half: from outside edge to split face.
    if side == 'left':
        x_cut = (-ro - 1.0, 0.5)
    else:
        x_cut = (-0.5, ro + 1.0)

    z_breaks = [0.0]
    for zc in z_centres:
        z_breaks += [max(0.0, zc - channel_h/2), min(height, zc + channel_h/2)]
    z_breaks += [height]
    z_breaks = sorted(set(round(z, 4) for z in z_breaks))

    tris: List[Tri] = []
    for z0, z1 in zip(z_breaks[:-1], z_breaks[1:]):
        if z1 <= z0:
            continue
        mid = (z0 + z1) / 2
        shape = collar_shape
        in_channel_band = any(abs(mid - zc) <= channel_h/2 + 1e-6 for zc in z_centres)
        if in_channel_band:
            cuts = []
            for y in y_positions:
                cuts.append(box(x_cut[0], y - channel_w_y/2, x_cut[1], y + channel_w_y/2))
            shape = shape.difference(unary_union(cuts))
        tris += extrude_polygon(shape, z0, z1)
    # Lateral clamp pivot boss reduced to 8 mm thickness, matching the cradle pivot tab.
    tris += extrude_polygon(boss_shape, 0, PIVOT_TAB_THICKNESS)
    return tris


def write_readme() -> None:
    readme = OUT/'README_DUAL_MONITOR_ARM.txt'
    readme.write_text(f"""Dual monitor arm/cradle prototype STL set for PLA

Generated parts:
- pole_clamp_left_half_35mm_pole_M4.stl
- pole_clamp_right_half_35mm_pole_M4.stl
- arm_180x20x20_M4_pivots_PRINT_4.stl  (print 4 copies)
- open_monitor_cradle_180x25x15_M4_pivot_PRINT_2.stl  (print 2 copies; mirror in slicer if required)
- M4_pivot_washer_12x2_PRINT_AS_NEEDED.stl

Key dimensions:
- Pole clearance: {POLE_CLEAR:.1f} mm for a 35 mm pole.
- Clamp bolts: four transverse M4 channels, not longitudinal/vertical holes.
- M4 clearance holes/channels: {M4_CLEAR:.1f} mm nominal clearance.
- Link arms: 180 mm long x 20 mm wide x 20 mm high centre body.
- Arm pivot tabs, clamp pivot bosses and cradle pivot tabs: {PIVOT_TAB_THICKNESS:.1f} mm matched thickness.
- Cradle: 180 mm long, 25 mm high, 15 mm internal monitor slot, 4 mm walls.

PLA print notes:
- This is a prototype, not a certified load-bearing monitor mount.
- Use high infill/perimeters for load-bearing parts, e.g. 5+ walls and 50-80% infill.
- Print arms flat on the bed so layer lines run through the length/width, not across the bolt load path.
- Use washers under M4 bolt heads/nuts.
- Test with the monitor supported by hand first; do not leave valuable equipment unsupported until physically verified.
- If monitors are heavy, use PETG/ABS/nylon or metal reinforcement rather than PLA.
""", encoding='utf-8')


def main():
    files = []
    parts = {
        'pole_clamp_left_half_35mm_pole_M4.stl': clamp_half_mesh('left'),
        'pole_clamp_right_half_35mm_pole_M4.stl': clamp_half_mesh('right'),
        'arm_180x20x20_M4_pivots_PRINT_4.stl': arm_mesh(),
        'open_monitor_cradle_180x25x15_M4_pivot_PRINT_2.stl': cradle_mesh(),
        'M4_pivot_washer_12x2_PRINT_AS_NEEDED.stl': pivot_spacer_mesh(),
    }
    manifest = []
    for name, tris in parts.items():
        path = OUT/name
        write_binary_stl(path, tris, name)
        files.append(path)
        manifest.append({'file': name, 'triangles': len(tris), 'bytes': path.stat().st_size})
    write_readme()
    files.append(OUT/'README_DUAL_MONITOR_ARM.txt')
    files.append(Path(__file__).resolve())
    (OUT/'manifest.json').write_text(json.dumps(manifest, indent=2), encoding='utf-8')
    files.append(OUT/'manifest.json')
    zip_path = OUT/'dual_monitor_arm_cradle_stl_pack.zip'
    with zipfile.ZipFile(zip_path, 'w', zipfile.ZIP_DEFLATED) as z:
        for p in files:
            z.write(p, arcname=p.name)
    print('Generated:')
    for item in manifest:
        print(f"- {item['file']}: {item['triangles']} triangles, {item['bytes']} bytes")
    print(f'- {zip_path.name}: {zip_path.stat().st_size} bytes')

if __name__ == '__main__':
    main()