"""Shared helpers for the six image-derived shaft studies.

All source dimensions are millimetres.  Parts are authored on +Z and rotated
onto +X at export so their long axis matches the supplied reference sheets.
"""

from build123d import Align, Axis, Box, Cylinder, Pos


JOIN_OVERLAP = 0.15


def stepped_body(segments):
    """Fuse [(diameter, length), ...] into one robust coaxial solid."""
    body = None
    z = 0.0
    for index, (diameter, length) in enumerate(segments):
        overlap = JOIN_OVERLAP if index else 0.0
        piece = Pos(0, 0, z - overlap) * Cylinder(
            diameter / 2.0,
            length + overlap,
            align=(Align.CENTER, Align.CENTER, Align.MIN),
        )
        body = piece if body is None else body + piece
        z += length
    return body, z


def cut_end_bore(body, total_length, diameter, depth):
    tool = Pos(0, 0, total_length - depth) * Cylinder(
        diameter / 2.0,
        depth + 1.0,
        align=(Align.CENTER, Align.CENTER, Align.MIN),
    )
    return body - tool


def cut_side_flat(body, diameter, z_start, length, flat_depth):
    """Cut one axial wrench/key flat into the +X side of a round shaft."""
    x_start = diameter / 2.0 - flat_depth
    tool = Pos(x_start, 0, z_start) * Box(
        diameter,
        diameter * 1.5,
        length,
        align=(Align.MIN, Align.CENTER, Align.MIN),
    )
    return body - tool


def add_thread_rings(body, z_start, length, major_diameter, pitch=2.0, rib_width=0.65):
    """Add simplified annular thread indication (not a true helical thread)."""
    count = max(1, int(length // pitch))
    for index in range(count):
        z = z_start + index * pitch
        ring = Pos(0, 0, z) * Cylinder(
            major_diameter / 2.0,
            min(rib_width, z_start + length - z),
            align=(Align.CENTER, Align.CENTER, Align.MIN),
        )
        body = body + ring
    return body


def export_horizontal(body, label):
    result = body.rotate(Axis.Y, 90)
    result.label = label
    return result