from nhf.build import Model, TargetKind, target, assembly, submodel
from nhf.materials import Role, Material
import nhf.utils

import math
from dataclasses import dataclass
import cadquery as Cq

@dataclass
class Onbashira(Model):

    n_side: int = 6
    # Dimensions of each side panel
    side_width: float = 200.0
    side_length: float = 600.0

    side_thickness: float = 25.4 / 8

    # Dimensions of gun barrels
    barrel_diam: float = 45.0
    barrel_length: float = 300.0
    # Radius from barrel centre to axis
    rotation_radius: float = 75.0
    # Radius of ball bearings
    bearing_ball_diam: float = 30.0
    bearing_ball_gap: float = 1.0
    bearing_height: float = 40.0
    bearing_thickness: float = 20.0

    material_side: Material = Material.WOOD_BIRCH
    material_bearing: Material = Material.PLASTIC_PLA
    material_bearing_ball: Material = Material.ACRYLIC_TRANSPARENT
    material_brace: Material = Material.ALUMINUM

    def __post_init__(self):
        assert self.n_side >= 3
        # Bulk must be large enough for the barrel + bearing to rotate
        assert self.bulk_radius - self.side_thickness - self.bearing_thickness - self.bearing_diam > self.rotation_radius + self.barrel_diam / 2
        assert self.bearing_height > self.bearing_diam

    @property
    def angle_side(self) -> float:
        return 360 / self.n_side
    @property
    def angle_dihedral(self) -> float:
        return 180 - self.angle_side
    @property
    def bulk_radius(self) -> float:
        """
        Calculate radius of the bulk to the centre
        """
        return self.side_width / 2 / math.tan(math.radians(self.angle_side / 2))
    @property
    def bearing_diam(self) -> float:
        return self.bearing_ball_diam + self.bearing_ball_gap
    @property
    def bearing_radius(self) -> float:
        return self.bulk_radius - self.side_thickness - self.bearing_thickness - self.bearing_diam / 2

    @target(name="side-panel", kind=TargetKind.DXF)
    def profile_side_panel(self) -> Cq.Sketch:
        return (
            Cq.Sketch()
            .rect(self.side_width, self.side_length)
        )

    def side_panel(self) -> Cq.Workplane:
        w = self.side_width
        l = self.side_length
        result = (
            Cq.Workplane()
            .placeSketch(self.profile_side_panel())
            .extrude(self.side_thickness)
        )
        intersector = (
            Cq.Workplane('XZ')
            .polyline([
                (-w/2, 0),
                (w/2, 0),
                (0, self.bulk_radius),
            ])
            .close()
            .extrude(l)
            .translate(Cq.Vector(0, l/2,0))
        )
        # Intersect the side panel
        return result * intersector

    def bearing_channel(self) -> Cq.Solid:
        """
        Generates a toroidal channel for the ball bearings
        """
        return Cq.Solid.makeTorus(
            radius1=self.bearing_radius,
            radius2=self.bearing_diam/2,
        )
    @target(name="inner-rotor")
    def inner_rotor(self) -> Cq.Workplane:
        r_outer = self.bearing_radius
        base = Cq.Solid.makeCylinder(
            radius=r_outer,
            height=self.bearing_height
        ).translate(Cq.Vector(0,0,-self.bearing_height/2))
        r_rot = self.rotation_radius
        channel = self.bearing_channel()
        return (
            Cq.Workplane()
            .add(base - channel)
            .faces(">Z")
            .workplane()
            .polygon(
                nSides=self.n_side,
                diameter=2 * r_rot,
                forConstruction=True
            )
            .vertices()
            .hole(self.barrel_diam)
        )
        return base - channel
    @target(name="outer-rotor")
    def outer_rotor(self) -> Cq.Workplane:
        polygon_radius = (self.bulk_radius - self.side_thickness) / math.cos(math.radians(self.angle_side / 2))
        profile = (
            Cq.Sketch()
            .regularPolygon(
                r=polygon_radius,
                n=self.n_side,
            )
        )
        inner = Cq.Solid.makeCylinder(
            radius=self.bearing_radius,
            height=self.bearing_height,
        )
        base = (
            Cq.Workplane()
            .placeSketch(profile)
            .extrude(self.bearing_height)
            .cut(inner)
            .translate(Cq.Vector(0,0,-self.bearing_height/2))
            .cut(self.bearing_channel())
        )
        r = self.bearing_radius * 2
        subtractor = Cq.Solid.makeBox(
            length=r * 2,
            width=r * 2,
            height=self.bearing_height,
        ).translate(Cq.Vector(-r, -r, -self.bearing_height))
        return base - subtractor
    def bearing_ball(self) -> Cq.Solid:
        return Cq.Solid.makeSphere(radius=self.bearing_ball_diam/2, angleDegrees1=-90)

    def assembly_bearing(self) -> Cq.Assembly:
        a = (
            Cq.Assembly()
            .addS(
                self.inner_rotor(),
                name="inner",
                material=self.material_bearing,
                role=Role.ROTOR,
            )
            .addS(
                self.outer_rotor(),
                name="outer",
                material=self.material_bearing,
                role=Role.STATOR,
            )
        )
        for i in range(self.n_side):
            ball = self.bearing_ball()
            a = a.addS(
                ball,
                name=f"bearing_ball{i}",
                material=self.material_bearing_ball,
                role=Role.BEARING,
                loc=Cq.Location.rot2d(i * 360/self.n_side) * Cq.Location(self.bearing_radius, 0, 0),
            )
        return a

    def assembly(self) -> Cq.Assembly:
        a = Cq.Assembly()
        side = self.side_panel()
        r = self.bulk_radius
        for i in range(6):
            a = a.addS(
                side,
                name=f"side{i}",
                material=self.material_side,
                role=Role.STRUCTURE | Role.DECORATION,
                loc=Cq.Location(0,0,0,0,i*60,0) * Cq.Location(0,0,-r)
            )
        return a