import math
import cadquery as Cq
from nhf import Material
from nhf.handle import Handle

def trident_assembly(
        handle: Handle,
        handle_segment_length: float = 24*25.4,
        terminal_height=100):
    def segment():
        return handle.segment(handle_segment_length)

    terminal = (
        handle
        .one_side_connector(height=terminal_height)
        .faces(">Z")
        .hole(15, terminal_height + handle.insertion_length - 10)
    )
    mat_i = Material.PLASTIC_PLA
    mat_s = Material.ACRYLIC_BLACK
    assembly = (
        Cq.Assembly()
        .add(handle.insertion(), name="i0", color=mat_i.color)
        .constrain("i0", "Fixed")
        .add(segment(), name="s1", color=mat_s.color)
        .constrain("i0?rim", "s1?mate1", "Plane", param=0)
        .add(handle.insertion(), name="i1", color=mat_i.color)
        .add(handle.connector(), name="c1", color=mat_i.color)
        .add(handle.insertion(), name="i2", color=mat_i.color)
        .constrain("s1?mate2", "i1?rim", "Plane", param=0)
        .constrain("i1?mate", "c1?mate1", "Plane")
        .constrain("i2?mate", "c1?mate2", "Plane")
        .add(segment(), name="s2", color=mat_s.color)
        .constrain("i2?rim", "s2?mate1", "Plane", param=0)
        .add(handle.insertion(), name="i3", color=mat_i.color)
        .constrain("s2?mate2", "i3?rim", "Plane", param=0)
        .add(handle.one_side_connector(), name="head", color=mat_i.color)
        .constrain("i3?mate", "head?mate", "Plane")
        .add(terminal, name="terminal", color=mat_i.color)
        .constrain("i0?mate", "terminal?mate", "Plane")
    )
    return assembly.solve()