import cadquery as cq
import cadkit as ck

# Parameters
# Shaft dimensions
S1_D = 20.0    # Diameter of first step
S1_L = 40.0    # Length of first step
S2_D = 25.0    # Diameter of second (middle) step
S2_L = 80.0    # Length of second step
S3_D = 20.0    # Diameter of third step
S3_L = 40.0    # Length of third step

# Keyway parameters
KEY_W = 6.0    # Keyway width
KEY_L = 50.0   # Keyway length
KEY_H = 3.0    # Keyway depth (standard for 25mm shaft is approx 3-4mm)

# General
CHAMFER = 1.0  # End chamfers

# Build the stepped shaft using cadkit.shaft
# steps=[(diameter, length), ...]
# Note: cadkit.shaft stacks steps upward from z=0.
# We define the sequence of steps.
shaft_geometry = ck.shaft(
    d=S1_D, 
    length=S1_L, 
    chamfer=CHAMFER, 
    steps=[
        (S2_D, S2_L), 
        (S3_D, S3_L)
    ]
)

# The keyway is on the middle step (S2).
# The middle step starts at z = S1_L and ends at z = S1_L + S2_L.
# We center the keyway length within the middle step.
keyway_z_start = S1_L + (S2_L - KEY_L) / 2
keyway_z_end = keyway_z_start + KEY_L

# Create the keyway cutter
# A box centered on the shaft axis, shifted up to the surface
keyway_cutter = (
    cq.Workplane("XY")
    .workplane(offset=S2_D / 2 - KEY_H)
    .box(KEY_W, KEY_L, KEY_H * 2, centered=(True, True, False))
    .translate((0, 0, keyway_z_start))
)

# Subtract keyway from shaft
final_shaft = shaft_geometry.cut(keyway_cutter)

# Add final chamfer to the top end (cadkit.shaft handles bottom and transitions, 
# but we ensure the top face is chamfered)
final_shaft = final_shaft.edges(">Z").chamfer(CHAMFER)

# Assembly
assy = cq.Assembly()
assy.add(final_shaft, name="stepped_shaft", color=cq.Color(0.7, 0.7, 0.75))

result = assy