"""
This schematics file contains all designs related to tool handles
"""
from dataclasses import dataclass, field
from typing import Union, Optional
import cadquery as Cq
import nhf.parts.metric_threads as metric_threads
import nhf.utils

class Mount:
    """
    Describes the internal connection between two cylinders
    """

    def diam_insertion_internal(self) -> float:
        """
        Diameter of the internal cavity in the insertion
        """
    def diam_connector_external(self) -> float:
        """
        Diameter of the external size of the connector
        """
    def external_thread(self, length: float) -> Cq.Shape:
        """
        Generates the external connector
        """
    def internal_thread(self, length: float) -> Cq.Shape:
        """
        Generates the internal connector
        """

@dataclass
class ThreadedMount(Mount):

    pitch: float = 3

    # Major diameter of the internal threads, following ISO metric screw thread
    # standard. This determines the wall thickness of the insertion.
    diam_threading: float = 27

    def diam_insertion_internal(self) -> float:
        r = metric_threads.metric_thread_major_radius(
            self.diam_threading,
            self.pitch,
            internal=True)
        return r * 2
    def diam_connector_external(self) -> float:
        r = metric_threads.metric_thread_minor_radius(
            self.diam_threading,
            self.pitch)
        return r * 2

    def external_thread(self, length: float):
        return metric_threads.external_metric_thread(
            self.diam_threading,
            self.pitch,
            length,
            top_lead_in=True)
    def internal_thread(self, length: float):
        return metric_threads.internal_metric_thread(
            self.diam_threading,
            self.pitch,
            length)

@dataclass
class BayonetMount(Mount):
    """
    Bayonet type connection
    """
    diam_outer: float = 30
    diam_inner: float = 27

    # Angular span (in degrees) of the slider
    pin_span: float = 15
    pin_height: float = 5
    # Wall at the bottom of the slot
    gap: float = 3

    # Angular span (in degrees) of the slot
    slot_span: float = 90

    # Number of pins equally distributed along a circle
    n_pin: int = 2


    def __post_init__(self):
        assert self.diam_outer > self.diam_inner
        assert self.n_pin * self.slot_span < 360
        assert self.slot_span > self.pin_span

    def diam_insertion_internal(self) -> float:
        return self.diam_outer
    def diam_connector_external(self) -> float:
        return self.diam_inner

    def external_thread(self, length: float):
        assert length > self.pin_height + self.gap
        pin = (
            Cq.Workplane('XY')
            .cylinder(
                height=self.pin_height,
                radius=self.diam_outer / 2,
                angle=self.pin_span,
                centered=(True, True, False))
            .copyWorkplane(Cq.Workplane('XY'))
            .cylinder(
                height=self.pin_height,
                radius=self.diam_inner / 2,
                centered=(True, True, False),
                combine="cut")
            .val()
        )
        result = (
            Cq.Workplane('XY')
            .workplane(offset=self.gap)
            .polarArray(radius=0, startAngle=0, angle=360, count=self.n_pin)
            .eachpoint(lambda loc: pin.located(loc), combine='a')
            .clean()
        )
        return result
    def internal_thread(self, length: float):
        assert length > self.pin_height + self.gap
        slot = (
            Cq.Workplane('XY')
            .cylinder(
                height=length - self.gap,
                radius=self.diam_outer / 2,
                angle=self.pin_span,
                centered=(True, True, False)
            )
            .copyWorkplane(Cq.Workplane('XY'))
            .cylinder(
                height=self.pin_height,
                radius=self.diam_outer / 2,
                angle=self.slot_span,
                centered=(True, True, False)
            )
            .val()
        )
        result = (
            Cq.Workplane('XY')
            .cylinder(
                height=length,
                radius=self.diam_outer / 2,
                centered=(True, True, False),
            )
            .copyWorkplane(Cq.Workplane('XY'))
            .workplane(offset=self.gap)
            .polarArray(radius=0, startAngle=self.slot_span, angle=360, count=self.n_pin)
            .cutEach(lambda loc: slot.located(loc))
            .clean()
            .copyWorkplane(Cq.Workplane('XY'))
            .cylinder(
                height=length,
                radius=self.diam_inner / 2,
                centered=(True, True, False),
                combine="cut"
            )
        )
        return result


@dataclass
class Handle:
    """
    Characteristic of a tool handle

    This assumes the handle segment material does not have threads. Each segment
    attaches to two insertions, which have threads on the inside. A connector
    has threads on the outside and mounts two insertions.

    Note that all the radial sizes are diameters (in mm).
    """

    # Outer and inner radii for the handle usually come in standard sizes
    diam: float = 38
    diam_inner: float = 33

    mount: Optional[Mount] = field(default_factory=lambda: ThreadedMount())

    # Internal cavity diameter. This determines the wall thickness of the connector
    diam_connector_internal: float = 18.0

    # If set to true, do not generate the connections
    simplify_geometry: bool = True

    # Length for the rim on the female connector
    rim_length: float = 5

    insertion_length: float = 30

    # Amount by which the connector goes into the segment
    connector_length: float = 60

    def __post_init__(self):
        assert self.diam > self.diam_inner, "Material thickness cannot be <= 0"
        if self.mount:
            assert self.diam_inner > self.mount.diam_insertion_internal(), "Threading radius is too big"
            assert self.mount.diam_insertion_internal() >= self.mount.diam_connector_external()
            assert self.mount.diam_connector_external() > self.diam_connector_internal, "Internal diameter is too large"
        assert self.insertion_length > self.rim_length

    def segment(self, length: float):
        result = (
            Cq.Workplane()
            .cylinder(
                radius=self.diam / 2,
                height=length)
            .faces(">Z")
            .hole(self.diam_inner)
        )
        result.faces("<Z").tag("mate1")
        result.faces(">Z").tag("mate2")
        return result

    def insertion(self, holes=[]):
        """
        This type of mount is used to connect two handlebar pieces. Each handlebar
        piece is a tube which cannot be machined, so the mount connects to the
        handle by glue.

        Tags:
        * lip: Co-planar Mates to the rod
        * mate: Mates to the connector

        WARNING: A tolerance lower than the defualt (maybe 5e-4) is required for
        STL export.

        Set `holes` to the heights for drilling holes into the model for resin
        to flow out.
        """
        result = (
            Cq.Workplane('XY')
            .cylinder(
                radius=self.diam_inner / 2,
                height=self.insertion_length - self.rim_length,
                centered=[True, True, False])
        )
        result.faces(">Z").tag("rim")
        if self.rim_length > 0:
            result = (
                result.faces(">Z")
                .workplane()
                .circle(self.diam / 2)
                .extrude(self.rim_length)
                .faces(">Z")
                .hole(self.mount.diam_insertion_internal())
            )
        result.faces(">Z").tag("mate")
        result.copyWorkplane(Cq.Workplane('XY')).tagPlane("dir", "+X")
        if not self.simplify_geometry:
            thread = self.mount.internal_thread(self.insertion_length).val()
            result = result.union(thread)
        for h in holes:
            cyl = Cq.Solid.makeCylinder(
                radius=2,
                height=self.diam * 2,
                pnt=(-self.diam, 0, h),
                dir=(1, 0, 0))
            result = result.cut(cyl)
        return result

    def connector(self, solid: bool = True):
        """
        Tags:
        * mate{1,2}: Mates to the connector

        WARNING: A tolerance lower than the defualt (maybe 2e-4) is required for
        STL export.
        """
        result = (
            Cq.Workplane('XY')
            .cylinder(
                radius=self.diam / 2,
                height=self.connector_length,
            )
        )
        for (tag, selector) in [("mate1", "<Z"), ("mate2", ">Z")]:
            result.faces(selector).tag(tag)
            result = (
                result
                .faces(selector)
                .workplane()
                .circle(self.mount.diam_connector_external() / 2)
                .extrude(self.insertion_length)
            )
        if not solid:
            result = result.faces(">Z").hole(self.diam_connector_internal)
        if not self.simplify_geometry:
            thread = self.mount.external_thread(self.insertion_length).val()
            result = (
                result
                .union(
                    thread
                    .located(Cq.Location((0, 0, -self.connector_length))))
                .union(
                    thread
                    .rotate((0,0,0), (1,0,0), angleDegrees=180)
                    .located(Cq.Location((0, 0, self.connector_length))))
            )
        result.copyWorkplane(Cq.Workplane('XY')).tagPlane("dir", "+X")
        return result

    def one_side_connector(self, height=None):
        if height is None:
            height = self.rim_length
        result = (
            Cq.Workplane('XY')
            .cylinder(
                radius=self.diam / 2,
                height=height,
                centered=(True, True, False)
            )
        )
        result.faces(">Z").tag("mate")
        result.faces("<Z").tag("base")
        result = (
            result
            .faces(">Z")
            .workplane()
            .circle(self.mount.diam_connector_external() / 2)
            .extrude(self.insertion_length)
        )
        result.copyWorkplane(Cq.Workplane('XY')).tagPlane("dir", "+X")
        if not self.simplify_geometry:
            thread = self.mount.external_thread(self.insertion_length).val()
            result = (
                result
                .union(
                    thread
                    # Avoids collision in some mating cases
                    .rotate((0,0,0), (1,0,0), angleDegrees=180)
                    .located(Cq.Location((0, 0, height + self.insertion_length))))
            )
        return result

    def threaded_core(self, length):
        """
        Generates a threaded core for unioning with other components
        """
        result = (
            Cq.Workplane('XY')
            .cylinder(
                radius=self.mount.diam_connector_external / 2,
                height=length,
                centered=(True, True, False),
            )
        )
        result.faces(">Z").tag("mate")
        result.faces("<Z").tag("base")
        if not self.simplify_geometry:
            thread = self.mount.external_thread(length=length).val()
            result = (
                result
                .union(thread)
            )
        return result

    def connector_insertion_assembly(self):
        connector_color = Cq.Color(0.8,0.8,0.5,0.3)
        insertion_color = Cq.Color(0.7,0.7,0.7,0.3)
        result = (
            Cq.Assembly()
            .add(self.connector(), name="c", color=connector_color)
            .add(self.insertion(), name="i1", color=insertion_color)
            .add(self.insertion(), name="i2", color=insertion_color)
            .constrain("c?mate1", "i1?mate", "Plane")
            .constrain("c?mate2", "i2?mate", "Plane")
            .constrain("c?dir", "i1?dir", "Axis")
            .constrain("c?dir", "i2?dir", "Axis")
            .solve()
        )
        return result
    def connector_one_side_insertion_assembly(self):
        connector_color = Cq.Color(0.8,0.8,0.5,0.3)
        insertion_color = Cq.Color(0.7,0.7,0.7,0.3)
        result = (
            Cq.Assembly()
            .add(self.insertion(), name="i", color=connector_color)
            .add(self.one_side_connector(), name="c", color=insertion_color)
            .constrain("i?mate", "c?mate", "Plane")
            .constrain("c?dir", "i?dir", "Axis")
            .solve()
        )
        return result