from types import SimpleNamespace
from cadquery import Workplane
from math import sin, cos, pi

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

nothing = 0.00001
origin = (0, 0, 0)
XY = Workplane("XY")
center_xy = (True, True, False)
m3hole = 2.8 / 2

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 - nothing
        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_usbc_board(board, usbc, tol, usbc_overhang, hole_radius, hole_offset,
                     bolt_radius, bolt_head_radius, bolt_depth, usbc_height):
    def gen(tol):
        part = box(*(board + V3(tol, tol, tol) * 2).xyz()).translate((0, 0, -tol))
        part_usbc = build_usbc(**usbc.__dict__)
        if tol != 0:
            part_usbc = part_usbc.bounds
        part = part.union(part_usbc.rotate(origin, (0, 0, 1), 180)\
                          .translate((0, board.y / 2 + usbc_overhang, usbc_height)))
        for x in (-1, 1):
            if tol != 0:
                part = part.union(cylinder(bolt_head_radius, 10)
                        .translate((x * (board.x / 2 - hole_offset), board.y / 2 - hole_offset, board.z + tol)))
                part = part.union(cylinder(bolt_radius, board.z + 2 * tol + bolt_depth)
                        .translate((x * (board.x / 2 - hole_offset), board.y / 2 - hole_offset, -tol-bolt_depth)))
            else:
                part = part.cut(cylinder(hole_radius, board.z)
                        .translate((x * (board.x / 2 - hole_offset), board.y / 2 - hole_offset, 0)))
        return part.translate((0, -board.y/2, 0))
    return Ext(gen(0), gen(tol))

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_oled(outer, screen, tol, hole, pins, pinlen, ribbon):
    def gen(tol):
        slab = outer.z - screen.z
        lower = V3(outer.x + 2 * tol, outer.y + 2 * tol, slab + 2 * tol)
        upper = screen + V3(tol, tol, (10 if tol != 0 else 0)) * 2
        part2 = box(*upper.xyz()).translate((0, 0, lower.z))
        part = box(*lower.xyz()).union(part2).translate((0, 0, -tol))
        for y in range(-1, 2, 2):
            for x in range(-1, 2, 2):
                part = part.cut(cylinder(hole - tol, outer.z - tol)
                    .translate((x * 21 / 2, y * 22 / 2, 0)))
        part = part.union(box(*pins.xyz()).translate((0, outer.y / 2 - pins.y / 2, slab)))
        part = part.union(box(pins.x, pins.y, pinlen)\
                    .translate((0, outer.y / 2 - pins.y / 2, slab - pinlen)))
        if tol != 0:
            part = part.union(box(*ribbon.xyz()).translate((0, -outer.y / 2 + ribbon.y / 2 + tol, slab)))
        return part
    return Ext(gen(0), gen(tol))

def build_antenna(radius, height, top_offset, top_space, top_radius, tol, hex_height, hex_radius):
    def gen(tol):
        part = cylinder(radius + tol, height + 2 * tol).translate((0, 0, -tol))
        if tol != 0:
            part = part.union(cylinder(top_radius + 2 * tol, top_space + tol).translate((0, 0, top_offset)))
            part = part.union(cylinder(top_radius + tol + 4, height + tol - top_offset - top_space) \
                              .translate((0, 0, top_offset + top_space)))
            part = part.union(XY.polygon(6, hex_radius + tol * 2) \
                              .extrude(10 + hex_height + tol).translate((0, 0, -10)))
        else:
            part = part.union(cylinder(top_radius, height - top_offset).translate((0, 0, top_offset)))
            part = part.union(XY.polygon(6, hex_radius).extrude(hex_height))
        return part
    return Ext(gen(0), gen(tol))

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

    tower_radius = bolt.radius + wall.x + 1
    tower_height = inner.z

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

    antenna_wall_radius = lid.antenna.radius + 4
    antenna_rod = cylinder(antenna_wall_radius+wall.x, tower_height) \
        .cut(cylinder(antenna_wall_radius, tower_height))
    part = part.union(antenna_rod.translate((lid.antenna_offset + V3(0, -outer.y / 2, floor)).xyz()))
    part = part.union(antenna_rod.translate((lid.antenna_offset + V3(0, -outer.y / 2, floor) - V3(2 * lid.antenna_offset.x, 0, 0)).xyz()))
    part = part.union(box(abs(lid.antenna_offset.x)*2-2*nothing, (antenna_wall_radius+wall.x)*2, tower_height)\
                      .translate((0, lid.antenna_offset.y - outer.y / 2, floor)))

    board_tol = 2 * tol
    main_wall_dist = board_space.x + 2 * board_tol + wall.x
    part = part.union(box(wall.x, inner.y, tower_height).translate((-main_wall_dist/2, 0, floor)))
    part = part.union(box(wall.x, inner.y, tower_height).translate((main_wall_dist/2, 0, floor)))

    part = part.cut(box(main_wall_dist - wall.x, inner.y, tower_height).translate((0, 0, floor)))
    part = part.cut(box(abs(lid.antenna_offset.x)*2, antenna_wall_radius*2, tower_height)\
                      .translate((0, lid.antenna_offset.y - outer.y / 2, floor)))

    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()))

    # rod cut-out
    rod_pos = rod_shaft.pos + V3(0, 0, floor)
    base_rod_outer = rod.outer + 6
    part = part.union(cylinder(base_rod_outer, tower_height).translate(rod_pos.xyz()))

    fill_y = inner.y / 2 - rod_shaft.pos.y
    part = part.union(box(main_wall_dist, 2 * base_rod_outer - 4, tower_height).translate(rod_pos.xyz()).translate((0, 1, 0)))
    part_rod = cylinder(rod.outer + tol, tower_height).translate(rod_pos.xyz())
    part = part.cut(part_rod)
    part = part.union(box(main_wall_dist, fill_y, 3).translate(rod_pos.xyz())\
                      .translate((0, -rod_shaft.pos.y+inner.y/2-fill_y/2, 0)) \
                      .intersect(part_rod))
    #part = part.cut(cylinder(rod_hole.radius, 50).rotate(origin, (1, 0, 0), 90)
    #                .translate((rod_pos + V3(0, 0, rod_hole.height)).xyz()))
    part = part.cut(box(20, 25, 8).translate(rod_pos.xyz()).translate((0, -rod.outer - wall.x/2)))
    part = part.cut(box(8, 40, 5).translate(rod_pos.xyz()).translate((0, 10)))

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

    board_hole = V(outer = 3, inner = m3hole, height = 4)
    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 + board_tol, 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()))
            dist = board_hole.outer + board_tol
            if y == 1:
                part = part.union(box(dist, board_hole.outer * 2, board_hole.height-nothing)
                            .translate((pos.x + x * dist / 2, pos.y, pos.z)))
            else:
                dist2 = dist + board_hole.outer
                part = part.union(box(dist, dist2, board_hole.height)
                            .translate((pos.x + x * dist / 2,
                                        pos.y - dist + dist2 / 2, pos.z)))
                part = part.union(box(dist2, dist, board_hole.height)
                            .translate((pos.x + x * (dist - dist2 / 2),
                                        pos.y - dist / 2, pos.z)))
            part = part.cut(cylinder(board_hole.inner, board_hole.height).translate(pos.xyz()))

    part_usbc = build_usbc_board(**usbc_board.__dict__)
    usbc_board_pos = V3(0, outer.y / 2 - usbc_board.usbc_overhang, floor + 4)
    part = part.union(box(usbc_board.board.x, usbc_board.board.y, 4)\
                      .translate((usbc_board_pos - V3(0, usbc_board.board.y / 2, 4)).xyz()))
    part = part.cut(part_usbc.bounds.translate(usbc_board_pos.xyz()))
    if debug:
        part = part.union(part_usbc.translate(usbc_board_pos.xyz()))

    return part

def build_lid(base, rod, rod_shaft, wall, bolt, bolt_depth, tol, hole_spacing,
              oled, oled_offset, antenna, antenna_offset, rod_support, debug=False):
    outer = base
    inner = base - V21(wall + tol * 3)
    rod_wall = V(outer = rod.outer + tol * 2 + wall, inner = rod.outer + tol * 2, z = rod_shaft.height)

    part_bolt = build_bolt(bolt, bolt_depth, tol)

    def gen(tol):
        outer_padded = outer + V3(tol, tol, tol) * 4
        inner_padded = inner - V3(tol, tol, tol) * 4
        part = box(*outer_padded.xyz()).translate((0, 0, -tol*2))
        part = part.cut(box(*inner_padded.xyz()).translate((0, 0, -tol*2)))
        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 * 2
                head_height = bolt.head_height + wall + tol * 2
                offset_x = (outer.x - 2 * wall - hole_spacing.x) / 2 + tol * 2
                offset_y = (outer.y - 2 * wall - hole_spacing.y) / 2 + tol * 2
                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()))
                bolt_pos = (pos + V3(0, 0, -bolt_depth - bolt.head_height)).xyz()
                if tol != 0:
                    part = part.union(part_bolt.bounds.translate(bolt_pos))
                else:
                    part = part.cut(part_bolt.bounds.translate(bolt_pos))
                    part = part.cut(cylinder(bolt.radius + tol, bolt_depth + bolt.head_height).translate(bolt_pos))

        oled_pos = (oled_offset + V3(0, -outer_padded.y / 2, outer_padded.z - oled.outer.z)).xyz()
        part_oled = build_oled(**oled.__dict__)
        part = part.union(box(*oled.outer.xyz()).translate(oled_pos))
        part = part.cut(part_oled.bounds.translate(oled_pos))
        if debug:
            part = part.union(part_oled.translate(oled_pos))

        antenna_pos = (antenna_offset + V3(0, -outer_padded.y / 2, outer_padded.z - antenna.top_offset - antenna.top_space)).xyz()
        part_antenna = build_antenna(**antenna.__dict__)
        part = part.union(cylinder(antenna.top_radius + 2, antenna.height).translate(antenna_pos))
        part = part.cut(part_antenna.bounds.translate(antenna_pos))
        if debug:
            part = part.union(part_antenna.translate(antenna_pos))

        part = part.union(cylinder(rod.outer + tol * 2 + rod_support.thickness, rod_support.height + 2 * tol)
                        .translate((rod_shaft.pos + V3(0, 0, inner.z - rod_support.height - 2 * tol)).xyz()))

        return part

    part = gen(0)
    part_rod_wall = cylinder(rod_wall.outer + 5 + tol * 2, 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 + V3(0, 0, inner.z - rod_support.height - 2 * tol)).xyz()))

    part_bounds = gen(tol)

    return Ext(part, part_bounds)

def build_sleeve(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, 23, 0), height = 35)

    tol = 0.05
    rod = V(outer = 25.2 / 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 = m3hole)

    oled = V(outer = V3(25.5, 26, 4), screen = V3(24.9, 16.6, 1.6),
             tol = tol * 2, hole = 1, pins = V3(4 * 2.54, 2.54, 0.5),
             pinlen = 8, ribbon = V3(9, 5, 1))

    antenna = V(radius = 6.2/2, height = 13, top_offset = 7.6 + 0.5,
            top_radius = 7.6/2, top_space = 3, tol = tol * 2,
            hex_height = 2, hex_radius = 9)

    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 # let bolt go through

    usbc = V(width = 9, height = 3, depth = 7.3, tol = tol)
    usbc_board = V(usbc = usbc, board = V3(22, 13, 2), usbc_overhang = 1,
                   hole_radius = 3.2/2, hole_offset = 2.8, tol = tol,
                   bolt_radius = 3/2, bolt_head_radius = 3, bolt_depth = 4, usbc_height = 1.60)

    board_space = V3(36, 50, 10)
    rod_support = V(thickness = thin_wall, height = 9)
    #rod_hole = V(height = 12, radius = m3hole)

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

    sleeve_ctx = V(rod=rod, rod_shaft=rod_shaft, tol=tol, wall=thin_wall)
    #show_object(build_sleeve(**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, oled=oled, oled_offset=V3(0, 20, 0),
                antenna=antenna, antenna_offset=V3(-30, 20, 0), rod_support=rod_support)
    #show_object(build_lid(**lid_ctx.__dict__).translate((0, 0, 50)))

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

    #show_object(build_oled(**oled.__dict__).bounds)

    #show_object(build_antenna(**antenna.__dict__).bounds)

    #show_object(build_usbc_board(**usbc_board.__dict__).bounds)

main()