from dataclasses import dataclass
import math
import cadquery as Cq
from nhf import Item, Role
import nhf.utils

@dataclass(frozen=True)
class FlatHeadBolt(Item):
    diam_head: float
    height_head: float
    diam_thread: float
    height_thread: float

    @property
    def name(self) -> str:
        return f"Bolt M{int(self.diam_thread)} h{int(self.height_thread)}mm"

    @property
    def role(self) -> Role:
        return Role.CONNECTION


    def generate(self) -> Cq.Assembly:
        head = Cq.Solid.makeCylinder(
            radius=self.diam_head / 2,
            height=self.height_head,
        )
        rod = (
            Cq.Workplane('XY')
            .cylinder(
                radius=self.diam_thread/ 2,
                height=self.height_thread,
                centered=(True, True, False))
        )
        rod.faces("<Z").tag("tip")
        rod.faces(">Z").tag("root")
        rod = rod.union(head.located(Cq.Location((0, 0, self.height_thread))))
        return rod


@dataclass(frozen=True)
class ThreaddedKnob(Item):
    """
    A threaded rod with knob on one side
    """
    diam_thread: float
    height_thread: float
    diam_knob: float

    diam_neck: float
    height_neck: float
    height_knob: float

    @property
    def name(self) -> str:
        return f"Knob M{int(self.diam_thread)} h{int(self.height_thread)}mm"

    def generate(self) -> Cq.Assembly:
        knob = Cq.Solid.makeCylinder(
            radius=self.diam_knob / 2,
            height=self.height_knob,
        )
        neck = Cq.Solid.makeCylinder(
            radius=self.diam_neck / 2,
            height=self.height_neck,
        )
        thread = (
            Cq.Workplane('XY')
            .cylinder(
                radius=self.diam_thread / 2,
                height=self.height_thread,
                centered=(True, True, False))
        )
        thread.faces("<Z").tag("tip")
        thread.faces(">Z").tag("root")

        return (
            Cq.Assembly()
            .addS(thread, name="thread", role=Role.CONNECTION)
            .addS(neck, name="neck", role=Role.HANDLE,
                  loc=Cq.Location((0, 0, self.height_thread)))
            .addS(knob, name="knob", role=Role.HANDLE,
                  loc=Cq.Location((0, 0, self.height_thread + self.height_neck)))
        )


@dataclass(frozen=True)
class HexNut(Item):
    diam_thread: float
    pitch: float

    thickness: float
    width: float

    def __post_init__(self):
        assert self.width > self.diam_thread

    @property
    def name(self):
        return f"HexNut M{int(self.diam_thread)}-{self.pitch}"

    @property
    def role(self) -> Role:
        return Role.CONNECTION

    def generate(self) -> Cq.Workplane:
        r = self.width / math.sqrt(3)
        result = (
            Cq.Workplane("XY")
            .sketch()
            .regularPolygon(r=r, n=6)
            .circle(r=self.diam_thread/2, mode='s')
            .finalize()
            .extrude(self.thickness)
        )
        result.faces("<Z").tag("bot")
        result.faces(">Z").tag("top")
        result.copyWorkplane(Cq.Workplane('XY')).tagPlane("dirX", direction="+X")
        return result