from types import SimpleNamespace
from cadquery import Workplane

if "show_object" not in globals():
    show_object = lambda v: ()

nothing = 0.0001
origin = (0, 0, 0)
XY = Workplane("XY")
center_xy = (True, True, False)

class V(SimpleNamespace):
    def __add__(self, rhs):
        assert(set(self.__dict__.keys()) <= set(rhs.__dict__.keys()))
        items = [(k, v + rhs.__dict__[k]) for k,v in self.__dict__.items()]
        return V(**dict(items))

    def __sub__(self, rhs):
        assert(set(self.__dict__.keys()) <= set(rhs.__dict__.keys()))
        items = [(k, v - rhs.__dict__[k]) for k,v in self.__dict__.items()]
        return V(**dict(items))

    def __mul__(self, rhs):
        assert(isinstance(rhs, int))
        items = [(k, v *rhs) for k,v in self.__dict__.items()]
        return V(**dict(items))

    def xyz(self):
        return (self.x, self.y, self.z)

def Ext(part, bounds):
    setattr(part, "bounds", bounds)
    return part

def V3(x, y, z):
    return V(x=x, y=y, z=z)

def V21(v):
    return V3(2 * v, 2 * v, v)

def box(*args):
    return XY.box(*args, centered=center_xy)

def cylinder(r, h):
    return XY.circle(r).extrude(h)

def build_usbc(width, height, depth, tol):
    def shape(width, height, depth):
        radius = (height + 2 * tol) / 2
        offset = width / 2 - radius
        part = box(width - radius * 2, height, depth)
        part = part.union(cylinder(radius, depth).translate((-offset, 0, 0)))
        part = part.union(cylinder(radius, depth).translate((offset, 0, 0)))
        return part
    def gen(tol):
        dims = (width + 2 * tol, height + 2 * tol, depth + 2 * tol)
        part = shape(*dims).translate((0, 0, -tol))
        if tol != 0:
            part = part.union(shape(*dims[:2], 20).translate((0, 0, -20)))
        return part
    part = gen(0).translate((0, -height / 2, 0)).rotate(origin, (1, 0, 0), -90)
    bounds = gen(tol).translate((0, -height / 2, 0)).rotate(origin, (1, 0, 0), -90)
    return Ext(part, bounds)

def build_bolt(bolt, shaft, tol):
    height = shaft + bolt.head_height

    def gen(tol):
        part = cylinder(bolt.head_radius + tol, bolt.head_height + 2 * tol)\
            .translate((0, 0, height + 2 * tol - bolt.head_height - 2 * tol))
        if tol != 0:
            part = part.union(cylinder(bolt.cut_radius + tol, height + 2 * tol))
            part = part.translate((0, 0, -tol))
        else:
            part = part.union(cylinder(bolt.radius, height))
        return part

    return Ext(gen(0), gen(tol))

def build_rod(rod, tol):
    part = XY.circle(rod.outer).circle(rod.inner).extrude(rod.len)

    bounds = cylinder(rod.outer + tol, rod.len + 2 * tol).translate((0, 0, -tol))

    return Ext(part, bounds)

def build_base(base, rod, rod_shaft, usbc, bolt, wall, tol, lid):
    outer = base
    inner = outer - V3(2 * wall.x, 2 * wall.y, wall.z) + V21(tol) 

    floor = wall.z - tol
    part = box(*base.xyz())
    part = part.cut(box(*inner.xyz()).translate((0, 0, floor)))

    tower_radius = bolt.radius + wall.x + 1
    tower_height = inner.z
    for y in range(-1, 2, 2):
        for x in range(-1, 2, 2):
            pos = V3(x * lid.hole_spacing.x / 2, y * lid.hole_spacing.y / 2, floor)
            offset_x = (inner.x - lid.hole_spacing.x) / 2
            offset_y = (inner.y - lid.hole_spacing.y) / 2
            part = part.union(box(offset_x, offset_y + tower_radius, tower_height) \
                              .translate((pos + V3(offset_x*x/2, y*(offset_y-(offset_y+tower_radius)/2), 0)).xyz()))
            part = part.union(box(offset_x + tower_radius, offset_y, tower_height) \
                                  .translate((pos + V3(x*(offset_x-(offset_x+tower_radius)/2), offset_y*y/2, 0)).xyz()))
            part = part.union(cylinder(tower_radius, tower_height).translate(pos.xyz()))

    part = part.union(box(wall.x, inner.y, tower_height).translate((-20, 0, floor)))
    part = part.union(box(wall.x, inner.y, tower_height).translate((20, 0, floor)))

    # rod cut-out
    rod_pos = (rod_shaft.pos + V3(0, 0, floor)).xyz()
    base_rod_outer = rod.outer + 6
    part = part.union(cylinder(base_rod_outer, tower_height).translate(rod_pos))
    fill_y = inner.y / 2 - rod_shaft.pos.y
    part = part.union(box(40, fill_y, tower_height).translate(rod_pos)\
                      .translate((0, -rod_shaft.pos.y+inner.y/2-fill_y/2, 0)))
    part = part.cut(cylinder(rod.outer + tol, tower_height).translate(rod_pos))
    part = part.union(box(40, fill_y, 3).translate(rod_pos)\
                      .translate((0, -rod_shaft.pos.y+inner.y/2-fill_y/2, 0)))
    part = part.cut(box(20, 25, 8).translate(rod_pos).translate((0, -rod.outer - wall.x/2)))

    part = part.cut(build_lid(**lid.__dict__).bounds.translate((0, 0, base.z - lid.base.z)))

    board_hole = V(outer = 3, inner = 2.9 / 2 + tol, height = 2)
    board_space = V3(40 - wall.x, 60, 0)
    print("Board Space", board_space)
    board_space_inner = board_space - V3(2 * board_hole.outer, 2 * board_hole.outer, 0)
    board_center = V3(0, -inner.y / 2 + board_space.y / 2, floor)
    for y in (-1, 1):
        for x in (-1, 1):
            pos = board_center + V3(x * board_space_inner.x / 2, y * board_space_inner.y / 2, 0)
            part = part.union(cylinder(board_hole.outer, board_hole.height).translate(pos.xyz()))
            pos2 = pos + V3(x * board_hole.outer / 2, y * board_hole.outer / 2, 0)
            part = part.union(box(board_hole.outer, board_hole.outer, board_hole.height-nothing).translate(pos2.xyz()))
            part = part.cut(cylinder(board_hole.inner, board_hole.height).translate(pos.xyz()))

    part_usbc = build_usbc(**usbc.__dict__, tol=tol)
    part = part.cut(part_usbc.bounds.translate((board_center + V3(0, -board_space.y / 2 - 1.5, 4)).xyz()))

    return part

def build_lid(base, rod, rod_shaft, wall, bolt, bolt_depth, tol, hole_spacing):
    outer = base
    inner = base - V21(wall + tol)
    rod_wall = V(outer = rod.outer + tol + wall, inner = rod.outer + tol, z = rod_shaft.height)

    part_bolt = build_bolt(bolt, bolt_depth, tol)

    def gen(tol):
        outer_padded = outer + V3(tol, tol, tol) * 2
        inner_padded = inner - V3(tol, tol, tol) * 2
        part = box(*outer_padded.xyz()).translate((0, 0, -tol))
        part = part.cut(box(*inner_padded.xyz()).translate((0, 0, -tol)))
        for y in range(-1, 2, 2):
            for x in range(-1, 2, 2):
                pos = V3(x * hole_spacing.x / 2, y * hole_spacing.y / 2, outer.z)
                head_radius = bolt.head_radius + wall + tol
                head_height = bolt.head_height + wall + tol
                offset_x = (outer.x - 2 * wall - hole_spacing.x) / 2 + tol
                offset_y = (outer.y - 2 * wall - hole_spacing.y) / 2 + tol
                part = part.union(cylinder(head_radius, head_height) \
                                  .translate((pos + V3(0, 0, -head_height)).xyz()))
                part = part.union(box(offset_x, offset_y + head_radius, head_height) \
                                  .translate((pos + V3(offset_x*x/2, y*(offset_y-(offset_y+head_radius)/2), -head_height)).xyz()))
                part = part.union(box(offset_x + head_radius, offset_y, head_height) \
                                  .translate((pos + V3(x*(offset_x-(offset_x+head_radius)/2), offset_y*y/2, -head_height)).xyz()))
                if tol != 0:
                    part = part.union(part_bolt.bounds.translate((pos + V3(0, 0, -bolt_depth - bolt.head_height)).xyz()))
                else:
                    part = part.cut(part_bolt.bounds.translate((pos + V3(0, 0, -bolt_depth - bolt.head_height)).xyz()))
        return part

    part = gen(0)
    part_rod_wall = cylinder(rod_wall.outer + 5 + tol, 10)
    part = part.cut(part_rod_wall.translate((rod_shaft.pos + V3(0, 0, outer.z - 1)).xyz()))
    part = part.cut(build_rod(rod, tol).bounds.translate(rod_shaft.pos.xyz()))

    part_bounds = gen(tol)

    return Ext(part, part_bounds)

def build_sheath(rod, rod_shaft, tol, wall):
    rod_wall = V(outer = rod.outer + tol + wall, inner = rod.outer + tol, z = rod_shaft.height)
    part = cylinder(rod_wall.outer + 5, 2)
    part = part.union(cylinder(rod_wall.outer, rod_wall.z))
    part = part.cut(build_rod(rod, tol).bounds)
    return part

def main():
    # TEST PRINT
    base = V3(80, 80, 15)
    rod_shaft = V(pos = V3(0, 0, 0), height = 20)

    base = V3(120, 120, 30)
    rod_shaft = V(pos = V3(0, 30, 0), height = 35)
    tol = 0.05
    rod = V(outer = 25.3 / 2, inner = 21 / 2, len = 500)
    thin_wall = 2.4
    thick_wall = 4
    bolt = V(radius = 3 / 2, head_radius = 5.4 / 2, head_height = 3, cut_radius = 2.9 / 2)

    lid = V3(base.x, base.y, 10)
    lid_hole_spacing = V3(lid.x - 15, lid.y - 15, 0)
    insert = V(radius=5 / 2, height=5)
    base_wall = V3(thick_wall, thick_wall, thin_wall)

    bolt_depth = base.z - bolt.head_height # - base_wall.z # let bolt go through

    usbc = V(width = 9, height = 3, depth = 7.3)

    #show_object(build_rod(rod).translate((0, 0, 10)))

    sheath_ctx = V(rod=rod, rod_shaft=rod_shaft, tol=tol, wall=thin_wall)
    show_object(build_sheath(**sheath_ctx.__dict__).translate((-100, 0, 0)))

    lid_ctx = V(base=lid, rod=rod, rod_shaft=rod_shaft,
                wall=thin_wall, bolt=bolt, bolt_depth=bolt_depth, tol=tol,
                hole_spacing=lid_hole_spacing)
    show_object(build_lid(**lid_ctx.__dict__).translate((150, 0, 0)))

    base_ctx = V(base=base, rod=rod, rod_shaft=rod_shaft, bolt=bolt,
                 wall=base_wall, tol=tol, usbc=usbc)
    show_object(build_base(**base_ctx.__dict__, lid=lid_ctx))

main()