V2 machine — drum + feeder

from __future__ import annotations

import math
import sys

import cadquery as cq

from spec import SPEC

if hasattr(sys.stdout, "reconfigure"):
    sys.stdout.reconfigure(encoding="utf-8", errors="replace")


# Drum axis is +X. Drum shell runs from x=0 to x=drum_length.
# +Y = drive side. +Z = up.


def hollow_cylinder_x(length: float, od: float, id_: float) -> cq.Workplane:
    outer = cq.Workplane("YZ").cylinder(length, od / 2.0, centered=(True, True, True))
    inner = cq.Workplane("YZ").cylinder(length + 2.0, id_ / 2.0, centered=(True, True, True))
    return outer.cut(inner)


def cyl_x(length: float, diameter: float) -> cq.Workplane:
    return cq.Workplane("YZ").cylinder(length, diameter / 2.0, centered=(True, True, True))


def cyl_y(length: float, diameter: float) -> cq.Workplane:
    return cq.Workplane("XZ").cylinder(length, diameter / 2.0, centered=(True, True, True))


def cyl_z(length: float, diameter: float) -> cq.Workplane:
    return cq.Workplane("XY").cylinder(length, diameter / 2.0, centered=(True, True, True))


def box_xyz(x: float, y: float, z: float) -> cq.Workplane:
    return cq.Workplane("XY").box(x, y, z)


def radial_hole_cutter_x(hole_d: float, radial_length: float, axial_x: float, radial_center_r: float, theta_deg: float) -> cq.Workplane:
    return (
        cq.Workplane("XZ")
        .cylinder(radial_length, hole_d / 2.0, centered=(True, True, True))
        .translate((axial_x, radial_center_r, 0))
        .rotate((0, 0, 0), (1, 0, 0), theta_deg)
    )


def axial_hole_cutter_x(hole_d: float, axial_length: float, y: float, z: float) -> cq.Workplane:
    return cq.Workplane("YZ").cylinder(axial_length, hole_d / 2.0, centered=(True, True, True)).translate((0, y, z))


def cut_radial_bolt_pattern(part: cq.Workplane, count: int, hole_d: float, pcd: float, radial_length: float) -> cq.Workplane:
    for i in range(count):
        part = part.cut(radial_hole_cutter_x(hole_d, radial_length, 0.0, pcd / 2.0, 360.0 * i / count))
    return part


def make_structural_ring(s: dict) -> cq.Workplane:
    ring = hollow_cylinder_x(s["ring_thickness_mm"], s["ring_od_mm"], s["ring_id_mm"])
    ring = cut_radial_bolt_pattern(
        ring,
        int(s["ring_bolt_count"]),
        float(s["ring_bolt_clearance_d_mm"]),
        float(s["ring_bolt_pcd_mm"]),
        float(s["ring_od_mm"] - s["ring_id_mm"] + 40.0),
    )
    drain_r = float(s["ring_drain_pcd_mm"]) / 2.0
    for theta in (-98.0, -82.0)[: int(s["ring_drain_count"])]:
        y = drain_r * math.cos(math.radians(theta))
        z = drain_r * math.sin(math.radians(theta))
        ring = ring.cut(axial_hole_cutter_x(float(s["ring_drain_d_mm"]), float(s["ring_thickness_mm"]) + 4.0, y, z))
    return ring


def make_running_band(s: dict) -> cq.Workplane:
    band = hollow_cylinder_x(s["band_axial_width_mm"], s["band_od_mm"], s["band_id_mm"])
    return cut_radial_bolt_pattern(
        band,
        int(s["band_bolt_count"]),
        float(s["band_bolt_clearance_d_mm"]),
        float(s["band_bolt_pcd_mm"]),
        float(s["band_od_mm"] - s["band_id_mm"] + 40.0),
    )


def make_side_wall(s: dict) -> cq.Workplane:
    return hollow_cylinder_x(float(s["side_wall_thickness_mm"]), float(s["side_wall_od_mm"]), float(s["side_wall_id_mm"]))


def make_inlet_cone(s: dict) -> cq.Workplane:
    length = float(s["cone_length_mm"])
    outer = (
        cq.Workplane("YZ")
        .circle(float(s["cone_mouth_od_mm"]) / 2.0)
        .workplane(offset=length)
        .circle(float(s["cone_drum_end_od_mm"]) / 2.0)
        .loft(combine=True)
    )
    inner = (
        cq.Workplane("YZ")
        .circle(float(s["cone_mouth_id_mm"]) / 2.0)
        .workplane(offset=length)
        .circle(float(s["cone_drum_end_id_mm"]) / 2.0)
        .loft(combine=True)
    )
    return outer.cut(inner).translate((-length, 0, 0))


def make_stringer(s: dict) -> cq.Workplane:
    length = float(s["stringer_length_mm"])
    width = float(s["stringer_width_mm"])
    height = float(s["stringer_height_mm"])
    clearance = float(s["stringer_end_clearance_d_mm"])
    return (
        cq.Workplane("XY")
        .box(length, width, height)
        .faces(">Z")
        .workplane()
        .pushPoints([(-length / 2.0 + 55.0, 0), (length / 2.0 - 55.0, 0)])
        .hole(clearance)
    )


def oriented_internal_part(part: cq.Workplane, x_pos: float, theta_deg: float, radial_z_local: float) -> cq.Workplane:
    return part.translate((0, 0, radial_z_local)).rotate((0, 0, 0), (1, 0, 0), theta_deg).translate((x_pos, 0, 0))


def make_lifter_body(s: dict) -> cq.Workplane:
    return cq.Workplane("XY").box(
        float(s["lifter_axial_length_mm"]),
        float(s["lifter_plate_thickness_mm"]),
        float(s["lifter_inward_height_mm"]),
    )


def make_lifter_wear_bar(s: dict) -> cq.Workplane:
    bar = cq.Workplane("XY").box(
        float(s["lifter_axial_length_mm"]),
        float(s["lifter_wear_bar_thickness_mm"]),
        float(s["lifter_wear_bar_height_mm"]),
    )
    x_span = float(s["lifter_axial_length_mm"]) * 0.32
    return bar.faces(">Y").workplane().pushPoints([(-x_span, 0), (x_span, 0)]).hole(11.0)


def make_lifter_tab(s: dict) -> cq.Workplane:
    lx, ly, lz = [float(v) for v in s["lifter_tab_mm"]]
    tab = cq.Workplane("XY").box(lx, ly, lz)
    return tab.faces(">Y").workplane().pushPoints([(-18.0, 0.0), (18.0, 0.0)]).hole(float(s["lifter_tab_bolt_clearance_d_mm"]))


def add_box(asm: cq.Assembly, shape: cq.Workplane, loc: tuple[float, float, float], name: str, color: cq.Color) -> None:
    asm.add(shape.translate(loc), name=name, color=color)


def raceway_contact_center(s: dict, theta_deg: float, roller_radius: float) -> tuple[float, float]:
    r = float(s["band_od_mm"]) / 2.0 + roller_radius
    return r * math.cos(math.radians(theta_deg)), r * math.sin(math.radians(theta_deg))


def make_pillow_block() -> cq.Workplane:
    base = box_xyz(80.0, 95.0, 24.0).translate((0, 0, -32.0))
    upright = box_xyz(46.0, 82.0, 56.0)
    bore = cyl_x(50.0, 42.0)
    return base.union(upright).cut(bore)


def make_anchor_bolt_stack() -> cq.Workplane:
    bolt = cyl_z(34.0, 18.0).translate((0, 0, 17.0))
    washer = cyl_z(4.0, 42.0).translate((0, 0, 36.0))
    nut = cq.Workplane("XY").polygon(6, 34.0).extrude(16.0).translate((0, 0, 40.0))
    return bolt.union(washer).union(nut)


def make_jack_foot(s: dict, extension_extra: float = 0.0) -> cq.Workplane:
    foot_x, foot_y, foot_t = [float(v) for v in s["frame_foot_plate_mm"]]
    screw_h = float(s["frame_jack_screw_height_mm"]) + extension_extra
    plate = box_xyz(foot_x, foot_y, foot_t)
    screw = cyl_z(screw_h, float(s["frame_jack_screw_d_mm"])).translate((0, 0, foot_t / 2.0 + screw_h / 2.0))
    pad = cyl_z(12.0, 72.0).translate((0, 0, foot_t / 2.0 + screw_h + 6.0))
    locknut = cq.Workplane("XY").polygon(6, 70.0).extrude(18.0).translate((0, 0, foot_t / 2.0 + screw_h * 0.55))
    return plate.union(screw).union(pad).union(locknut)


def make_motor_mount_plate(s: dict) -> cq.Workplane:
    plate = box_xyz(520.0, 780.0, 16.0)
    return (
        plate.faces(">Z")
        .workplane()
        .pushPoints([(-180.0, -110.0), (-180.0, 110.0), (180.0, -110.0), (180.0, 110.0)])
        .hole(float(s["gearmotor_mount_clearance_d_mm"]))
    )


def make_visual_bolt_z(d: float = 12.0, h: float = 12.0) -> cq.Workplane:
    washer = cyl_z(3.0, d * 2.1).translate((0, 0, 1.5))
    head = cq.Workplane("XY").polygon(6, d * 1.7).extrude(h).translate((0, 0, 3.0))
    return washer.union(head)


def make_belt_span_x(width_x: float, length: float, thickness: float, angle_deg_from_y: float) -> cq.Workplane:
    return box_xyz(width_x, length, thickness).rotate((0, 0, 0), (1, 0, 0), angle_deg_from_y)


def add_frame(asm: cq.Assembly, s: dict) -> None:
    rail_len = float(s["frame_length_mm"])
    y_rail = float(s["frame_rail_y_mm"])
    z_rail = float(s["frame_rail_z_mm"])
    ptr_w = float(s["frame_ptr_w_mm"])
    ptr_h = float(s["frame_ptr_h_mm"])
    x_center = 2000.0

    for y, side in ((-y_rail, "passive -Y"), (y_rail, "drive +Y")):
        add_box(
            asm,
            box_xyz(rail_len, ptr_w, ptr_h),
            (x_center, y, z_rail),
            f"FRAME · FRM-90001 long rail {side} — PTR 80×40 painted/coated steel (welded weldment, anchored to floor at jack-screw feet)",
            cq.Color(0.12, 0.16, 0.18),
        )

    for idx, x in enumerate(s["frame_crossmember_x_mm"], start=1):
        add_box(
            asm,
            box_xyz(ptr_h, 2.0 * y_rail + ptr_w, ptr_w),
            (float(x), 0.0, z_rail),
            f"FRAME · FRM-90002 crossmember {idx}/{len(s['frame_crossmember_x_mm'])} — PTR 80×40 painted/coated steel (welded to long rails)",
            cq.Color(0.13, 0.17, 0.19),
        )

    for side_y, label in ((-y_rail, "passive"), (y_rail, "drive")):
        brace = box_xyz(1450.0, 24.0, 32.0).rotate((0, 0, 0), (0, 1, 0), -10.0)
        add_box(
            asm,
            brace,
            (2000.0, side_y, z_rail - 95.0),
            f"FRAME · FRM-90003 diagonal brace {label} side — painted steel (welded to FRM-90001 rail and FRM-90002 crossmembers)",
            cq.Color(0.10, 0.13, 0.15),
        )

    leg_h = float(s["frame_leg_height_mm"])
    ground_z = float(s["ground_z_mm"])
    for xi, x in enumerate(s["frame_leg_x_mm"], start=1):
        for yi, y in enumerate(s["frame_leg_y_mm"], start=1):
            feed_end = float(x) < 0.0
            extension = 90.0 if feed_end else 20.0
            add_box(
                asm,
                box_xyz(ptr_w, ptr_h, leg_h),
                (float(x), float(y), z_rail - ptr_h / 2.0 - leg_h / 2.0),
                f"FRAME · FRM-90004 vertical post/leg x{xi} y{yi} — PTR 80×40 (welded to frame weldment, bears on jack screw)",
                cq.Color(0.12, 0.16, 0.18),
            )
            asm.add(
                make_jack_foot(s, extension).translate((float(x), float(y), ground_z)),
                name=(
                    f"FRAME · FRM-90005 adjustable jack screw/foot x{xi} y{yi} — one of 4 corner jacks, "
                    f"{'feed-end high jack' if feed_end else 'discharge-end low jack'}, 0–6° tilt range "
                    f"(bolted/anchored to floor via {s['frame_anchor_bolt']})"
                ),
                color=cq.Color(0.22, 0.22, 0.20),
            )
            asm.add(
                box_xyz(110.0, 6.0, 46.0).translate((float(x) + 95.0, float(y), ground_z + 190.0)),
                name=f"FRAME · FRM-90007 angle scale plate x{xi} y{yi} — 0–6° drum-frame tilt indicator (bolted to jack bracket)",
                color=cq.Color(0.85, 0.78, 0.28),
            )
            for bi, (ax, ay) in enumerate(((-65.0, -45.0), (-65.0, 45.0), (65.0, -45.0), (65.0, 45.0)), start=1):
                asm.add(
                    make_anchor_bolt_stack().translate((float(x) + ax, float(y) + ay, ground_z + 8.0)),
                    name=f"FRAME · FRM-90006 M16 anchor bolt foot{xi}{yi} pos{bi}/4 — Ø18 clearance (anchors jack foot to deck)",
                    color=cq.Color(0.05, 0.05, 0.055),
                )


def add_support_rollers(asm: cq.Assembly, s: dict) -> None:
    rr = float(s["support_roller_d_mm"]) / 2.0
    width = float(s["support_roller_face_width_mm"])
    shaft_d = float(s["support_roller_shaft_d_mm"])
    angle = float(s["support_roller_angle_from_bottom_deg"])
    z_rail = float(s["frame_rail_z_mm"])

    for station_idx, x in enumerate(s["support_station_x_mm"], start=1):
        for theta, side in ((-90.0 - angle, "passive -Y"), (-90.0 + angle, "drive +Y")):
            y, z = raceway_contact_center(s, theta, rr)
            asm.add(
                cyl_x(width, float(s["support_roller_d_mm"])).translate((float(x), y, z)),
                name=(
                    f"ROLLER · ROL-70001 support roller station {station_idx} {side} — Ø{s['support_roller_d_mm']:.0f} × "
                    f"{width:.0f} hardened steel Rc45 (shaft supported by pillow blocks bolted to FRAME crossmember)"
                ),
                color=cq.Color(0.46, 0.47, 0.49),
            )
            asm.add(
                cyl_x(width + 130.0, shaft_d).translate((float(x), y, z)),
                name=f"ROLLER · ROL-70002 support shaft station {station_idx} {side} — Ø{shaft_d:.0f} steel shaft (captured by two pillow blocks)",
                color=cq.Color(0.20, 0.21, 0.22),
            )
            for dx, end_name in ((-(width / 2.0 + 45.0), "inboard"), (width / 2.0 + 45.0, "outboard")):
                bx = float(x) + dx
                bz = z - 48.0
                pedestal_h = max(40.0, (bz - 44.0) - z_rail)
                asm.add(
                    box_xyz(95.0, 120.0, pedestal_h).translate((bx, y, z_rail + pedestal_h / 2.0)),
                    name=(
                        f"ROLLER · BRK-70004 heavy bearing pedestal {end_name}, station {station_idx} {side} — "
                        f"painted steel bracket (welded/bolted to FRM-90002 frame crossmember at x={float(x):.0f})"
                    ),
                    color=cq.Color(0.10, 0.12, 0.13),
                )
                asm.add(
                    make_pillow_block().translate((bx, y, bz)),
                    name=(
                        f"ROLLER · BRG-70003 sealed pillow-block bearing {end_name}, station {station_idx} {side} — "
                        f"{s['support_roller_bearing']} (bolted to FRM-90002 frame crossmember via BRK-70004, 4 × M12)"
                    ),
                    color=cq.Color(0.08, 0.10, 0.11),
                )
            shield_y = y + (38.0 if y > 0 else -38.0)
            asm.add(
                box_xyz(width + 180.0, 16.0, 150.0).translate((float(x), shield_y, z + 5.0)),
                name=f"SAFETY GUARD · GRD-13001 support roller sand shield station {station_idx} {side} — transparent yellow shield (bolted to roller pedestal)",
                color=cq.Color(0.92, 0.72, 0.18, 0.45),
            )


def add_motor_cover(asm: cq.Assembly, s: dict, mx: float, my: float, mz: float, ml: float, mw: float, mh: float) -> None:
    c = float(s["motor_cover_clearance_mm"])
    t = 2.0
    ox = ml + 2.0 * c
    oy = mw + 2.0 * c
    oz = mh + c
    top_z = mz + mh / 2.0 + c
    bottom_open_z = mz - mh / 2.0

    panels = [
        ("top", box_xyz(ox, oy, t).translate((mx, my, top_z))),
        ("drive long side +Y", box_xyz(ox, t, oz).translate((mx, my + oy / 2.0, bottom_open_z + oz / 2.0))),
        ("inside long side -Y", box_xyz(ox, t, oz).translate((mx, my - oy / 2.0, bottom_open_z + oz / 2.0))),
        ("upstream short side", box_xyz(t, oy, oz).translate((mx - ox / 2.0, my, bottom_open_z + oz / 2.0))),
        ("downstream short side", box_xyz(t, oy, oz).translate((mx + ox / 2.0, my, bottom_open_z + oz / 2.0))),
    ]
    for label, panel in panels:
        asm.add(
            panel,
            name=f"SAFETY GUARD · GRD-13004 gearmotor 5-side cover {label} — 2 mm A36 bolt-on wrap, bottom open for ventilation (bolted to motor skid via M8 tabs)",
            color=cq.Color(0.62, 0.70, 0.78, 0.38),
        )

    for i in range(int(s["motor_cover_slot_vent_count"])):
        asm.add(
            box_xyz(90.0, 3.0, 14.0).translate((mx - 180.0 + i * 72.0, my + oy / 2.0 + 2.0, mz + 10.0)),
            name=f"SAFETY GUARD · GRD-13005 gearmotor cover slot vent {i + 1}/{s['motor_cover_slot_vent_count']} — black visual opening on +Y side",
            color=cq.Color(0.02, 0.02, 0.02),
        )


def add_drive_roller_and_transmission(asm: cq.Assembly, s: dict) -> None:
    drive_x = float(s["drive_roller_x_mm"])
    dr = float(s["drive_roller_d_mm"]) / 2.0
    drive_theta = -6.0
    drive_y, drive_z = raceway_contact_center(s, drive_theta, dr)
    width = float(s["drive_roller_face_width_mm"])
    z_rail = float(s["frame_rail_z_mm"])

    asm.add(
        cyl_x(width, float(s["drive_roller_d_mm"])).translate((drive_x, drive_y, drive_z)),
        name=f"DRIVE · DRV-80010 lagged drive roller — Ø{s['drive_roller_d_mm']:.0f} × {width:.0f}, {s['drive_roller_lagging']} (frame-mounted friction drive on downstream raceway)",
        color=cq.Color(0.02, 0.02, 0.025),
    )
    asm.add(
        cyl_x(width + 240.0, float(s["drive_roller_shaft_d_mm"])).translate((drive_x, drive_y, drive_z)),
        name=f"DRIVE · DRV-80011 drive roller shaft — Ø{s['drive_roller_shaft_d_mm']:.0f} steel (captured by frame-mounted pillow blocks)",
        color=cq.Color(0.16, 0.17, 0.18),
    )

    for dx, label in ((-(width / 2.0 + 75.0), "inboard"), (width / 2.0 + 75.0, "pulley side")):
        bx = drive_x + dx
        bz = drive_z - 58.0
        pedestal_h = max(120.0, bz - z_rail - 44.0)
        asm.add(
            box_xyz(110.0, 130.0, pedestal_h).translate((bx, drive_y, z_rail + pedestal_h / 2.0)),
            name=f"DRIVE · BRK-80014 drive bearing pedestal {label} — heavy steel bracket (bolted/welded to downstream FRM-90002 frame crossmember, not to plate)",
            color=cq.Color(0.10, 0.12, 0.13),
        )
        asm.add(
            make_pillow_block().translate((bx, drive_y, bz)),
            name=f"DRIVE · BRG-80012 drive roller pillow-block bearing {label} — sealed 2RS (bolted to FRAME crossmember via BRK-80014, 4 × M12)",
            color=cq.Color(0.08, 0.10, 0.11),
        )

    screw = cyl_y(260.0, 28.0).translate((drive_x, drive_y + 105.0, drive_z - 20.0))
    spring = cyl_y(115.0, 60.0).translate((drive_x, drive_y + 48.0, drive_z - 20.0))
    bracket = box_xyz(220.0, 18.0, 190.0).translate((drive_x, drive_y + 210.0, drive_z - 20.0))
    asm.add(
        screw.union(spring).union(bracket),
        name=f"DRIVE · DRV-80013 screw/spring preload assembly — target {s['drive_roller_nominal_preload_kn'][0]:.1f}–{s['drive_roller_nominal_preload_kn'][1]:.1f} kN (bolted to frame drive pedestal)",
        color=cq.Color(0.18, 0.18, 0.16),
    )

    nip_guard = box_xyz(width + 260.0, 18.0, 260.0).rotate((0, 0, 0), (1, 0, 0), -8.0)
    asm.add(
        nip_guard.translate((drive_x, drive_y - 68.0, drive_z + 80.0)),
        name="SAFETY GUARD · GRD-13002 drive roller nip guard — transparent yellow shield (bolted to drive pedestal)",
        color=cq.Color(0.95, 0.74, 0.16, 0.45),
    )

    mx, my, mz = [float(v) for v in s["gearmotor_location_mm"]]
    ml, mw, mh = [float(v) for v in s["gearmotor_envelope_mm"]]
    asm.add(
        box_xyz(ml, mw, mh).translate((mx, my, mz)),
        name=f"DRIVE · DRV-80020 gearmotor 3.0 kW — {s['gearmotor_material']}, envelope {ml:.0f} × {mw:.0f} × {mh:.0f} (bolted to skid via 4 × M12)",
        color=cq.Color(0.12, 0.28, 0.48),
    )
    asm.add(
        make_motor_mount_plate(s).translate((mx, my, mz - mh / 2.0 - 18.0)),
        name="DRIVE · DRV-80021 gearmotor skid/base plate — widened steel plate reaches drive-side frame rail (bolted to FRAME rail via M12 slotted tension bolts)",
        color=cq.Color(0.18, 0.20, 0.21),
    )
    for bx, by in ((-180.0, -110.0), (-180.0, 110.0), (180.0, -110.0), (180.0, 110.0)):
        asm.add(
            make_visual_bolt_z(12.0, 12.0).translate((mx + bx, my + by, mz - mh / 2.0 - 7.0)),
            name=f"DRIVE · DRV-80022 M12 gearmotor mount bolt visual ({bx:+.0f},{by:+.0f}) — slotted base locknut",
            color=cq.Color(0.04, 0.04, 0.045),
        )
    add_motor_cover(asm, s, mx, my, mz, ml, mw, mh)

    pulley_x = drive_x + width / 2.0 + 125.0
    pulley_d = float(s["pulley_d_mm"])
    pulley_w = float(s["pulley_face_width_mm"])
    driven_pulley_center = (pulley_x, drive_y, drive_z)
    motor_pulley_center = (pulley_x, my, mz + 20.0)

    for center, label in ((driven_pulley_center, "drive-roller shaft"), (motor_pulley_center, "gearmotor shaft")):
        px, py, pz = center
        asm.add(
            cyl_x(pulley_w, pulley_d).translate((px, py, pz)),
            name=f"DRIVE · DRV-80030 Ø{pulley_d:.0f} sheave/pulley on {label} — steel sheave (keyed/bolted to shaft)",
            color=cq.Color(0.07, 0.07, 0.075),
        )
        asm.add(
            cyl_x(pulley_w + 24.0, 74.0).translate((px, py, pz)),
            name=f"DRIVE · DRV-80031 pulley hub on {label} — steel hub with M10 shaft-bolt callout",
            color=cq.Color(0.16, 0.16, 0.17),
        )

    dy = motor_pulley_center[1] - driven_pulley_center[1]
    dz = motor_pulley_center[2] - driven_pulley_center[2]
    cdist = math.sqrt(dy * dy + dz * dz)
    angle = math.degrees(math.atan2(dz, dy))
    normal_y = -math.sin(math.radians(angle))
    normal_z = math.cos(math.radians(angle))
    belt_offset = pulley_d / 2.0 - 8.0
    for sign, label in ((1.0, "outer tight span"), (-1.0, "inner return span")):
        ymid = (motor_pulley_center[1] + driven_pulley_center[1]) / 2.0 + sign * normal_y * belt_offset
        zmid = (motor_pulley_center[2] + driven_pulley_center[2]) / 2.0 + sign * normal_z * belt_offset
        asm.add(
            make_belt_span_x(pulley_w + 10.0, cdist, 14.0, angle).translate((pulley_x, ymid, zmid)),
            name=f"DRIVE · DRV-80032 V-belt visual {label} — {s['belt_type']}",
            color=cq.Color(0.01, 0.01, 0.012),
        )

    asm.add(
        box_xyz(pulley_w + 70.0, abs(dy) + pulley_d + 160.0, abs(dz) + pulley_d + 140.0).translate(
            (pulley_x, (motor_pulley_center[1] + driven_pulley_center[1]) / 2.0, (motor_pulley_center[2] + driven_pulley_center[2]) / 2.0)
        ),
        name=f"SAFETY GUARD · GRD-13003 full belt/pulley guard — transparent yellow envelope, {s['belt_guard_material']} (bolted to gearmotor skid/frame)",
        color=cq.Color(0.95, 0.74, 0.16, 0.33),
    )


def add_axial_guides(asm: cq.Assembly, s: dict) -> None:
    band_w = float(s["band_axial_width_mm"])
    sw_t = float(s["side_wall_thickness_mm"])
    guide_d = float(s["guide_roller_d_mm"])
    guide_w = float(s["guide_roller_face_width_mm"])
    lip_r = float(s["side_wall_od_mm"]) / 2.0
    guide_zs = [120.0, -120.0]
    end_defs = (
        ("upstream", float(s["band_x_inset_mm"]) - band_w / 2.0 - sw_t, -1.0),
        ("downstream", float(s["drum_length_mm"]) - float(s["band_x_inset_mm"]) + band_w / 2.0 + sw_t, 1.0),
    )
    for end_label, x_face, axial_sign in end_defs:
        for i, z in enumerate(guide_zs, start=1):
            y_surface = math.sqrt(max(lip_r * lip_r - z * z, 0.0))
            guide_y = y_surface + guide_w / 2.0 + 2.0
            guide_x = x_face + axial_sign * (guide_d / 2.0 + 2.0)
            asm.add(
                cyl_y(guide_w, guide_d).translate((guide_x, guide_y, z)),
                name=f"ROLLER · ROL-70020 axial guide roller {end_label} {i}/2 — Ø{guide_d:.0f} × {guide_w:.0f}, acts on PRT-80002 lip only (bolted to frame bracket)",
                color=cq.Color(0.42, 0.43, 0.45),
            )
            arm = box_xyz(120.0, 22.0, 32.0).translate((guide_x, guide_y + 54.0, z))
            bracket = box_xyz(80.0, 18.0, 140.0).translate((guide_x, guide_y + 118.0, z - 28.0))
            asm.add(
                arm.union(bracket),
                name=f"ROLLER · BRG-70021 axial-guide adjustable bracket {end_label} {i}/2 — slotted painted steel (bolted to FRAME rail)",
                color=cq.Color(0.10, 0.12, 0.13),
            )


def add_vfd_panel(asm: cq.Assembly, s: dict) -> None:
    vx, vy, vz = [float(v) for v in s["vfd_location_mm"]]
    vw, vd, vh = [float(v) for v in s["vfd_enclosure_mm"]]
    asm.add(
        box_xyz(vw, vd, vh).translate((vx, vy, vz)),
        name=f"HMI · ELE-16001 NEMA 4X VFD/control enclosure — stainless/polymer washdown cabinet, {s['vfd_spec']} (bolted to frame/service post)",
        color=cq.Color(0.80, 0.82, 0.78),
    )
    asm.add(
        box_xyz(vw - 36.0, 8.0, vh - 60.0).translate((vx, vy - vd / 2.0 - 5.0, vz)),
        name="HMI · ELE-16002 control-panel door face — Spanish labels, display, hour meter, lockable disconnect (hinged to ELE-16001)",
        color=cq.Color(0.68, 0.70, 0.66),
    )
    asm.add(cyl_y(24.0, 56.0).translate((vx - 135.0, vy - vd / 2.0 - 18.0, vz + 140.0)), name="HMI · ELE-16003 lockable main disconnect handle — panel-mounted", color=cq.Color(0.10, 0.10, 0.10))
    asm.add(cyl_y(28.0, 70.0).translate((vx + 135.0, vy - vd / 2.0 - 20.0, vz + 130.0)), name="HMI · ELE-16004 emergency stop mushroom — panel-mounted", color=cq.Color(0.85, 0.02, 0.02))
    for idx, (cx, cz, col, label) in enumerate(
        [
            (-90.0, 30.0, cq.Color(0.02, 0.55, 0.12), "START green"),
            (0.0, 30.0, cq.Color(0.88, 0.80, 0.05), "RESET/JOG amber"),
            (90.0, 30.0, cq.Color(0.85, 0.02, 0.02), "STOP red"),
        ],
        start=1,
    ):
        asm.add(cyl_y(20.0, 34.0).translate((vx + cx, vy - vd / 2.0 - 18.0, vz + cz)), name=f"HMI · ELE-16005-{idx} {label} button — IP-rated panel operator", color=col)

    tower_x = vx + vw / 2.0 - 60.0
    tower_base_z = vz + vh / 2.0 + 35.0
    asm.add(cyl_z(70.0, 20.0).translate((tower_x, vy, tower_base_z)), name="HMI · ELE-16006 light-tower mast — bolted to control enclosure", color=cq.Color(0.08, 0.08, 0.08))
    for i, (col, label) in enumerate([(cq.Color(0.10, 0.75, 0.18), "RUN green"), (cq.Color(0.95, 0.75, 0.05), "WARN amber"), (cq.Color(0.88, 0.02, 0.02), "FAULT red")], start=1):
        asm.add(cyl_z(36.0, 54.0).translate((tower_x, vy, tower_base_z + 35.0 + (i - 1) * 38.0)), name=f"HMI · ELE-16007-{i} 3-stage light tower {label} lens — stacked on mast", color=col)

    asm.add(box_xyz(160.0, 70.0, 120.0).translate((3300.0, -760.0, -340.0)), name="HMI · ELE-16008 remote e-stop station enclosure — maintenance-side box (bolted to frame)", color=cq.Color(0.82, 0.82, 0.76))
    asm.add(cyl_y(28.0, 62.0).translate((3300.0, -800.0, -330.0)), name="HMI · ELE-16009 emergency stop mushroom — red, discharge/maintenance side (panel-mounted)", color=cq.Color(0.85, 0.02, 0.02))


def bar_between_yz(x_center: float, p1_yz: tuple[float, float], p2_yz: tuple[float, float], x_width: float, z_thick: float) -> cq.Workplane:
    y1, z1 = p1_yz
    y2, z2 = p2_yz
    dy = y2 - y1
    dz = z2 - z1
    length = math.sqrt(dy * dy + dz * dz)
    angle = math.degrees(math.atan2(dz, dy))
    return box_xyz(x_width, length, z_thick).rotate((0, 0, 0), (1, 0, 0), angle).translate((x_center, (y1 + y2) / 2.0, (z1 + z2) / 2.0))


def cylinder_between_yz_xlocal(x_center: float, p1_yz: tuple[float, float], p2_yz: tuple[float, float], diameter: float) -> cq.Workplane:
    y1, z1 = p1_yz
    y2, z2 = p2_yz
    dy = y2 - y1
    dz = z2 - z1
    length = math.sqrt(dy * dy + dz * dz)
    angle = math.degrees(math.atan2(dz, dy))
    return cyl_y(length, diameter).rotate((0, 0, 0), (1, 0, 0), angle).translate((x_center, (y1 + y2) / 2.0, (z1 + z2) / 2.0))


def make_compression_spring_visual(x: float, start_yz: tuple[float, float], end_yz: tuple[float, float], spring_d: float) -> cq.Workplane:
    y1, z1 = start_yz
    y2, z2 = end_yz
    dy = y2 - y1
    dz = z2 - z1
    length = math.sqrt(dy * dy + dz * dz)
    angle = math.degrees(math.atan2(dz, dy))
    core = cyl_y(length, spring_d * 0.18).rotate((0, 0, 0), (1, 0, 0), angle).translate((x, (y1 + y2) / 2.0, (z1 + z2) / 2.0))
    spring = core
    coils = 8
    for i in range(coils):
        t = (i + 0.5) / coils
        cy = y1 + dy * t
        cz = z1 + dz * t
        ring = cyl_y(8.0, spring_d).rotate((0, 0, 0), (1, 0, 0), angle).translate((x, cy, cz))
        spring = spring.union(ring)
    return spring


def add_brushes(asm: cq.Assembly, s: dict) -> None:
    brush_len = float(s["brush_bristle_length_mm"])
    brush_d = float(s["brush_bristle_d_mm"])
    brush_r = brush_d / 2.0
    band_r = float(s["band_od_mm"]) / 2.0
    overlap = float(s.get("brush_bristle_overlap_mm", 5.0))
    contact_theta = -45.0

    # One and only one tangential cylindrical bristle bundle per band.
    # Bundle axis is parallel to drum X. Bristles touch the running-band OD at approx. 4 o'clock with ~5 mm overlap.
    bristle_center_r = band_r + brush_r - overlap
    bristle_y = bristle_center_r * math.cos(math.radians(contact_theta))
    bristle_z = bristle_center_r * math.sin(math.radians(contact_theta))

    # Frame-mounted pivot is radially outward from the bristle bundle, close to the drive-side rail/crossmember.
    pivot_r = float(s.get("brush_frame_pivot_radial_mm", 916.0))
    pivot_y = pivot_r * math.cos(math.radians(contact_theta))
    pivot_z = pivot_r * math.sin(math.radians(contact_theta))

    flat_bar_x, flat_bar_z = [float(v) for v in s.get("brush_pivot_arm_flat_bar_mm", [70.0, 18.0])]
    spring_d = float(s.get("brush_spring_d_mm", 50.0))

    band_centers = [
        (float(s["band_x_inset_mm"]), "upstream"),
        (float(s["drum_length_mm"]) - float(s["band_x_inset_mm"]), "downstream"),
    ]

    for idx, (x, label) in enumerate(band_centers, start=1):
        bristle_yz = (bristle_y, bristle_z)
        pivot_yz = (pivot_y, pivot_z)
        arm_mid_yz = ((bristle_y + pivot_y) / 2.0, (bristle_z + pivot_z) / 2.0)

        # Fixed bracket/lug on the frame side for the compression spring; positioned so the spring is visibly loading the arm inward.
        spring_fixed_yz = (pivot_y - 24.0, pivot_z - 190.0)

        asm.add(
            cyl_x(brush_len, brush_d).translate((x, bristle_y, bristle_z)),
            name=(
                f"AB-10001-{idx} anti-blinding bristle bundle, {label} running band — ONE Ø{brush_d:.0f} × {brush_len:.0f} "
                f"tangential cylinder, axis parallel to drum X, bristles overlap PRT-80001 OD by {overlap:.0f} mm at 4-o'clock"
            ),
            color=cq.Color(0.06, 0.055, 0.035),
        )

        asm.add(
            cyl_y(150.0, float(s["brush_pivot_shaft_d_mm"])).translate((x, pivot_y, pivot_z)),
            name=(
                f"AB-10002-{idx} frame-mounted brush pivot pin, {label} band — Ø{s['brush_pivot_shaft_d_mm']:.0f} horizontal pin "
                f"perpendicular to drum axis, captured by welded frame bracket at drive-side rail/crossmember"
            ),
            color=cq.Color(0.18, 0.18, 0.17),
        )

        asm.add(
            bar_between_yz(x, bristle_yz, pivot_yz, flat_bar_x, flat_bar_z),
            name=(
                f"AB-10003-{idx} single radial pivot arm, {label} band — flat bar from bristle-bundle midpoint to frame pivot; "
                f"no duplicate arms, no floating parts"
            ),
            color=cq.Color(0.14, 0.14, 0.13),
        )

        asm.add(
            make_compression_spring_visual(x, arm_mid_yz, spring_fixed_yz, spring_d),
            name=(
                f"AB-10004-{idx} compression spring preload, {label} band — Ø{spring_d:.0f} spring between pivot-arm midpoint "
                f"and fixed frame bracket, loading bristles onto running band"
            ),
            color=cq.Color(0.72, 0.70, 0.62),
        )

        bracket_web = bar_between_yz(x, (float(s["frame_rail_y_mm"]), float(s["frame_rail_z_mm"])), pivot_yz, 80.0, 14.0)
        pivot_lug = box_xyz(120.0, 26.0, 90.0).translate((x, pivot_y, pivot_z))
        spring_lug = box_xyz(95.0, 22.0, 75.0).translate((x, spring_fixed_yz[0], spring_fixed_yz[1]))
        asm.add(
            bracket_web.union(pivot_lug).union(spring_lug),
            name=(
                f"AB-10005-{idx} welded frame bracket for brush, {label} band — anchored to drive-side FRAME rail/crossmember; "
                f"carries pivot pin and spring fixed lug"
            ),
            color=cq.Color(0.10, 0.12, 0.13),
        )


def add_sand_deflector(asm: cq.Assembly, s: dict) -> None:
    length = float(s["sand_deflector_length_mm"])
    span = float(s["sand_deflector_span_mm"])
    thick = float(s["sand_deflector_thickness_mm"])
    tilt = float(s["sand_deflector_tilt_deg"])
    x_center = float(s["sand_deflector_x_center_mm"])
    z_center = float(s["sand_deflector_z_center_mm"])
    plate = box_xyz(length, span, thick).rotate((0, 0, 0), (1, 0, 0), tilt)
    asm.add(
        plate.translate((x_center, 0.0, z_center)),
        name=f"DEF-90040 single inclined sand-deflector plate — {s['sand_deflector_material']}, {length:.0f} × {span:.0f} × {thick:.0f}, tilted {tilt:.0f}° toward {s['sand_deflector_slope_side']} (bolted to frame tabs, old multi-plate pan removed)",
        color=cq.Color(0.58, 0.62, 0.62, 0.72),
    )
    low_y = -span * math.cos(math.radians(tilt)) / 2.0 - 80.0
    low_z = z_center - span * math.sin(math.radians(tilt)) / 2.0 - 15.0
    asm.add(
        box_xyz(length, 28.0, 42.0).translate((x_center, low_y, low_z)),
        name="DEF-90041 low-side drip edge / collection lip — 304 stainless angle (welded to DEF-90040, directs sand to side collection cart)",
        color=cq.Color(0.46, 0.50, 0.50),
    )


def add_misc_guards_and_reference(asm: cq.Assembly, s: dict) -> None:
    for y, side in ((-710.0, "passive -Y"), (710.0, "drive +Y")):
        asm.add(
            box_xyz(3850.0, 18.0, 480.0).translate((2000.0, y, -250.0)),
            name=f"SAFETY GUARD · GRD-13010 removable side guard {side} — transparent yellow roller/stringer access panel (bolted to frame, tool-only removal)",
            color=cq.Color(0.95, 0.74, 0.16, 0.22),
        )

    cx, cy, cz = [float(v) for v in s["cart_660l_location_mm"]]
    cl, cw, ch = [float(v) for v in s["cart_660l_envelope_mm"]]
    asm.add(
        box_xyz(cl, cw, ch).translate((cx, cy, cz)),
        name=f"REF-90060 660 L wheeled waste cart clearance envelope — translucent reference only, {cl:.0f} × {cw:.0f} × {ch:.0f} mm, centered 250 mm beyond drum discharge end",
        color=cq.Color(0.10, 0.45, 0.18, 0.18),
    )


def build_assembly(s: dict = SPEC) -> cq.Assembly:
    asm = cq.Assembly(name=s["slug"])

    drum_len = float(s["drum_length_mm"])
    drum_od = float(s["drum_od_mm"])
    drum_id = float(s["drum_id_mm"])
    drum_inner_r = drum_id / 2.0

    asm.add(
        hollow_cylinder_x(drum_len, drum_od, drum_id).translate((drum_len / 2.0, 0, 0)),
        name=f"PRT-80005 drum shell / screen carrier — 8 mm shell with {s['screen_panel_material']}, Ø{drum_od:.0f} OD × Ø{drum_id:.0f} ID × {drum_len:.0f}, Ø4 mm perforation callout (screen panels bolted/welded to shell stringers)",
        color=cq.Color(0.72, 0.76, 0.78, 0.55),
    )

    seam_len = drum_len - 2.0 * float(s["ring_x_inset_mm"])
    for i in range(4):
        theta = i * 90.0 + 45.0
        seam = oriented_internal_part(cq.Workplane("XY").box(seam_len, 6.0, 4.0), drum_len / 2.0, theta, drum_inner_r - 2.0)
        asm.add(
            seam,
            name=f"PRT-80005 screen panel seam/backing strip {i + 1}/4 — 316L stainless backing strip (welded/bolted to screen shell)",
            color=cq.Color(0.55, 0.60, 0.62),
        )

    # Ring positions: upstream + midspan + downstream (3 rings total).
    # Midspan ring added 2026-04-30 per torsion/wobble analysis — drum length /
    # OD ratio = 4000/1100 = 3.6, industrial trommels at this ratio typically use
    # 3-4 rings, not 2. Mid-span ring shifts first natural mode away from operating
    # RPM and reduces drum body deflection under asymmetric sargassum load.
    ring_positions = (
        (float(s["ring_x_inset_mm"]),                      "upstream"),
        (drum_len / 2.0,                                    "midspan"),
        (drum_len - float(s["ring_x_inset_mm"]),            "downstream"),
    )
    for x_pos, label in ring_positions:
        asm.add(
            make_structural_ring(s).translate((x_pos, 0, 0)),
            name=f"PRT-80003 {label} structural ring — {s['ring_material']}, 14 mm, OD {s['ring_od_mm']:.0f}/ID {s['ring_id_mm']:.0f}, 12 × M10 PCD {s['ring_bolt_pcd_mm']:.0f} (welded to drum shell)",
            color=cq.Color(0.33, 0.34, 0.36),
        )

    band_centers = ((float(s["band_x_inset_mm"]), "upstream"), (drum_len - float(s["band_x_inset_mm"]), "downstream"))
    for x_pos, label in band_centers:
        asm.add(
            make_running_band(s).translate((x_pos, 0, 0)),
            name=f"PRT-80001 {label} running band / raceway — {s['band_material']}, OD {s['band_od_mm']:.0f}/ID {s['band_id_mm']:.0f} × {s['band_axial_width_mm']:.0f} wide (bolted through shell to PRT-80003 ring via same 12 × M10 PCD)",
            color=cq.Color(0.50, 0.50, 0.54),
        )

    sw_t = float(s["side_wall_thickness_mm"])
    band_w = float(s["band_axial_width_mm"])
    for band_x, band_label in band_centers:
        for side_name, offset in (("inner", -(band_w / 2.0 + sw_t / 2.0)), ("outer", band_w / 2.0 + sw_t / 2.0)):
            asm.add(
                make_side_wall(s).translate((band_x + offset, 0, 0)),
                name=f"PRT-80002 {band_label} {side_name} side wall / axial flange — {s['side_wall_material']}, OD {s['side_wall_od_mm']:.0f}/ID {s['side_wall_id_mm']:.0f} × {s['side_wall_thickness_mm']:.0f} (bolted to PRT-80003 ring via 12 × M10 PCD 1120)",
                color=cq.Color(0.64, 0.66, 0.68),
            )

    if bool(s.get("include_inlet_cone", False)):
        asm.add(
            make_inlet_cone(s),
            name=f"PRT-80004 upstream inlet cone — {s['cone_material']}, axial length {s['cone_length_mm']:.0f} (welded to upstream PRT-80002 side wall)",
            color=cq.Color(0.38, 0.42, 0.44),
        )

    n_str = int(s["stringer_count"])
    stringer_base = make_stringer(s)
    str_h = float(s["stringer_height_mm"])
    for i in range(n_str):
        theta = i * 360.0 / n_str
        asm.add(
            oriented_internal_part(stringer_base, drum_len / 2.0, theta, drum_inner_r - str_h / 2.0),
            name=f"PRT-80006 longitudinal stringer / carrier rail {i + 1}/{n_str} — {s['stringer_material']}, 40 × 10 × {s['stringer_length_mm']:.0f}, M8 end holes (welded to upstream + downstream PRT-80003 rings)",
            color=cq.Color(0.29, 0.30, 0.31),
        )

    cols = int(s["lifter_angular_cols"])
    rows = int(s["lifter_axial_rows"])
    total = cols * rows
    x_start = float(s["lifter_x_start_mm"])
    x_end = float(s["lifter_x_end_mm"])
    x_step = (x_end - x_start) / max(1, rows - 1)
    theta_step = 360.0 / cols
    helix_step = float(s["lifter_helix_angle_deg_per_row"])
    lift_h = float(s["lifter_inward_height_mm"])
    wear_h = float(s["lifter_wear_bar_height_mm"])
    tab_lx, tab_ly, tab_lz = [float(v) for v in s["lifter_tab_mm"]]

    lifter_body_base = make_lifter_body(s)
    wear_bar_base = make_lifter_wear_bar(s)
    tab_base = make_lifter_tab(s)

    placed = 0
    for row in range(rows):
        x_pos = x_start + row * x_step
        helix_offset = row * helix_step
        for col in range(cols):
            theta = col * theta_step + helix_offset
            placed += 1
            asm.add(
                oriented_internal_part(lifter_body_base, x_pos, theta, drum_inner_r - lift_h / 2.0),
                name=f"PRT-80007 lifter body {placed}/{total} — {s['lifter_body_material']}, row {row + 1}/{rows}, col {col + 1}/{cols}, θ={theta:.1f}° (welded to nearest stringer via tab)",
                color=cq.Color(0.55, 0.42, 0.25),
            )
            asm.add(
                oriented_internal_part(tab_base, x_pos, theta, drum_inner_r - tab_lz / 2.0),
                name=f"PRT-80007 lifter connection tab {placed}/{total} — small visible tab with 2 × M8 Ø9 hole positions (welded to nearest stringer via tab)",
                color=cq.Color(0.62, 0.48, 0.28),
            )
            asm.add(
                oriented_internal_part(wear_bar_base, x_pos, theta, drum_inner_r - lift_h + wear_h / 2.0),
                name=f"PRT-80007 Hardox wear bar {placed}/{total} — {s['lifter_wear_bar_material']} (bolted to lifter body via 2 × M10 countersunk clearance Ø11)",
                color=cq.Color(0.18, 0.18, 0.19),
            )

    add_frame(asm, s)
    add_support_rollers(asm, s)
    add_drive_roller_and_transmission(asm, s)
    add_axial_guides(asm, s)
    add_vfd_panel(asm, s)
    add_brushes(asm, s)
    add_sand_deflector(asm, s)
    add_misc_guards_and_reference(asm, s)

    return asm


def main() -> None:
    asm = build_assembly(SPEC)
    out = SPEC["outputs"]

    asm.save(out["step"])
    asm.save(out["glb"])

    print(
        f"{SPEC['title']}: exported {out['step']} and {out['glb']} with corrected frame-mounted anti-blinding brushes. "
        "STL export intentionally skipped."
    )


if __name__ == "__main__":
    main()

from cadquery_part import build_assembly, main
from spec import SPEC

assembly = build_assembly(SPEC)

if __name__ == "__main__":
    main()