bl_info = {
    "name": "Snap Connector Generator",
    "author": "Your Name",
    "version": (1, 3),
    "blender": (2, 80, 0),
    "location": "View3D > Add > Mesh > Snap Connector",
    "description": "Generates a configurable snap-fit connector pin and a socket",
    "category": "Add Mesh",
}

import bpy
import bmesh
from math import pi, cos, sin

import numpy as np

def generate_rounded_ridge_profile(taper_r, ridge_z, ridge_height, ridge_radius_factor, steps=10):
   
    z_start = ridge_z - ridge_height / 2
    z_end = ridge_z #+ ridge_height / 2
    z_vals = np.linspace(z_start, z_end, steps)

    profile = []
    for z in z_vals:
        t = (z - z_start) / (z_end - z_start)  # Normalized 0 → 1
        bulge = (1 - np.cos(t * pi)) / 2       # 0 → 1 → 0
        r = taper_r * (1 + (ridge_radius_factor - 1) * bulge)        
        profile.append((z, r))

    z_start = z_end
    z_end += ridge_height / 1.5
    z_vals = np.linspace(z_start, z_end, steps)
    for z in z_vals:
        t = (z - z_start) / (z_end - z_start)  # Normalized 0 → 1
        bulge = (1 - np.cos(t * pi)) / 2       # 0 → 1 → 0
        bulge = 1 - bulge
        r = taper_r * (1 + (ridge_radius_factor - 1) * bulge)        
        profile.append((z, r))

    return profile



def create_tapered_pin(radius, height, radius_ratio, ridge_height, ridge_radius_factor, taper_ratio, ridge_position_ratio, segments, slot_width, fill_top, fill_bottom, steps, add_second_slot):
    mesh = bpy.data.meshes.new("Tapered_Pin")
    obj = bpy.data.objects.new("Snap_Pin", mesh)
    bpy.context.collection.objects.link(obj)

    bm = bmesh.new()

    ridge_z = height * ridge_position_ratio

    # Profile steps (Z, Radius)
    steps_list = [(0.0, radius)]
    steps_list += generate_rounded_ridge_profile(
        taper_r=radius * taper_ratio,
        ridge_z=height * ridge_position_ratio,
        ridge_height=ridge_height,
        ridge_radius_factor=ridge_radius_factor,
        steps=steps  # Use the passed steps parameter
    )
    steps_list.append((height, radius * taper_ratio))

    ring_verts = []

    for i in range(segments):
        angle = 2 * pi * i / segments
        ring = []
        for z, r in steps_list:
            x = r * cos(angle)
            y = r * sin(angle)
            ring.append(bm.verts.new((x, y, z)))
        ring_verts.append(ring)

    bm.verts.ensure_lookup_table()

    # Side faces
    for i in range(segments):
        next_i = (i + 1) % segments
        for j in range(len(steps_list) - 1):
            v1 = ring_verts[i][j]
            v2 = ring_verts[next_i][j]
            v3 = ring_verts[next_i][j + 1]
            v4 = ring_verts[i][j + 1]
            bm.faces.new((v1, v2, v3, v4))

    # Fill bottom
    if fill_bottom:
        bottom_ring = [ring[0] for ring in ring_verts]
        bmesh.ops.contextual_create(bm, geom=bottom_ring)

    # Fill top
    if fill_top:
        top_ring = [ring[-1] for ring in ring_verts]
        bmesh.ops.contextual_create(bm, geom=top_ring)

    bm.to_mesh(mesh)
    bm.free()

    obj.location = (-radius * 2, 0, 0)

    # === SLOT: Vertical cut ===
    slot_z_center = obj.location.z + height * 0.8
    slot_height = height

    bpy.ops.mesh.primitive_cube_add(
        size=1,
        location=(obj.location.x, obj.location.y, slot_z_center)
    )
    slot = bpy.context.active_object
    slot.name = "Flex_Slot"
    slot.scale[0] = slot_width / 2
    slot.scale[1] = radius * 2.1
    slot.scale[2] = slot_height

    slot_bottom_z = slot_z_center - slot_height * 0.5

    bool_mod = obj.modifiers.new(name="Slot_Cut", type='BOOLEAN')
    bool_mod.operation = 'DIFFERENCE'
    bool_mod.object = slot
    bpy.context.view_layer.objects.active = obj
    bpy.ops.object.modifier_apply(modifier=bool_mod.name)
    bpy.data.objects.remove(slot, do_unlink=True)

    # === SLOT 2: Vertical cut, rotated 90 degrees ===
    if add_second_slot:
        bpy.ops.mesh.primitive_cube_add(
            size=1,
            location=(obj.location.x, obj.location.y, slot_z_center)
        )
        slot2 = bpy.context.active_object
        slot2.name = "Flex_Slot_2"
        slot2.scale[0] = slot_width / 2
        slot2.scale[1] = radius * 2.1
        slot2.scale[2] = slot_height
        # Rotate 90 degrees around Z-axis
        slot2.rotation_euler[2] = pi / 2  # 90 degrees in radians

        bool_mod2 = obj.modifiers.new(name="Slot_Cut_2", type='BOOLEAN')
        bool_mod2.operation = 'DIFFERENCE'
        bool_mod2.object = slot2
        bpy.context.view_layer.objects.active = obj
        bpy.ops.object.modifier_apply(modifier=bool_mod2.name)
        bpy.data.objects.remove(slot2, do_unlink=True)


    # === ROUND CUTOUT: Vertical cylinder to split top ===
    cut_radius = radius * radius_ratio
    cut_height = height * 1.25

    bpy.ops.mesh.primitive_cylinder_add(
        radius=cut_radius,
        depth=cut_height,
        vertices=segments,
        location=(obj.location.x, obj.location.y, slot_bottom_z + cut_height / 2),
        rotation=(0, 0, 0)
    )
    cutter = bpy.context.active_object
    cutter.name = "Flex_Cut"

    bool_cut = obj.modifiers.new(name="Vertical_Cylinder_Cut", type='BOOLEAN')
    bool_cut.operation = 'DIFFERENCE'
    bool_cut.object = cutter
    bpy.context.view_layer.objects.active = obj
    bpy.ops.object.modifier_apply(modifier=bool_cut.name)
    bpy.data.objects.remove(cutter, do_unlink=True)

    return obj

def create_socket(radius, height, taper_ratio, ridge_position_ratio, ridge_height, ridge_radius_factor, segments, tolerance, steps):
    mesh = bpy.data.meshes.new("Snap_Socket")
    obj = bpy.data.objects.new("Snap_Socket", mesh)
    bpy.context.collection.objects.link(obj)

    bm = bmesh.new()

    # Outer profile with tolerance
    base_r = radius + tolerance
    taper_r = base_r * taper_ratio
    ridge_z = height * ridge_position_ratio - tolerance  # ridge sits lower
    ridge_height_with_tol = ridge_height + tolerance * 2  # taller ridge cut

    # Profile steps (Z, Radius) using the same rounded ridge profile as the pin
    steps_list = [(0.0, base_r)]
    steps_list += generate_rounded_ridge_profile(
        taper_r=taper_r,
        ridge_z=ridge_z,
        ridge_height=ridge_height_with_tol,
        ridge_radius_factor=ridge_radius_factor + tolerance / taper_r,  # Adjust factor for tolerance
        steps=steps  # Use the passed steps parameter
    )
    steps_list.append((height, taper_r))

    ring_verts = []

    for i in range(segments):
        angle = 2 * pi * i / segments
        ring = []
        for z, r in steps_list:
            x = r * cos(angle)
            y = r * sin(angle)
            ring.append(bm.verts.new((x, y, z)))
        ring_verts.append(ring)

    bm.verts.ensure_lookup_table()

    # Side faces
    for i in range(segments):
        next_i = (i + 1) % segments
        for j in range(len(steps_list) - 1):
            v1 = ring_verts[i][j]
            v2 = ring_verts[next_i][j]
            v3 = ring_verts[next_i][j + 1]
            v4 = ring_verts[i][j + 1]
            bm.faces.new((v1, v2, v3, v4))

    # Fill bottom
    bottom_ring = [ring[0] for ring in ring_verts]
    bmesh.ops.contextual_create(bm, geom=bottom_ring)

    # Fill top
    top_ring = [ring[-1] for ring in ring_verts]
    bmesh.ops.contextual_create(bm, geom=top_ring)

    bm.to_mesh(mesh)
    bm.free()

    obj.location = (radius * 2, 0, 0)
    return obj




class MESH_OT_generate_snap_connector(bpy.types.Operator):
    bl_idname = "mesh.generate_snap_connector"
    bl_label = "Generate Snap Connector"
    bl_options = {'REGISTER', 'UNDO'}

    connector_radius: bpy.props.FloatProperty(
        name="Connector Radius",
        default=5.0,
        min=0.1,
        max=100.0
    )
    
    inner_cut_radius_ratio: bpy.props.FloatProperty(
        name="Inner Cut Radius Ratio",
        description="Ratio of the cylindrical cutout radius to the connector radius",
        default=0.5,
        min=0.05,
        max=0.95
    )
    
    taper_ratio: bpy.props.FloatProperty(
        name="Taper Ratio",
        default=0.95,
        min=0.5,
        max=1.0
    )
    ridge_position: bpy.props.FloatProperty(
        name="Ridge Position (0-1)",
        default=0.8,
        min=0.1,
        max=0.99
    )
    ridge_segments: bpy.props.IntProperty(
        name="Segments",
        default=32,
        min=5,
        max=128
    )
    slot_width: bpy.props.FloatProperty(
        name="Slot Width",
        default=4.0,
        min=0.1,
        max=10.0
    )
    
    add_second_slot: bpy.props.BoolProperty(
        name="Four slots",
        description="Add a second slot rotated 90 degrees around Z-axis",
        default=False
    )    
    
    fill_top: bpy.props.BoolProperty(
        name="Fill Top",
        default=True
    )
    fill_bottom: bpy.props.BoolProperty(
        name="Fill Bottom",
        default=False
    )

    tolerance: bpy.props.FloatProperty(
        name="Tolerance",
        description="Extra spacing between pin and socket",
        default=0.1,
        min=0.001,
        max=1.0
    )
    
    ridge_steps: bpy.props.IntProperty(
        name="Ridge Steps",
        description="Number of steps for the ridge profile smoothness",
        default=5,
        min=2,
        max=24
    )    

    def execute(self, context):
        r = self.connector_radius
        h = r * 2
        ridge_height = r * 0.2
        ridge_radius_factor = 1.1

        create_tapered_pin(
            radius=r,
            height=h,
            radius_ratio=self.inner_cut_radius_ratio,
            ridge_height=ridge_height,
            ridge_radius_factor=ridge_radius_factor,
            taper_ratio=self.taper_ratio,
            ridge_position_ratio=self.ridge_position,
            segments=self.ridge_segments,
            slot_width=self.slot_width,
            fill_top=self.fill_top,
            fill_bottom=self.fill_bottom,
            steps=self.ridge_steps,
            add_second_slot=self.add_second_slot 
        )

        create_socket(
            radius=r,
            height=h,
            taper_ratio=self.taper_ratio,
            ridge_position_ratio=self.ridge_position,
            ridge_height=r * 0.2,
            ridge_radius_factor=1.1,
            segments=self.ridge_segments,
            tolerance=self.tolerance,
            steps=self.ridge_steps  # Pass the new parameter
        )
        return {'FINISHED'}



def menu_func(self, context):
    self.layout.operator(MESH_OT_generate_snap_connector.bl_idname, icon='MESH_CYLINDER')


def register():
    bpy.utils.register_class(MESH_OT_generate_snap_connector)
    bpy.types.VIEW3D_MT_mesh_add.append(menu_func)


def unregister():
    bpy.utils.unregister_class(MESH_OT_generate_snap_connector)
    bpy.types.VIEW3D_MT_mesh_add.remove(menu_func)


if __name__ == "__main__":
    register()