from dataclasses import dataclass
import cadquery as Cq
from nhf import Role
from nhf.build import Model, target
import nhf.parts.springs as springs
import nhf.utils

TOL = 1e-6

@dataclass
class DiskJoint(Model):
    """
    Sandwiched disk joint for the wrist and elbow
    """

    radius_housing: float = 22.0
    radius_disk: float = 20.0
    radius_spring: float = 9 / 2
    radius_axle: float = 3.0

    housing_thickness: float = 5.0
    disk_thickness: float = 5.0
    # Gap between disk and the housing
    #disk_thickness_gap: float = 0.1
    spring_thickness: float = 1.3
    spring_height: float = 6.5
    spring_tail_length: float = 45.0
    spring_angle: float = 90.0
    spring_angle_shift: float = 0
    wall_inset: float = 2.0

    # Angular span of movement
    movement_angle: float = 120.0
    # Angular span of tongue on disk
    tongue_span: float = 30.0
    tongue_length: float = 10.0

    generate_inner_wall: bool = False


    def __post_init__(self):
        super().__init__(name="disk-joint")
        assert self.housing_thickness > self.wall_inset
        assert self.radius_housing > self.radius_disk
        assert self.radius_disk > self.radius_axle
        assert self.housing_upper_carve_offset > 0

    def spring(self):
        return springs.torsion_spring(
            radius=self.radius_spring,
            height=self.spring_height,
            thickness=self.spring_thickness,
            tail_length=self.spring_tail_length,
            right_handed=False,
        )

    @property
    def total_thickness(self):
        return self.housing_thickness * 2 + self.disk_thickness

    @property
    def opening_span(self):
        return self.movement_angle + self.tongue_span

    @property
    def housing_upper_carve_offset(self):
        return self.housing_thickness + self.disk_thickness - self.spring_height

    @property
    def radius_spring_internal(self):
        return self.radius_spring - self.spring_thickness

    @target(name="disk")
    def disk(self) -> Cq.Workplane:
        cut = (
            Cq.Solid.makeBox(
                length=self.spring_tail_length,
                width=self.spring_thickness,
                height=self.disk_thickness,
            )
            .located(Cq.Location((0, self.radius_spring_internal, 0)))
            .rotate((0, 0, 0), (0, 0, 1), self.spring_angle_shift)
        )
        result = (
            Cq.Workplane('XY')
            .cylinder(
                height=self.disk_thickness,
                radius=self.radius_disk,
                centered=(True, True, False)
            )
            .copyWorkplane(Cq.Workplane('XY'))
            .cylinder(
                height=self.disk_thickness,
                radius=self.radius_spring,
                centered=(True, True, False),
                combine='cut',
            )
            .cut(cut)
        )
        plane = result.copyWorkplane(Cq.Workplane('XY'))
        plane.tagPlane("dir", direction="+X")
        plane.workplane(offset=self.disk_thickness).tagPlane("mate_top")
        result.copyWorkplane(Cq.Workplane('YX')).tagPlane("mate_bot")

        radius_tongue = self.radius_disk + self.tongue_length
        tongue = (
            Cq.Solid.makeCylinder(
                height=self.disk_thickness,
                radius=radius_tongue,
                angleDegrees=self.tongue_span,
            ).cut(Cq.Solid.makeCylinder(
                height=self.disk_thickness,
                radius=self.radius_disk,
            ))
        )
        result = result.union(tongue, tol=TOL)
        return result

    def wall(self) -> Cq.Compound:
        height = self.disk_thickness + self.wall_inset
        wall = Cq.Solid.makeCylinder(
            radius=self.radius_housing,
            height=height,
            angleDegrees=360 - self.opening_span,
        ).cut(Cq.Solid.makeCylinder(
            radius=self.radius_disk,
            height=height,
        )).rotate((0, 0, 0), (0, 0, 1), self.opening_span)
        return wall

    @target(name="housing-lower")
    def housing_lower(self) -> Cq.Workplane:
        result = (
            Cq.Workplane('XY')
            .cylinder(
                radius=self.radius_housing,
                height=self.housing_thickness,
                centered=(True, True, False),
            )
            .cut(Cq.Solid.makeCylinder(
                radius=self.radius_axle,
                height=self.housing_thickness,
            ))
        )
        result.faces(">Z").tag("mate")
        result.faces(">Z").workplane().tagPlane("dir", direction="+X")
        result = result.cut(
            self
            .wall()
            .mirror("XY")
            .located(Cq.Location((0, 0, self.housing_thickness + self.disk_thickness)))
            #.rotate((0, 0, 0), (1, 0, 0), 180)
            #.located(Cq.Location((0, 0, self.disk_thickness + self.housing_thickness)))
        )
        return result

    @target(name="housing-upper")
    def housing_upper(self) -> Cq.Workplane:
        carve = (
            Cq.Solid.makeCylinder(
                radius=self.radius_spring,
                height=self.housing_thickness
            ).fuse(Cq.Solid.makeBox(
                length=self.spring_tail_length,
                width=self.spring_thickness,
                height=self.housing_thickness
            ).located(Cq.Location((0, self.radius_spring_internal, 0))))
        ).rotate((0, 0, 0), (0, 0, 1), self.spring_angle - self.spring_angle_shift)
        result = (
            Cq.Workplane('XY')
            .cylinder(
                radius=self.radius_housing,
                height=self.housing_thickness,
                centered=(True, True, False),
            )
        )
        result.faces(">Z").tag("mate")
        result.copyWorkplane(Cq.Workplane('XY')).tagPlane("dir", direction="+X")
        result = result.faces(">Z").hole(self.radius_axle * 2)

        # tube which holds the spring interior
        if self.generate_inner_wall:
            tube = (
                Cq.Solid.makeCylinder(
                    radius=self.radius_spring_internal,
                    height=self.disk_thickness + self.housing_thickness,
                ).cut(Cq.Solid.makeCylinder(
                    radius=self.radius_axle,
                    height=self.disk_thickness + self.housing_thickness,
                ))
            )
            result = result.union(tube)
        wall = (
            self.wall()
            .rotate((0, 0, 0), (1, 0, 0), 180)
            .located(Cq.Location((0, 0, self.disk_thickness + self.housing_thickness + self.wall_inset)))
        )
        result = (
            result
            .union(wall, tol=TOL)
            .cut(carve.located(Cq.Location((0, 0, self.housing_upper_carve_offset))))
        )
        return result


    def assembly(self) -> Cq.Assembly:
        result = (
            Cq.Assembly()
            .add(self.disk(), name="disk", color=Role.CHILD.color)
            .add(self.housing_lower(), name="housing_lower", color=Role.PARENT.color)
            .add(self.housing_upper(), name="housing_upper", color=Role.PARENT.color)
            .constrain("disk?mate_bot", "housing_lower?mate", "Plane")
            .constrain("disk?mate_top", "housing_upper?mate", "Plane")
            .solve()
        )
        return result

if __name__ == '__main__':
    p = DiskJoint()
    p.build_all()