from dataclasses import dataclass, field
import cadquery as Cq
from nhf.parts.joints import HirthJoint
from nhf import Material, Role
from nhf.build import Model, TargetKind, target, assembly
import nhf.utils

@dataclass(frozen=True, kw_only=True)
class Mannequin:
    """
    A mannequin for calibration
    """

    shoulder_width: float = 400
    shoulder_to_waist: float = 440
    waist_width: float = 250
    head_height: float = 220.0
    neck_height: float = 105.0
    neck_diam: float = 140
    head_diam: float = 210
    torso_thickness: float = 150

    def generate(self) -> Cq.Workplane:
        head_neck = (
            Cq.Workplane("XY")
            .cylinder(
                radius=self.neck_diam/2,
                height=self.neck_height,
                centered=(True, True, False))
            .faces(">Z")
            .workplane()
            .cylinder(
                radius=self.head_diam/2,
                height=self.head_height,
                combine=True, centered=(True, True, False))
        )
        result = (
            Cq.Workplane("XY")
            .rect(self.waist_width, self.torso_thickness)
            .workplane(offset=self.shoulder_to_waist)
            .rect(self.shoulder_width, self.torso_thickness)
            .loft(combine=True)
            .union(head_neck.translate((0, 0, self.shoulder_to_waist)))
        )
        return result.translate((0, self.torso_thickness / 2, 0))


BASE_POS_X = 60.0
BASE_POS_Y = 100.0

@dataclass(kw_only=True)
class Harness(Model):
    thickness: float = 25.4 / 8
    width: float = 220.0
    height: float = 310.0
    fillet: float = 10.0

    wing_base_pos: list[tuple[str, float, float]] = field(default_factory=lambda: [
        ("r1", BASE_POS_X + 10, BASE_POS_Y),
        ("l1", -BASE_POS_X - 10, BASE_POS_Y),
        ("r2", BASE_POS_X + 10, 0),
        ("l2", -BASE_POS_X - 10, 0),
        ("r3", BASE_POS_X, -BASE_POS_Y),
        ("l3", -BASE_POS_X, -BASE_POS_Y),
    ])
    # Holes drilled onto harness for attachment with HS joint
    harness_to_root_conn_diam: float = 6

    hs_hirth_joint: HirthJoint = field(default_factory=lambda: HirthJoint(
        radius=25.0,
        radius_inner=15.0,
        tooth_height=7.0,
        base_height=5.0,
        n_tooth=24,
    ))

    hs_joint_base_width: float = 85
    hs_joint_base_thickness: float = 10
    hs_joint_corner_fillet: float = 5
    hs_joint_corner_cbore_diam: float = 12
    hs_joint_corner_cbore_depth: float = 2
    hs_joint_corner_inset: float = 12

    hs_joint_axis_diam: float = 12.0
    hs_joint_axis_cbore_diam: float = 20
    hs_joint_axis_cbore_depth: float = 3

    mannequin: Mannequin = Mannequin()

    def __post_init__(self):
        super().__init__(name="harness")

    @target(name="profile", kind=TargetKind.DXF)
    def profile(self) -> Cq.Sketch:
        """
        Creates the harness shape
        """
        w, h = self.width / 2, self.height / 2
        sketch = (
            Cq.Sketch()
            .polygon([
                (w, h),
                (w, -h),
                (-w, -h),
                (-w, h),
                #(0.7 * w, h),
                #(w, 0),
                #(0.7 * w, -h),
                #(0.7 * -w, -h),
                #(-w, 0),
                #(0.7 * -w, h),
            ])
            #.rect(self.harness_width, self.harness_height)
            .vertices()
            .fillet(self.fillet)
        )
        for tag, x, y in self.wing_base_pos:
            conn = [(px + x, py + y) for px, py in self.hs_joint_harness_conn()]
            sketch = (
                sketch
                .push(conn)
                .tag(tag)
                .circle(self.harness_to_root_conn_diam / 2, mode='s')
                .reset()
            )
        return sketch

    def surface(self) -> Cq.Workplane:
        """
        Creates the harness shape
        """
        result = (
            Cq.Workplane('XZ')
            .placeSketch(self.profile())
            .extrude(self.thickness)
        )
        result.faces(">Y").tag("mount")
        plane = result.faces(">Y").workplane()
        for tag, x, y in self.wing_base_pos:
            conn = [(px + x, py + y) for px, py
                    in self.hs_joint_harness_conn()]
            for i, (px, py) in enumerate(conn):
                plane.moveTo(px, py).tagPoint(f"{tag}_{i}")
        return result

    def hs_joint_harness_conn(self) -> list[tuple[int, int]]:
        """
        Generates a set of points corresponding to the connectorss
        """
        dx = self.hs_joint_base_width / 2 - self.hs_joint_corner_inset
        return [
            (dx, dx),
            (dx, -dx),
            (-dx, -dx),
            (-dx, dx),
        ]

    @target(name="hs-joint-parent")
    def hs_joint_parent(self):
        """
        Parent part of the Houjuu-Scarlett joint, which is composed of a Hirth
        coupling, a cylindrical base, and a mounting base.
        """
        hirth = self.hs_hirth_joint.generate()
        conn = self.hs_joint_harness_conn()
        result = (
            Cq.Workplane('XY')
            .box(
                self.hs_joint_base_width,
                self.hs_joint_base_width,
                self.hs_joint_base_thickness,
                centered=(True, True, False))
            .translate((0, 0, -self.hs_joint_base_thickness))
            .edges("|Z")
            .fillet(self.hs_joint_corner_fillet)
            .faces(">Z")
            .workplane()
            .pushPoints(conn)
            .cboreHole(
                diameter=self.harness_to_root_conn_diam,
                cboreDiameter=self.hs_joint_corner_cbore_diam,
                cboreDepth=self.hs_joint_corner_cbore_depth)
        )
        # Creates a plane parallel to the holes but shifted to the base
        plane = result.faces(">Z").workplane(offset=-self.hs_joint_base_thickness)

        for i, (px, py) in enumerate(conn):
            plane.moveTo(px, py).tagPoint(f"h{i}")
        result = (
            result
            .faces(">Z")
            .workplane()
            .union(hirth, tol=0.1)
            .clean()
        )
        result = (
            result.faces("<Z")
            .workplane()
            .cboreHole(
                diameter=self.hs_joint_axis_diam,
                cboreDiameter=self.hs_joint_axis_cbore_diam,
                cboreDepth=self.hs_joint_axis_cbore_depth,
            )
            .clean()
        )
        result.faces("<Z").tag("base")
        return result

    @assembly()
    def assembly(self) -> Cq.Assembly:
        harness = self.surface()
        mannequin_z = self.mannequin.shoulder_to_waist * 0.6
        result = (
            Cq.Assembly()
            .addS(
                harness, name="base",
                material=Material.WOOD_BIRCH,
                role=Role.STRUCTURE)
            .constrain("base", "Fixed")
            .addS(
                self.mannequin.generate(),
                name="mannequin",
                role=Role.FIXTURE,
                loc=Cq.Location((0, -self.thickness, -mannequin_z), (0, 0, 1), 180))
            .constrain("mannequin", "Fixed")
        )
        for name in ["l1", "l2", "l3", "r1", "r2", "r3"]:
            j = self.hs_joint_parent()
            result.addS(
                j, name=name,
                role=Role.PARENT,
                material=Material.PLASTIC_PLA)
            for i in range(4):
                result.constrain(f"base?{name}_{i}", f"{name}?h{i}", "Point")
        result.solve()
        return result