"""
Wheel geometry generation using build123d.

Creates worm wheel with two options:
1. Helical: Pure helical gear teeth (flat root) - simpler, point contact
2. Hobbed/Throated: Arc-bottomed teeth that match worm curvature - better contact
"""

import math
from typing import Optional, Literal
from build123d import *
from ..io.loaders import WheelParams, WormParams, AssemblyParams
from .features import (
    BoreFeature,
    KeywayFeature,
    SetScrewFeature,
    HubFeature,
    add_bore_and_keyway,
    create_hub
)

# Profile types per DIN 3975
# ZA: Straight flanks in axial section (Archimedean) - best for CNC machining
# ZK: Slightly convex flanks - better for 3D printing (reduces stress concentrations)
ProfileType = Literal["ZA", "ZK", "ZI"]


class WheelGeometry:
    """
    Generates 3D geometry for a worm wheel.

    Supports two tooth types:
    - helical: Pure helical gear teeth (no throat cut) - simpler geometry
    - hobbed: Helical teeth with toroidal throat cut - attempts to match worm

    Optionally adds bore and keyway features.
    """

    def __init__(
        self,
        params: WheelParams,
        worm_params: WormParams,
        assembly_params: AssemblyParams,
        face_width: float = None,
        throated: bool = False,
        bore: Optional[BoreFeature] = None,
        keyway: Optional[KeywayFeature] = None,
        ddcut: Optional['DDCutFeature'] = None,
        set_screw: Optional[SetScrewFeature] = None,
        hub: Optional[HubFeature] = None,
        profile: ProfileType = "ZA"
    ):
        """
        Initialize wheel geometry generator.

        Args:
            params: Wheel parameters from calculator
            worm_params: Worm parameters (needed for throating)
            assembly_params: Assembly parameters
            face_width: Wheel face width in mm (default: auto-calculated)
            throated: If True, apply throat cut (hobbed style); if False, pure helical
            bore: Optional bore feature specification
            keyway: Optional keyway feature specification (requires bore)
            set_screw: Optional set screw feature specification (requires bore)
            hub: Optional hub feature specification (flush/extended/flanged)
            profile: Tooth profile type per DIN 3975:
                     "ZA" - Straight flanks (trapezoidal) - best for CNC (default)
                     "ZK" - Slightly convex flanks - better for 3D printing
        """
        self.params = params
        self.worm_params = worm_params
        self.assembly_params = assembly_params
        self.throated = throated
        self.bore = bore
        self.keyway = keyway
        self.ddcut = ddcut
        self.set_screw = set_screw
        self.hub = hub
        self.profile = profile.upper() if isinstance(profile, str) else profile

        # Set keyway as hub type if specified
        if self.keyway is not None:
            self.keyway.is_shaft = False

        # Calculate face width if not provided
        if face_width is None:
            d1 = worm_params.pitch_diameter_mm
            ratio = assembly_params.ratio
            self.face_width = 0.73 * (d1 ** (1/3)) * math.sqrt(ratio)
            self.face_width = max(0.3 * d1, min(0.67 * d1, self.face_width))
        else:
            self.face_width = face_width

        # Cache for built geometry (avoids rebuilding on export)
        self._part = None

    def build(self) -> Part:
        """
        Build the complete wheel geometry.

        Returns:
            build123d Part object ready for export
        """
        # Return cached geometry if already built
        if self._part is not None:
            return self._part

        # Create helical gear (throating is built into the tooth profile)
        gear = self._create_helical_gear()

        # Add bore, keyway, and set screw if specified
        if self.bore is not None or self.keyway is not None or self.ddcut is not None or self.set_screw is not None:
            gear = add_bore_and_keyway(
                gear,
                part_length=self.face_width,
                bore=self.bore,
                keyway=self.keyway,
                ddcut=self.ddcut,
                set_screw=self.set_screw,
                axis=Axis.Z
            )

        # Add hub if specified (additive feature, comes after subtractive features)
        if self.hub is not None:
            bore_diameter = self.bore.diameter if self.bore is not None else None
            gear = create_hub(
                gear,
                hub=self.hub,
                wheel_face_width=self.face_width,
                wheel_root_diameter=self.params.root_diameter_mm,
                bore_diameter=bore_diameter,
                axis=Axis.Z
            )

        # Cache the built geometry
        self._part = gear
        return gear

    def _create_helical_gear(self) -> Part:
        """
        Create helical gear by extruding and twisting tooth space profiles.

        Uses extrusion with rotation rather than helix sweep to avoid
        self-intersection issues with tight helix pitches.

        For throated wheels, the tooth space profile has an arc at the root
        that matches the worm's curvature, creating the throat naturally
        as part of the tooth geometry.
        """
        z = self.params.num_teeth
        m = self.params.module_mm
        tip_radius = self.params.tip_diameter_mm / 2
        root_radius = self.params.root_diameter_mm / 2
        pitch_radius = self.params.pitch_diameter_mm / 2
        pressure_angle = math.radians(self.assembly_params.pressure_angle_deg)

        # The wheel's helix angle equals 90° - worm lead angle
        # This determines how much the teeth twist over the face width
        lead_angle = math.radians(self.worm_params.lead_angle_deg)

        # Total twist angle over face width
        # At pitch radius, the teeth advance by (face_width * tan(lead_angle)) axially
        # This corresponds to a rotation of: twist = face_width * tan(lead_angle) / pitch_radius (radians)
        twist_radians = self.face_width * math.tan(lead_angle) / pitch_radius
        twist_degrees = math.degrees(twist_radians)

        # Create gear blank
        blank = Cylinder(
            radius=tip_radius,
            height=self.face_width,
            align=(Align.CENTER, Align.CENTER, Align.CENTER)
        )

        # Calculate tooth space dimensions
        circular_pitch = math.pi * m
        space_width_pitch = circular_pitch / 2 + self.assembly_params.backlash_mm

        # Space is wider at tip, narrower at root (inverse of tooth shape)
        space_width_tip = space_width_pitch + 2 * (tip_radius - pitch_radius) * math.tan(pressure_angle)
        space_width_root = space_width_pitch - 2 * (pitch_radius - root_radius) * math.tan(pressure_angle)
        space_width_root = max(0.1 * m, space_width_root)

        half_root = space_width_root / 2
        half_tip = space_width_tip / 2

        # Extend cut beyond face to get clean edges
        extension = 0.5
        cut_height = self.face_width + 2 * extension

        # Number of sections for lofting the twisted extrusion
        num_sections = max(8, int(abs(twist_degrees) / 5) + 1)

        # For throated wheels, calculate the worm position for arc profile
        if self.throated:
            centre_distance = self.assembly_params.centre_distance_mm
            worm_tip_radius = self.worm_params.tip_diameter_mm / 2
            # Add small clearance for fit
            arc_radius = worm_tip_radius + 0.1

        # Cut tooth spaces
        gear = blank

        for i in range(z):
            base_angle = (360 / z) * i

            # Create sections along Z for lofting
            sections = []

            for s in range(num_sections):
                t = s / (num_sections - 1)  # 0 to 1
                z_pos = -cut_height / 2 + t * cut_height

                # Rotation at this Z position (linear interpolation of twist)
                rotation_at_z = -twist_degrees / 2 + t * twist_degrees
                section_angle = base_angle + rotation_at_z
                section_angle_rad = math.radians(section_angle)

                # Create profile plane at this Z, rotated appropriately
                radial = Vector(math.cos(section_angle_rad), math.sin(section_angle_rad), 0)
                tangent = Vector(-math.sin(section_angle_rad), math.cos(section_angle_rad), 0)

                # Profile plane: origin at pitch radius, X = radial outward, Y = tangential
                origin = Vector(
                    pitch_radius * math.cos(section_angle_rad),
                    pitch_radius * math.sin(section_angle_rad),
                    z_pos
                )
                profile_plane = Plane(origin=origin, x_dir=radial, z_dir=Vector(0, 0, 1))

                # Profile offsets from pitch radius (in radial direction)
                inner = root_radius - pitch_radius - 0.3  # Extend below root
                outer = tip_radius + 0.3 - pitch_radius

                # For throated wheels, the root depth varies with Z position
                # to match the worm's cylindrical surface
                if self.throated and abs(z_pos) < arc_radius:
                    # Calculate where the worm surface is at this Z
                    worm_surface_dist = centre_distance - math.sqrt(arc_radius**2 - z_pos**2)
                    throated_inner = worm_surface_dist - pitch_radius
                    # Use the shallower of the two (worm surface or calculated root)
                    actual_inner = max(inner, throated_inner)
                else:
                    actual_inner = inner

                with BuildSketch(profile_plane) as sk:
                    with BuildLine():
                        if self.profile == "ZA":
                            # ZA profile: Straight flanks (trapezoidal) per DIN 3975
                            # Best for CNC machining - simple, accurate, standard
                            root_left = (actual_inner, -half_root)
                            root_right = (actual_inner, half_root)
                            tip_left = (outer, -half_tip)
                            tip_right = (outer, half_tip)

                            Line(root_left, tip_left)      # Left flank (straight)
                            Line(tip_left, tip_right)      # Tip
                            Line(tip_right, root_right)    # Right flank (straight)
                            Line(root_right, root_left)    # Root (closes)

                        elif self.profile == "ZK":
                            # ZK profile: Circular arc flanks per DIN 3975 Type K
                            # Biconical grinding wheel profile - convex circular arc
                            # Better for 3D printing and reduces stress concentrations

                            # Generate circular arc flanks
                            num_points = 9  # Points per flank for smooth arc
                            left_flank = []
                            right_flank = []

                            # Arc radius typically 0.4-0.5 × module for biconical cutter
                            arc_radius = 0.45 * self.params.module_mm

                            # Calculate arc center position
                            flank_height = outer - actual_inner
                            flank_width_change = half_root - half_tip

                            # Angle of straight flank for reference
                            if flank_width_change > 0 and flank_height > 0:
                                flank_angle = math.atan(flank_width_change / flank_height)
                            else:
                                flank_angle = 0

                            # Generate arc points
                            for j in range(num_points):
                                t = j / (num_points - 1)
                                r_pos = actual_inner + t * flank_height

                                # Circular arc deviation from straight line
                                linear_width = half_root + t * (half_tip - half_root)

                                # Arc bulge (circular, not parabolic)
                                arc_param = t * math.pi  # 0 to π
                                arc_bulge = arc_radius * 0.15 * math.sin(arc_param)  # Circular arc approximation

                                width = linear_width + arc_bulge
                                left_flank.append((r_pos, -width))
                                right_flank.append((r_pos, width))

                            # Build profile with circular arc flanks
                            Spline(left_flank)
                            Line(left_flank[-1], right_flank[-1])  # Tip
                            Spline(list(reversed(right_flank)))
                            Line(right_flank[0], left_flank[0])    # Root (closes)

                        elif self.profile == "ZI":
                            # ZI profile: Involute helicoid per DIN 3975 Type I
                            # In axial section, appears as straight flanks (generatrix of involute helicoid)
                            # The involute shape is in normal section (perpendicular to thread)
                            # Manufactured by hobbing

                            root_left = (actual_inner, -half_root)
                            root_right = (actual_inner, half_root)
                            tip_left = (outer, -half_tip)
                            tip_right = (outer, half_tip)

                            Line(root_left, tip_left)      # Left flank (straight generatrix)
                            Line(tip_left, tip_right)      # Tip
                            Line(tip_right, root_right)    # Right flank (straight generatrix)
                            Line(root_right, root_left)    # Root (closes)

                        else:
                            raise ValueError(f"Unknown profile type: {self.profile}")
                    make_face()

                sections.append(sk.sketch.faces()[0])

            # Loft the sections to create twisted tooth space
            try:
                space = loft(sections, ruled=True)
                gear = gear - space
            except Exception as e:
                print(f"Warning: Tooth space {i} failed: {e}")

        return gear

    def show(self):
        """Display the wheel in OCP viewer."""
        wheel = self.build()
        try:
            from ocp_vscode import show as ocp_show
            ocp_show(wheel)
        except ImportError:
            try:
                show(wheel)
            except:
                print("No viewer available.")
        return wheel

    def export_step(self, filepath: str):
        """Export wheel to STEP file (builds if not already built)."""
        if self._part is None:
            self.build()

        if hasattr(self._part, 'export_step'):
            self._part.export_step(filepath)
        else:
            from build123d import export_step as exp_step
            exp_step(self._part, filepath)

        print(f"Exported wheel to {filepath}")
