from build123d import *
from bd_warehouse.fastener import ClearanceHole, SocketHeadCapScrew, CounterSunkScrew
import math
import copy
import sys

VIEWER = "viewer" in sys.argv

if VIEWER:
    from ocp_vscode import *

spacing = 21.5
size = spacing / math.sqrt(3)

FILLET_AMOUNT = 3
CHAMFER_AMOUNT = 1
PLATE_INTERFACE = 0.5
THIN_WALL = 1.5

switch_sit = 2.2
switch_snap = 1.2
top_thick = switch_snap  # 2.0 # 2.2 # 1.3
top_space = switch_sit - top_thick
pcb_thick = 1.6
bottom_space = 1
bottom_thick = 3
bottom_expand = 2


def floor_to(n, step=1):
    return math.floor(n / step) * step


oled_thick = 6.05
cover_thick = oled_thick + 0.25 + 1.5  # - top_thick - top_space
cover_fastener = SocketHeadCapScrew("M3-0.5", 8)
bottom_fastener = CounterSunkScrew(
    "M2-0.4", floor_to(bottom_thick + bottom_space + pcb_thick + 2), "iso14581"
)

bottom_seal_compress = 0.05

feet_d = 11
height = math.sqrt(3) / 2

layout = """
K C C C C C C C C C K K
 K K K K K K K K K K K K
K K K K K K K K K K K K
 K K K K K K K K K K K K
K K K K K K K K K K K K
"""

milling = True


def shape2points(shape, filter="KC"):
    lines = shape.strip("\n").split("\n")

    for y, line in enumerate(lines):
        for x, char in enumerate(line):
            if char in filter:
                yield Vector(x * spacing / 2, y * -1.5 * size)


def g(x, y, ox=0, oy=0):
    return (x * spacing + ox, y * size + oy)


class LayoutLocations(Locations):
    def __init__(self, layout, filter="KC"):
        super().__init__(list(shape2points(layout, filter)))


def fillet_andor_chamfer(
    shape,
    filletable,
    chamferable,
    chamfer_amount=CHAMFER_AMOUNT,
    fillet_amount=FILLET_AMOUNT,
):
    if milling:
        fillet(filletable(shape.edges()) | Edge.is_interior, fillet_amount)
        fillet(
            filletable(shape.edges()).filter_by(Edge.is_interior, reverse=True),
            fillet_amount + chamfer_amount,
        )
        if chamferable:
            chamfer(chamferable(shape.edges()), chamfer_amount)
    else:
        fillet(filletable(shape.edges()), fillet_amount)


BOTTOM_SCREWS = [
    g(1, -1.5),
    g(1, -4.5),
    g(5.5, -3.5),
    g(10, -1.5),
    g(10, -4.5),
]

COVER_SCREWS = [g(3.5, 0.2), g(6.5, 0.2)]

with BuildPart() as cover:
    exporter = ExportDXF()

    bottom_plane = Plane.XY.offset(top_space + top_thick)
    top_plane = bottom_plane.offset(cover_thick)

    with BuildSketch(bottom_plane.offset(-PLATE_INTERFACE)):
        Polygon(g(0.5, 0), g(0.5, 1), g(9, 1), g(9, 0), align=Align.NONE)
        with LayoutLocations(layout, "C"):
            RegularPolygon(size, 6, rotation=90)
    extrude(amount=cover_thick + PLATE_INTERFACE)
    # extrude(target=top_plane, until=Until.NEXT)
    fillet_andor_chamfer(
        cover, lambda e: e | Axis.Z, lambda e: e >> Axis.Z
    )  # | Plane.XY)

    cover.interface = cover.faces() << Axis.Z

    # mounting bosses
    with BuildSketch(top_plane.offset(-THIN_WALL)) as sk:
        outerWire = (cover.faces() << Axis.Z)[0].wire()

        add(outerWire.offset_2d(-THIN_WALL))  # (THIN_WALL - CHAMFER_AMOUNT)))
        make_face()

        with Locations(COVER_SCREWS):
            Rectangle(9, 20, mode=Mode.SUBTRACT)

        fillet(sk.vertices(Select.LAST), FILLET_AMOUNT)

    extrude(dir=(0, 0, -1), until=Until.LAST, mode=Mode.SUBTRACT)

    # mounting holes
    with Locations(top_plane):
        with Locations(COVER_SCREWS):
            ClearanceHole(cover_fastener)

    # cutouts
    with Locations(top_plane):
        # oled
        with BuildSketch(bottom_plane) as cutout:
            with Locations(g(2, 0.185, ox=2.5)):
                RectangleRounded(25.5, 9, 2)
        exporter.add_shape(cutout.sketch)
        extrude(until=Until.LAST, mode=Mode.SUBTRACT)
        if milling:
            chamfer(cover.edges(Select.LAST) >> Axis.Z, 1)

        # mode indicator
        with BuildSketch(bottom_plane) as indicator:
            with Locations(g(1, 0.185)):
                Circle(5.85)
        exporter.add_shape(indicator.sketch)
        extrude(until=Until.LAST, mode=Mode.SUBTRACT)

    exporter.add_shape(section(section_by=Plane((cover.faces() << Axis.Z)[0])))

    if VIEWER:
        cover_assembly = Compound(
            children=[cover.part]
            + [
                copy.copy(cover_fastener).moved(l)
                for l in cover_fastener.hole_locations
            ]
        )
        show_object(cover_assembly, "cover")
    else:
        exporter.write("cover.dxf")
        export_step(cover.part, "cover.step")


with BuildPart() as top:
    with BuildSketch(Plane.XY.offset(top_space)) as outline:
        with LayoutLocations(layout, "KC"):
            RegularPolygon(size, 6, rotation=90)

        # fill top
        Polygon(g(0, 0), g(0, 1), g(9, 1), g(9, 0), align=Align.NONE)

        # fill top-right
        Polygon(g(10, 0), g(10, 1), g(11, 1), g(11, 0), align=Align.NONE)
        # thin top-right
        # Polygon(
        #     g(9.5, 0.75),
        #     g(9.5, 1),
        #     g(11.5, 1),
        #     g(11.5, 0.75),
        #     align=Align.NONE,
        #     mode=Mode.SUBTRACT,
        # )

        # thin bottom
        Polygon(
            g(0.5, -6.5),
            g(0.5, -7),
            g(10.5, -7),
            g(10.5, -6.5),
            align=Align.NONE,
            mode=Mode.SUBTRACT,
        )
        # fill bottom
        Polygon(
            g(0, -6),
            g(0, -6.75),
            g(11, -6.75),
            g(11, -6),
            align=Align.NONE,
        )

        fillet(outline.vertices(), FILLET_AMOUNT)

    extrude(amount=top_thick)

    # fillet_andor_chamfer(
    #     top,
    #     lambda e: e | Axis.Z,
    #     None,
    #     chamfer_amount=PLATE_INTERFACE,
    # )
    # fillet(top.edges(Select.NEW) | Axis.Z, FILLET_AMOUNT - 0.001)

    outerWire = (top.faces() >> Axis.Z)[0].outer_wire()
    chamfer(outerWire.edges(), PLATE_INTERFACE)

    with BuildPart(mode=Mode.SUBTRACT) as cover_seat:
        with BuildSketch(Plane.XY.offset(top_space + top_thick)):
            Polygon(g(0.5, 0), g(0.5, 2), g(9.5, 2), g(9.5, 0), align=Align.NONE)
            with LayoutLocations(layout, "C"):
                RegularPolygon(size, 6, rotation=90)
        extrude(amount=-PLATE_INTERFACE)

        fillet_andor_chamfer(
            cover_seat,
            lambda e: e | Axis.Z,
            None,
        )

    # cover hole
    outerWire = cover.interface[0].outer_wire().offset_2d(-2)
    extrude(Face(outerWire), until=Until.LAST, mode=Mode.SUBTRACT)

    with BuildSketch() as holes:
        with LayoutLocations(layout, "K"):
            RectangleRounded(13.4, 13.4, 0.5)
    extrude(until=Until.LAST, mode=Mode.SUBTRACT)

    with BuildSketch() as holes:
        with Locations(BOTTOM_SCREWS):
            Circle(3)
    extrude(until=Until.LAST, mode=Mode.SUBTRACT)
    # chamfer(top.edges(Select.LAST) << Axis.Z, 0.5) # no need bc gap

    outer_edge = (top.faces() << Axis.Z)[0].outer_wire()
    tube = extrude(
        Face(outer_edge),
        amount=-(top_space + pcb_thick + bottom_space + PLATE_INTERFACE),
        mode=Mode.PRIVATE,
    )
    offset(
        tube,
        amount=-(THIN_WALL - PLATE_INTERFACE),
        openings=tube.faces() | Axis.Z,
        mode=Mode.ADD,
    )
    top.interface = tube.faces() << Axis.Z

    if VIEWER:
        show_object(top, "top")
    else:
        export_step(top.part, "plate.top.step")


with BuildPart() as pcb:
    pcb.color = (0, 0.5, 0)

    with BuildSketch() as outline:
        with LayoutLocations(layout, "KC"):
            RegularPolygon(size, 6, rotation=90)

        # fill top
        Polygon(g(0, 0), g(0, 1), g(9, 1), g(9, 0), align=Align.NONE)

        # fill top-right
        Polygon(g(10, 0), g(10, 1), g(11, 1), g(11, 0), align=Align.NONE)

        # thin bottom
        Polygon(
            g(0.5, -6.5),
            g(0.5, -7),
            g(10.5, -7),
            g(10.5, -6.5),
            align=Align.NONE,
            mode=Mode.SUBTRACT,
        )
        # fill bottom
        Polygon(
            g(0, -6),
            g(0, -6.75),
            g(11, -6.75),
            g(11, -6),
            align=Align.NONE,
        )

        fillet(outline.vertices(), FILLET_AMOUNT)
        offset(outline.faces(), -(THIN_WALL + 0.5))

    extrude(amount=-pcb_thick)

    if VIEWER:
        show_object(pcb, "pcb")
    else:
        exporter = ExportDXF()
        exporter.add_shape(outline.sketch)
        exporter.write("pcb.dxf")


with BuildPart() as bottom:
    top_plane = Plane.XY.offset(-(pcb_thick + bottom_space))
    bottom_plane = top_plane.offset(-bottom_thick)

    with BuildSketch(bottom_plane) as outline:
        with LayoutLocations(layout, "KC"):
            RegularPolygon(size, 6, rotation=90)

        # top edge
        Polygon(g(0, 0), g(0, 1), g(11, 1), g(11, 0), align=Align.NONE)

        # bottom edge
        Polygon(
            g(0, -6),
            g(0, -7),
            g(11, -7),
            g(11, -6),
            align=Align.NONE,
        )

        # left edge
        Polygon(
            g(-0.5, 0),
            g(-0.5, -6),
            g(0, -6),
            g(0, 0),
            align=Align.NONE,
        )

        # right edge
        Polygon(
            g(11, 1),
            g(11, -5),
            g(12, -5),
            g(12, -0),
            align=Align.NONE,
        )

        offset(outline.faces(), bottom_expand, kind=Kind.INTERSECTION)

    extrude(amount=bottom_thick)
    fillet_andor_chamfer(bottom, lambda e: e | Axis.Z, lambda e: e | Plane.XY)

    extrude(top.interface, until=Until.FIRST, mode=Mode.SUBTRACT)

    # bosses
    with BuildSketch(bottom_plane):
        with Locations(BOTTOM_SCREWS):
            Circle(2.5)
    extrude(amount=bottom_thick + bottom_space - bottom_seal_compress)

    # mounting screws
    with Locations(BOTTOM_SCREWS):
        with Locations(bottom_plane.reverse()):
            ClearanceHole(bottom_fastener)

    if VIEWER:
        bottom_assembly = Compound(
            children=[bottom.part]
            + [
                copy.copy(bottom_fastener).moved(l)
                for l in bottom_fastener.hole_locations
            ]
        )
        show_object(bottom_assembly, "bottom")
    else:
        export_step(bottom.part, "plate.bottom.step")