import math
from typing import Optional
from dataclasses import dataclass
import cadquery as Cq

@dataclass
class TorsionSpring:
    """
    A torsion spring with abridged geometry (since sweep is slow)
    """
    # Outer radius
    radius: float = 12.0
    height: float = 20.0
    thickness: float = 2.0

    # Angle (in degrees) between the two legs at neutral position
    angle_neutral: float = 90.0

    tail_length: float = 25.0
    right_handed: bool = False

    torsion_rate: Optional[float] = None

    @property
    def radius_inner(self) -> float:
        return self.radius - self.thickness

    def torque_at(self, theta: float) -> float:
        return self.torsion_rate * theta

    def generate(self, deflection: float = 0):
        omega = self.angle_neutral + deflection
        omega = -omega if self.right_handed else omega
        base = (
            Cq.Workplane('XY')
            .cylinder(height=self.height, radius=self.radius,
                      centered=(True, True, False))
        )
        base.faces(">Z").tag("top")
        base.faces("<Z").tag("bot")

        box_shift = -self.radius if self.right_handed else self.radius-self.thickness
        tail = Cq.Solid.makeCylinder(
            height=self.tail_length,
            radius=self.thickness / 2)
        # points cylinder to +X
        dy = self.radius - self.thickness / 2
        if self.right_handed:
            dy = -dy
        loc_dir_x = Cq.Location((0, 0, self.thickness / 2), (0, 1, 0), 90)
        loc_shift = Cq.Location((0, dy, 0))
        loc_top = Cq.Location((0, 0, self.height - self.thickness), (0, 0, 1), omega + 180)
        result = (
            base
            .cylinder(
                height=self.height,
                radius=self.radius - self.thickness,
                combine='s',
                centered=(True, True, True))
            .union(tail.located(loc_shift * loc_dir_x))
            .union(tail.located(loc_top * loc_shift.inverse * loc_dir_x))
            .clean()
        )
        r = -self.radius if self.right_handed else self.radius
        plane = result.copyWorkplane(Cq.Workplane('XY'))
        plane.polyline([(0, r, 0), (self.tail_length, r, 0)],
                       forConstruction=True).tag("dir_bot")
        omega = math.radians(omega)
        c, s = math.cos(omega), math.sin(omega)
        l = -self.tail_length
        plane.polyline([
            (-s * r,        c * r, self.height),
            (c * l - s * r, c * r + s * l, self.height)],
                       forConstruction=True).tag("dir_top")
        return result