import FreeCAD as App
import FreeCADGui as Gui
import Part
import math

#
# ========= USER PARAMETERS =========
#
shaft_diameter_in    = 3.0      # must match the shaft macro
num_notches          = 15       # must match the shaft macro
notch_radius_in      = 0.15     # must match the shaft macro

coupler_length_in    = 3.0      # axial length of coupler
wall_thickness_in    = 0.75     # radial wall thickness
#
# ===================================
#

INCH_TO_MM = 25.4

shaft_radius_mm   = (shaft_diameter_in / 2.0) * INCH_TO_MM
notch_radius_mm   = notch_radius_in * INCH_TO_MM
coupler_length_mm = coupler_length_in * INCH_TO_MM
wall_thickness_mm = wall_thickness_in * INCH_TO_MM

outer_radius_mm = shaft_radius_mm + wall_thickness_mm

doc = App.ActiveDocument
if doc is None:
    doc = App.newDocument("SplineCoupler")
else:
    doc = App.ActiveDocument

# ---------------- Coupler outer body ----------------
coupler_outer = doc.addObject("Part::Cylinder", "CouplerOuter")
coupler_outer.Radius = outer_radius_mm
coupler_outer.Height = coupler_length_mm
coupler_outer.Placement = App.Placement(
    App.Vector(0, 0, 0),
    App.Rotation(App.Vector(0, 0, 1), 0)
)

doc.recompute()

# ---------------- Internal "negative" of shaft profile ----------------

# Base inner cylinder (same OD as shaft)
inner_proto = doc.addObject("Part::Cylinder", "InnerSplineProto")
inner_proto.Radius = shaft_radius_mm
inner_proto.Height = coupler_length_mm + 2.0  # extend a bit to guarantee full cut
inner_proto.Placement = App.Placement(
    App.Vector(0, 0, -1.0),
    App.Rotation(App.Vector(0, 0, 1), 0)
)

doc.recompute()

# Create the notch cylinders (same pattern as shaft macro)
notch_objs = []

# NOTE: This matches your shaft macro line:
# notch_center_radius = shaft_radius_mm - notch_radius_mm * 0.5
notch_center_radius = shaft_radius_mm - notch_radius_mm * 0.5

for i in range(num_notches):
    angle = 2.0 * math.pi * i / float(num_notches)
    x = notch_center_radius * math.cos(angle)
    y = notch_center_radius * math.sin(angle)

    notch = doc.addObject("Part::Cylinder", f"CouplerNotch_{i+1}")
    notch.Radius = notch_radius_mm
    notch.Height = coupler_length_mm + 2.0
    notch.Placement = App.Placement(
        App.Vector(x, y, -1.0),
        App.Rotation(App.Vector(0, 0, 1), 0)
    )
    notch_objs.append(notch)

doc.recompute()

# Combine notches
notch_compound = doc.addObject("Part::Compound", "CouplerNotchCompound")
notch_compound.Links = notch_objs
doc.recompute()

# Cut notches into inner spline proto (to get internal shaft-like profile)
inner_spline = doc.addObject("Part::Cut", "CouplerInnerSpline")
inner_spline.Base = inner_proto
inner_spline.Tool = notch_compound
doc.recompute()

# ---------------- Final coupler: outer minus inner spline ----------------
coupler = doc.addObject("Part::Cut", "SplineCoupler")
coupler.Base = coupler_outer
coupler.Tool = inner_spline
doc.recompute()

# Hide construction geometry
inner_proto.ViewObject.Visibility = False
for n in notch_objs:
    n.ViewObject.Visibility = False
notch_compound.ViewObject.Visibility = False
coupler_outer.ViewObject.Visibility = False

coupler.ViewObject.Visibility = True

if Gui.ActiveDocument:
    Gui.ActiveDocument.ActiveView.fitAll()

App.Console.PrintMessage(
    "Done: coupler created.\n"
    "  Inner profile: matches 3\" shaft with %d notches (r=%.3f in)\n"
    "  Coupler length: %.3f in\n"
    "  Wall thickness: %.3f in\n"
    "  Outer diameter: %.3f in\n"
    % (
        num_notches,
        notch_radius_in,
        coupler_length_in,
        wall_thickness_in,
        shaft_diameter_in + 2.0 * wall_thickness_in
    )
)