import bpy
import bmesh
import math

def add_boolean(target, cutter, operation='DIFFERENCE', solver='EXACT'):
    mod = target.modifiers.new(name="Bool", type='BOOLEAN')
    mod.object = cutter
    mod.operation = operation
    mod.solver = solver
    bpy.context.view_layer.objects.active = target
    bpy.ops.object.modifier_apply(modifier=mod.name)
    bpy.data.objects.remove(cutter, do_unlink=True)

def generate_blocks():
    # Clear existing mesh objects
    bpy.ops.object.select_all(action='DESELECT')
    for obj in list(bpy.context.scene.objects):
        if obj.type == 'MESH':
            bpy.data.objects.remove(obj, do_unlink=True)

    block_size = 20.0
    neck_r = 3.0       # 6mm diameter neck for durability
    neck_l = 4.0       # 4mm neck length
    ball_r = 3.6       # 7.2mm ball (0.6mm radial undercut — gentle, easy-snap fit for toy kits)
    tol = 0.25         # 0.25mm clearance for smooth FDM 3D printing fit
    slit_w = 1.2       # 1.2mm compression slot width
    relief_r = 1.0     # Stress-relief hole radius at base of slit

    # ==========================================
    # BLOCK 1: PLUG (At Y = -15)
    # ==========================================
    bpy.ops.mesh.primitive_cube_add(size=block_size, location=(0, -15, block_size/2))
    b1 = bpy.context.active_object
    b1.name = "Block_Plug"

    mod_bev1 = b1.modifiers.new(name="Bevel", type='BEVEL')
    mod_bev1.width = 0.6
    mod_bev1.segments = 2
    bpy.context.view_layer.objects.active = b1
    bpy.ops.object.modifier_apply(modifier=mod_bev1.name)

    # 1. Base fillet cone at neck root to prevent layer shear
    bpy.ops.mesh.primitive_cone_add(
        vertices=48,
        radius1=neck_r + 1.0,
        radius2=neck_r,
        depth=1.5,
        location=(0, -15, block_size + 0.45)
    )
    neck_fillet = bpy.context.active_object
    add_boolean(b1, neck_fillet, operation='UNION')

    # 2. Neck cylinder
    bpy.ops.mesh.primitive_cylinder_add(
        vertices=48,
        radius=neck_r,
        depth=neck_l + 2.0,
        location=(0, -15, block_size + (neck_l - 2.0)/2)
    )
    neck = bpy.context.active_object
    add_boolean(b1, neck, operation='UNION')

    # 3. Ball at Z = block_size + neck_l
    ball_z = block_size + neck_l
    bpy.ops.mesh.primitive_uv_sphere_add(
        segments=48,
        ring_count=24,
        radius=ball_r,
        location=(0, -15, ball_z)
    )
    ball = bpy.context.active_object
    add_boolean(b1, ball, operation='UNION')

    # 4. Flatten top of ball by 0.6mm for easier entry
    bpy.ops.mesh.primitive_cube_add(size=10.0, location=(0, -15, ball_z + ball_r + 5.0 - 0.6))
    top_cut = bpy.context.active_object
    add_boolean(b1, top_cut, operation='DIFFERENCE')

    # 5. Cut compression slit along X axis
    slit_bottom_z = block_size + 1.2
    slit_top_z = ball_z + ball_r + 2.0
    slit_h = slit_top_z - slit_bottom_z
    bpy.ops.mesh.primitive_cube_add(
        size=1.0,
        location=(0, -15, slit_bottom_z + slit_h/2)
    )
    slit = bpy.context.active_object
    slit.scale = (ball_r * 3.0, slit_w, slit_h)
    bpy.ops.object.transform_apply(location=False, rotation=False, scale=True)
    add_boolean(b1, slit, operation='DIFFERENCE')

    # 6. Cylindrical stress-relief hole at base of slit (prevents fatigue cracking)
    bpy.ops.mesh.primitive_cylinder_add(
        vertices=32,
        radius=relief_r,
        depth=ball_r * 3.0,
        location=(0, -15, slit_bottom_z),
        rotation=(0, math.pi/2, 0)
    )
    relief = bpy.context.active_object
    add_boolean(b1, relief, operation='DIFFERENCE')


    # ==========================================
    # BLOCK 2: SOCKET (At Y = 15)
    # ==========================================
    bpy.ops.mesh.primitive_cube_add(size=block_size, location=(0, 15, block_size/2))
    b2 = bpy.context.active_object
    b2.name = "Block_Socket"

    mod_bev2 = b2.modifiers.new(name="Bevel", type='BEVEL')
    mod_bev2.width = 0.6
    mod_bev2.segments = 2
    bpy.context.view_layer.objects.active = b2
    bpy.ops.object.modifier_apply(modifier=mod_bev2.name)

    # 1. Socket Neck Channel
    bpy.ops.mesh.primitive_cylinder_add(
        vertices=48,
        radius=neck_r + tol,
        depth=neck_l + 2.0,
        location=(0, 15, block_size - neck_l/2 + 1.0)
    )
    c_neck = bpy.context.active_object
    add_boolean(b2, c_neck, operation='DIFFERENCE')

    # 2. Socket Spherical Cavity
    socket_ball_z = block_size - neck_l
    bpy.ops.mesh.primitive_uv_sphere_add(
        segments=48,
        ring_count=24,
        radius=ball_r + tol,
        location=(0, 15, socket_ball_z)
    )
    c_ball = bpy.context.active_object
    add_boolean(b2, c_ball, operation='DIFFERENCE')

    # 3. Entrance Lead-In Chamfer Cone
    bpy.ops.mesh.primitive_cone_add(
        vertices=48,
        radius1=neck_r + tol - 0.2,
        radius2=neck_r + tol + 1.4,
        depth=2.0,
        location=(0, 15, block_size - 0.5)
    )
    c_cone = bpy.context.active_object
    add_boolean(b2, c_cone, operation='DIFFERENCE')

    # 4. Relief Slot across the top of the socket
    bpy.ops.mesh.primitive_cube_add(
        size=1.0,
        location=(0, 15, block_size - neck_l/2)
    )
    s_slot = bpy.context.active_object
    s_slot.scale = ((ball_r + 2.5)*2, 1.0, neck_l + 1.0)
    bpy.ops.object.transform_apply(location=False, rotation=False, scale=True)
    add_boolean(b2, s_slot, operation='DIFFERENCE')

    for obj in [b1, b2]:
        bm = bmesh.new()
        bm.from_mesh(obj.data)
        bmesh.ops.remove_doubles(bm, verts=bm.verts, dist=0.001)
        bm.to_mesh(obj.data)
        bm.free()

if __name__ == "__main__":
    generate_blocks()