"""
Command-line interface for worm gear geometry generation.
"""

import argparse
import subprocess
import sys
import time
from pathlib import Path


class CLIProgressReporter:
    """Progress reporter for CLI with ETA calculation."""

    def __init__(self, task_name: str = "Processing"):
        self.task_name = task_name
        self.start_time = None
        self.last_percent = -1
        self.completed = False

    def __call__(self, message: str, percent: float):
        """Report progress with ETA."""
        if self.completed:
            return  # Don't print after completion

        if self.start_time is None:
            self.start_time = time.time()

        # Only update on significant progress (avoid spam)
        if percent >= 0 and (percent - self.last_percent >= 5 or percent >= 100):
            self.last_percent = percent
            elapsed = time.time() - self.start_time

            # Calculate ETA
            eta_str = ""
            if percent > 5 and percent < 100:
                total_estimated = elapsed / (percent / 100)
                remaining = total_estimated - elapsed
                if remaining > 60:
                    eta_str = f" (ETA: {remaining/60:.1f}min)"
                elif remaining > 0:
                    eta_str = f" (ETA: {remaining:.0f}s)"

            # Format progress bar
            bar_width = 20
            filled = int(bar_width * percent / 100)
            bar = "█" * filled + "░" * (bar_width - filled)

            # Print progress line (overwrite previous)
            print(f"\r  [{bar}] {percent:5.1f}%{eta_str}  ", end="", flush=True)

            if percent >= 100:
                self.completed = True
                total_time = time.time() - self.start_time
                if total_time > 60:
                    print(f"\n  Completed in {total_time/60:.1f} minutes")
                else:
                    print(f"\n  Completed in {total_time:.1f} seconds")

    def reset(self, task_name: str = None):
        """Reset for a new task."""
        if task_name:
            self.task_name = task_name
        self.start_time = None
        self.last_percent = -1
        self.completed = False


def get_version_string() -> str:
    """Get version string, including PR info if on a feature branch."""
    from importlib.metadata import PackageNotFoundError, version as pkg_version

    try:
        base_version = pkg_version("wormgear")
    except PackageNotFoundError:
        base_version = "0.0.0"

    # Try to detect git branch and PR
    try:
        # Get current branch
        result = subprocess.run(
            ["git", "rev-parse", "--abbrev-ref", "HEAD"],
            capture_output=True, text=True, timeout=5
        )
        branch = result.stdout.strip() if result.returncode == 0 else None

        if branch and branch != "main" and branch != "master":
            # On a feature branch - try to find PR number
            try:
                result = subprocess.run(
                    ["gh", "pr", "view", "--json", "number", "-q", ".number"],
                    capture_output=True, text=True, timeout=10
                )
                if result.returncode == 0 and result.stdout.strip():
                    pr_number = result.stdout.strip()
                    return f"{base_version} (PR#{pr_number})"
            except (subprocess.SubprocessError, OSError):
                pass
            # No PR found, show branch name
            return f"{base_version} ({branch})"

        # Check if working tree is dirty
        result = subprocess.run(
            ["git", "status", "--porcelain"],
            capture_output=True, text=True, timeout=5
        )
        if result.returncode == 0 and result.stdout.strip():
            return f"{base_version}-dev"

    except (subprocess.SubprocessError, OSError):
        pass

    return base_version

from ..io.loaders import (
    load_design_json,
    save_design_json,
    WormGearDesign,
    ManufacturingParams,
    ManufacturingFeatures
)
from ..io.package import generate_package, save_package_to_dir, create_package_zip
from ..enums import WormType, WormProfile, BoreType
from ..core.worm import WormGeometry
from ..core.globoid_worm import GloboidWormGeometry
from ..core.wheel import WheelGeometry
from ..core.virtual_hobbing import VirtualHobbingWheelGeometry
from ..core.features import (
    BoreFeature,
    KeywayFeature,
    DDCutFeature,
    SetScrewFeature,
    HubFeature,
    ReliefGrooveFeature,
    calculate_default_bore,
    calculate_default_ddcut,
    get_din_6885_keyway,
    get_set_screw_size
)
from ..core.mesh_alignment import (
    find_optimal_mesh_rotation,
    position_for_mesh,
    mesh_alignment_to_dict,
    MeshAlignmentResult,
)
from ..core.rim_thickness import (
    measure_rim_thickness,
    rim_thickness_to_dict,
    WHEEL_RIM_WARNING_THRESHOLD_MM,
    WORM_RIM_WARNING_THRESHOLD_MM,
)
from ..io.loaders import MeshAlignment, WormPosition, MeasuredGeometry, MeasurementPoint


def _parse_set_screw(args, json_features):
    """Parse set screw specification from CLI args or JSON features.

    Returns SetScrewFeature or None.
    """
    json_set_screw = json_features.set_screw if json_features else None
    if not (args.set_screw or json_set_screw):
        return None

    if args.set_screw_size:
        try:
            size_str = args.set_screw_size.upper()
            if size_str.startswith('M'):
                diameter = float(size_str[1:])
                return SetScrewFeature(
                    size=size_str,
                    diameter=diameter,
                    count=args.set_screw_count
                )
            else:
                print(f"  WARNING: Invalid set screw size '{args.set_screw_size}', using auto-size")
                return SetScrewFeature(count=args.set_screw_count)
        except ValueError:
            print(f"  WARNING: Could not parse set screw size '{args.set_screw_size}', using auto-size")
            return SetScrewFeature(count=args.set_screw_count)
    elif json_set_screw:
        size_str = json_set_screw.size.upper()
        try:
            diameter = float(size_str[1:]) if size_str.startswith('M') else None
            return SetScrewFeature(
                size=size_str,
                diameter=diameter,
                count=json_set_screw.count
            )
        except ValueError:
            return SetScrewFeature(count=json_set_screw.count)
    else:
        return SetScrewFeature(count=args.set_screw_count)


def _determine_bore(args, pitch_diameter, root_diameter, json_features, bore_cli_arg, ddcut_depth_arg):
    """Determine bore, keyway, DD-cut, and set screw features.

    Priority chain: CLI arg > JSON features > auto-calculate.

    Returns:
        tuple: (bore, keyway, ddcut, set_screw, bore_diameter, thin_rim_warning)
    """
    bore = None
    keyway = None
    ddcut = None
    set_screw = None
    bore_diameter = None
    thin_rim_warning = False

    if args.no_bore:
        return bore, keyway, ddcut, set_screw, bore_diameter, thin_rim_warning

    if bore_cli_arg is not None:
        # CLI override takes priority
        bore_diameter = bore_cli_arg
        actual_rim = (root_diameter - bore_diameter) / 2
        thin_rim_warning = actual_rim < 1.5 and actual_rim > 0
    elif json_features:
        # Check bore_type from JSON features
        json_bore_type = json_features.bore_type
        if json_bore_type == BoreType.NONE:
            bore_diameter = None
        elif json_bore_type == BoreType.CUSTOM and json_features.bore_diameter_mm is not None:
            bore_diameter = json_features.bore_diameter_mm
            actual_rim = (root_diameter - bore_diameter) / 2
            thin_rim_warning = actual_rim < 1.5 and actual_rim > 0
        else:
            bore_diameter, thin_rim_warning = calculate_default_bore(
                pitch_diameter, root_diameter
            )
    else:
        # No JSON features - auto-calculate bore
        bore_diameter, thin_rim_warning = calculate_default_bore(
            pitch_diameter, root_diameter
        )

    if bore_diameter is not None:
        bore = BoreFeature(diameter=bore_diameter)

        # Determine anti-rotation: CLI arg > JSON features > default (keyway if >= 6mm)
        json_anti_rotation = json_features.anti_rotation if json_features else None

        if args.dd_cut or (json_anti_rotation and json_anti_rotation.value == 'ddcut'):
            if ddcut_depth_arg:
                ddcut = DDCutFeature(depth=ddcut_depth_arg)
            else:
                depth_pct = json_features.ddcut_depth_percent if json_features else args.ddcut_depth_percent
                ddcut = calculate_default_ddcut(bore_diameter, depth_pct)
        elif args.no_keyway or (json_anti_rotation and json_anti_rotation.value == 'none'):
            pass
        elif bore_diameter >= 6.0:
            keyway = KeywayFeature()

        # Set screw
        set_screw = _parse_set_screw(args, json_features)

    return bore, keyway, ddcut, set_screw, bore_diameter, thin_rim_warning


def _build_features_desc(bore, bore_diameter, keyway, ddcut, set_screw):
    """Build the features description string for CLI output."""
    features_desc = ""
    if bore:
        features_desc += f", bore {bore_diameter}mm"
        if keyway:
            features_desc += " + keyway"
        elif ddcut:
            ddcut_depth = ddcut.get_depth(bore_diameter)
            features_desc += f" + DD-cut ({ddcut_depth:.1f}mm depth)"
        elif bore_diameter < 6.0:
            features_desc += " (no keyway - below DIN 6885 range)"
        if set_screw:
            screw_size, _ = set_screw.get_screw_specs(bore_diameter)
            screw_desc = f"{screw_size}"
            if set_screw.count > 1:
                screw_desc += f" x{set_screw.count}"
            features_desc += f" + set screw ({screw_desc})"
    return features_desc


def main():
    """Main CLI entry point."""
    parser = argparse.ArgumentParser(
        prog="wormgear",
        description="Generate CNC-ready STEP files for worm gear pairs from wormgear.studio designs",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  # Generate with auto-calculated bores and DIN 6885 keyways (default)
  wormgear design.json

  # Generate solid parts without bores
  wormgear design.json --no-bore

  # Override bore sizes (keyways auto-sized to match)
  wormgear design.json --worm-bore 8 --wheel-bore 12

  # Add set screw holes for shaft retention (auto-sized from bore)
  wormgear design.json --set-screw

  # Set screws with specific size and count
  wormgear design.json --set-screw --set-screw-size M4 --set-screw-count 2

  # Extended hub for bearing support
  wormgear design.json --hub-type extended --hub-length 15

  # Flanged hub with bolt holes for mounting
  wormgear design.json --hub-type flanged --flange-diameter 60 --flange-bolts 4

  # Bores but no keyways
  wormgear design.json --no-keyway

  # View in OCP viewer without saving
  wormgear design.json --view --no-save

  # Custom worm length and smoother geometry
  wormgear design.json --worm-length 50 --sections 72

  # Save extended JSON with all manufacturing features for reproducibility
  wormgear design.json --set-screw --hub-type extended --save-json complete_design.json

  # Generate globoid (hourglass) worm for 30-50% higher load capacity
  wormgear design.json --globoid

  # Tooth profiles per DIN 3975:
  wormgear design.json --profile ZA  # Straight flanks (default, CNC)
  wormgear design.json --profile ZK  # Circular arc (3D printing)
  wormgear design.json --profile ZI  # Involute (hobbing)

More info: https://wormgear.studio
        """
    )

    parser.add_argument(
        '-V', '--version',
        action='version',
        version=f'%(prog)s {get_version_string()}'
    )

    parser.add_argument(
        'design_file',
        type=str,
        help='JSON design file from wormgear.studio calculator'
    )

    parser.add_argument(
        '-o', '--output-dir',
        type=str,
        default='.',
        help='Output directory for STEP files (default: current directory)'
    )

    parser.add_argument(
        '--worm-length',
        type=float,
        default=40.0,
        help='Worm length in mm (default: 40)'
    )

    parser.add_argument(
        '--wheel-width',
        type=float,
        default=None,
        help='Wheel face width in mm (default: auto-calculated)'
    )

    parser.add_argument(
        '--wheel-tip-reduction',
        type=float,
        default=None,
        help='Reduce wheel tip diameter by this amount in mm (stub teeth)'
    )

    parser.add_argument(
        '--sections',
        type=int,
        default=36,
        help='Worm sections per turn for smoothness (default: 36)'
    )

    parser.add_argument(
        '--worm-only',
        action='store_true',
        help='Generate only the worm'
    )

    parser.add_argument(
        '--wheel-only',
        action='store_true',
        help='Generate only the wheel'
    )

    parser.add_argument(
        '--view',
        action='store_true',
        help='View in OCP viewer (requires ocp_vscode extension)'
    )

    parser.add_argument(
        '--no-save',
        action='store_true',
        help='Do not save STEP files (use with --view)'
    )

    parser.add_argument(
        '--zip-output',
        type=str,
        default=None,
        help='Create a ZIP archive instead of writing files to a directory'
    )

    parser.add_argument(
        '--no-3mf',
        action='store_true',
        help='Skip 3MF export (default: 3MF files generated)'
    )

    parser.add_argument(
        '--no-stl',
        action='store_true',
        help='Skip STL export (default: STL files generated)'
    )

    parser.add_argument(
        '--save-json',
        type=str,
        default=None,
        help='Save extended JSON with all manufacturing features (makes design fully reproducible)'
    )

    parser.add_argument(
        '--skip-mesh-alignment',
        action='store_true',
        help='Skip mesh alignment calculation (faster for complex geometry like virtual hobbing)'
    )

    parser.add_argument(
        '--hobbed',
        action='store_true',
        help='Generate hobbed wheel with throated teeth (default: helical without throating)'
    )

    parser.add_argument(
        '--globoid',
        action='store_true',
        help='Generate globoid (double-enveloping) worm with hourglass shape for 30-50%% higher load capacity'
    )

    parser.add_argument(
        '--profile',
        type=str,
        choices=['ZA', 'ZK', 'ZI', 'za', 'zk', 'zi'],
        default=None,
        help='Tooth profile type per DIN 3975: ZA=straight flanks/CNC (default), ZK=circular arc/3D print, ZI=involute/hobbing'
    )

    # Worm thread generation method
    parser.add_argument(
        '--generation-method',
        type=str,
        choices=['loft', 'sweep'],
        default='sweep',
        help='Worm thread generation method. '
             '"sweep" (default) sweeps a single profile along helix for cleaner topology and ~25x faster generation. '
             '"loft" creates sections along helix and lofts (legacy, slower).'
    )

    # Experimental virtual hobbing
    parser.add_argument(
        '--virtual-hobbing',
        action='store_true',
        help='EXPERIMENTAL: Use virtual hobbing simulation for wheel (more accurate but slower)'
    )

    parser.add_argument(
        '--hobbing-steps',
        type=int,
        default=72,
        help='Number of steps for virtual hobbing (default: 72, higher=more accurate but slower)'
    )

    parser.add_argument(
        '--trim-to-min-engagement',
        action='store_true',
        help='Trim wheel OD to throat minimum (globoid only) — removes flared edges'
    )

    # Bore and keyway options (defaults: auto-calculated bore with keyway)
    parser.add_argument(
        '--no-bore',
        action='store_true',
        help='Generate solid parts without bores (default: auto-calculated bores)'
    )

    parser.add_argument(
        '--no-keyway',
        action='store_true',
        help='Omit keyways (default: DIN 6885 keyways added with bores)'
    )

    # DD-cut options (alternative to keyways for small diameter bores)
    parser.add_argument(
        '--dd-cut',
        action='store_true',
        help='Use double-D cuts instead of keyways (for small diameter bores < 6mm)'
    )

    parser.add_argument(
        '--worm-ddcut-depth',
        type=float,
        default=None,
        help='Worm DD-cut flat depth in mm (default: auto ~15%% of bore diameter)'
    )

    parser.add_argument(
        '--wheel-ddcut-depth',
        type=float,
        default=None,
        help='Wheel DD-cut flat depth in mm (default: auto ~15%% of bore diameter)'
    )

    parser.add_argument(
        '--ddcut-depth-percent',
        type=float,
        default=15.0,
        help='DD-cut depth as percentage of bore diameter (default: 15.0 for ~70%% flat-to-flat). Common: 10%% (80%% f-t-f), 15%% (70%% f-t-f), 20%% (60%% f-t-f)'
    )

    parser.add_argument(
        '--worm-bore',
        type=float,
        default=None,
        help='Override worm bore diameter in mm (default: auto ~25%% of pitch diameter)'
    )

    parser.add_argument(
        '--wheel-bore',
        type=float,
        default=None,
        help='Override wheel bore diameter in mm (default: auto ~25%% of pitch diameter)'
    )

    # Set screw options (default: no set screws unless explicitly requested)
    parser.add_argument(
        '--set-screw',
        action='store_true',
        help='Add set screw holes for shaft retention (requires bore)'
    )

    parser.add_argument(
        '--set-screw-size',
        type=str,
        default=None,
        help='Override set screw size (e.g., "M3", "M4") - auto-sized from bore if not specified'
    )

    parser.add_argument(
        '--set-screw-count',
        type=int,
        default=1,
        help='Number of set screws (1-3, default: 1, evenly distributed)'
    )

    # Hub options (for wheel only, default: flush hub)
    parser.add_argument(
        '--hub-type',
        type=str,
        choices=['flush', 'extended', 'flanged'],
        default='flush',
        help='Hub type for wheel mounting (default: flush)'
    )

    parser.add_argument(
        '--hub-length',
        type=float,
        default=10.0,
        help='Hub extension length in mm for extended/flanged hubs (default: 10mm)'
    )

    parser.add_argument(
        '--flange-diameter',
        type=float,
        default=None,
        help='Flange outer diameter in mm (for flanged hub, default: auto-sized)'
    )

    parser.add_argument(
        '--flange-thickness',
        type=float,
        default=5.0,
        help='Flange thickness in mm (for flanged hub, default: 5mm)'
    )

    parser.add_argument(
        '--flange-bolts',
        type=int,
        default=4,
        help='Number of bolt holes in flange (for flanged hub, 0-8, default: 4)'
    )

    parser.add_argument(
        '--bolt-diameter',
        type=float,
        default=None,
        help='Bolt hole diameter in mm (for flanged hub, default: auto-sized)'
    )

    # Relief groove options (worm only)
    parser.add_argument(
        '--relief-groove',
        action='store_true',
        help='Add relief grooves at worm thread termination points (DIN 76 defaults)'
    )

    parser.add_argument(
        '--relief-groove-type',
        type=str,
        choices=['din76', 'full-radius'],
        default=None,
        help='Relief groove type: "din76" (rectangular undercut) or "full-radius" (semicircular)'
    )

    parser.add_argument(
        '--groove-width',
        type=float,
        default=None,
        help='DIN 76 groove width in mm (default: 2.5× axial pitch)'
    )

    parser.add_argument(
        '--groove-depth',
        type=float,
        default=None,
        help='DIN 76 groove depth in mm (default: 0.5× axial pitch)'
    )

    parser.add_argument(
        '--groove-radius',
        type=float,
        default=None,
        help='Full-radius groove radius in mm (default: 0.75× axial pitch)'
    )

    args = parser.parse_args()

    # Load design
    try:
        print(f"Loading design from {args.design_file}...")
        design = load_design_json(args.design_file)
    except Exception as e:
        print(f"Error loading design: {e}", file=sys.stderr)
        return 1

    # Determine what to generate
    generate_worm = not args.wheel_only
    generate_wheel = not args.worm_only

    worm = None
    wheel = None
    wheel_geo = None
    worm_rim_result = None
    wheel_rim_result = None
    worm_bore_diameter = None
    wheel_bore_diameter = None

    # Get features from JSON if present (source of truth)
    json_worm_features = design.features.worm if design.features else None
    json_wheel_features = design.features.wheel if design.features else None

    # Get manufacturing params from JSON (CLI args override)
    json_mfg = design.manufacturing
    use_globoid = args.globoid or (design.worm.type == WormType.GLOBOID if design.worm.type else False)
    # CLI profile takes precedence if explicitly provided, otherwise use JSON or default to ZA
    if args.profile is not None:
        use_profile = args.profile.upper()
    elif json_mfg and json_mfg.profile:
        use_profile = json_mfg.profile
    else:
        use_profile = 'ZA'
    use_virtual_hobbing = args.virtual_hobbing or (json_mfg.virtual_hobbing if json_mfg else False)
    use_hobbing_steps = args.hobbing_steps if args.hobbing_steps != 72 else (json_mfg.hobbing_steps if json_mfg else 72)
    use_sections = args.sections if args.sections != 36 else (json_mfg.sections_per_turn if json_mfg else 36)
    use_worm_length = args.worm_length if args.worm_length != 40.0 else (json_mfg.worm_length_mm if json_mfg and json_mfg.worm_length_mm else 40.0)
    use_wheel_width = args.wheel_width if args.wheel_width is not None else (json_mfg.wheel_width_mm if json_mfg and json_mfg.wheel_width_mm else None)
    use_trim_engagement = args.trim_to_min_engagement or (json_mfg.trim_to_min_engagement if json_mfg else False)
    # Generation method: CLI arg overrides JSON value; default 'sweep' for cylindrical
    use_generation_method = args.generation_method if args.generation_method != 'sweep' else (json_mfg.generation_method if json_mfg and json_mfg.generation_method else 'sweep')

    # Wheel tip reduction (stub teeth): CLI arg overrides JSON value
    # NOTE: The JSON tip_diameter_mm is already reduced by the calculator.
    # Only apply reduction here if the user provides a NEW value via CLI.
    if args.wheel_tip_reduction is not None:
        # User is overriding: undo any existing reduction, then apply new one
        if design.wheel.tip_reduction_mm and design.wheel.tip_reduction_mm > 0:
            design.wheel.tip_diameter_mm += design.wheel.tip_reduction_mm
            design.wheel.addendum_mm = (design.wheel.tip_diameter_mm - design.wheel.pitch_diameter_mm) / 2
        design.wheel.tip_reduction_mm = args.wheel_tip_reduction
        if args.wheel_tip_reduction > 0:
            design.wheel.tip_diameter_mm -= args.wheel_tip_reduction
            design.wheel.addendum_mm = (design.wheel.tip_diameter_mm - design.wheel.pitch_diameter_mm) / 2
    if design.wheel.tip_reduction_mm is not None and design.wheel.tip_reduction_mm > 0:
        original_tip = design.wheel.tip_diameter_mm + design.wheel.tip_reduction_mm
        print(f"  Wheel tip reduced: {original_tip:.2f}mm → {design.wheel.tip_diameter_mm:.2f}mm (stub teeth)")
    if use_trim_engagement and use_globoid:
        print(f"  Wheel OD trimmed to throat minimum (edges won't flare beyond engagement zone)")

    # Generate worm
    if generate_worm:
        worm_bore, worm_keyway, worm_ddcut, worm_set_screw, worm_bore_diameter, worm_thin_rim_warning = _determine_bore(
            args,
            pitch_diameter=design.worm.pitch_diameter_mm,
            root_diameter=design.worm.root_diameter_mm,
            json_features=json_worm_features,
            bore_cli_arg=args.worm_bore,
            ddcut_depth_arg=args.worm_ddcut_depth,
        )

        # Relief groove: CLI arg > JSON features > no groove
        worm_relief_groove = None
        json_relief_groove = json_worm_features.relief_groove if json_worm_features else None
        if args.relief_groove or json_relief_groove:
            groove_type = args.relief_groove_type or (json_relief_groove.type if json_relief_groove else "din76")
            worm_relief_groove = ReliefGrooveFeature(
                type=groove_type,
                width_mm=args.groove_width or (json_relief_groove.width_mm if json_relief_groove else None),
                depth_mm=args.groove_depth or (json_relief_groove.depth_mm if json_relief_groove else None),
                radius_mm=args.groove_radius or (json_relief_groove.radius_mm if json_relief_groove else None),
            )
            # Display groove info
            axial_pitch = design.worm.lead_mm / design.worm.num_starts
            if groove_type == "full-radius":
                r = worm_relief_groove.get_full_radius_dimensions(axial_pitch)
                print(f"  Relief groove (full-radius): radius={r:.2f}mm (axial pitch={axial_pitch:.2f}mm)")
            else:
                w, d, _ = worm_relief_groove.get_din76_dimensions(axial_pitch)
                print(f"  Relief groove (DIN 76): width={w:.2f}mm, depth={d:.2f}mm (axial pitch={axial_pitch:.2f}mm)")

        # Build description
        features_desc = _build_features_desc(worm_bore, worm_bore_diameter, worm_keyway, worm_ddcut, worm_set_screw)
        if worm_relief_groove:
            features_desc += f" + relief groove ({worm_relief_groove.type})"

        worm_type_desc = "globoid (hourglass)" if use_globoid else "cylindrical"
        profile_upper = use_profile
        if profile_upper == WormProfile.ZK or profile_upper == "ZK":
            profile_desc = "ZK/circular arc"
        elif profile_upper == WormProfile.ZI or profile_upper == "ZI":
            profile_desc = "ZI/involute"
        else:
            profile_desc = "ZA/straight"
        print(f"\nGenerating worm ({worm_type_desc}, {design.worm.num_starts}-start, module {design.worm.module_mm}mm, {profile_desc}{features_desc})...")

        profile = use_profile
        progress = CLIProgressReporter("Globoid worm") if use_globoid else None
        if use_globoid:
            worm_geo = GloboidWormGeometry(
                params=design.worm,
                assembly_params=design.assembly,
                wheel_pitch_diameter=design.wheel.pitch_diameter_mm,
                length=use_worm_length,
                sections_per_turn=use_sections,
                bore=worm_bore,
                keyway=worm_keyway,
                ddcut=worm_ddcut,
                set_screw=worm_set_screw,
                relief_groove=worm_relief_groove,
                profile=profile,
                progress_callback=progress
            )
        else:
            worm_geo = WormGeometry(
                params=design.worm,
                assembly_params=design.assembly,
                length=use_worm_length,
                sections_per_turn=use_sections,
                bore=worm_bore,
                keyway=worm_keyway,
                ddcut=worm_ddcut,
                set_screw=worm_set_screw,
                relief_groove=worm_relief_groove,
                profile=profile,
                generation_method=use_generation_method,
            )

        worm = worm_geo.build()
        print(f"  Volume: {worm.volume:.2f} mm³")

    # Generate wheel
    if generate_wheel:
        wheel_bore, wheel_keyway, wheel_ddcut, wheel_set_screw, wheel_bore_diameter, wheel_thin_rim_warning = _determine_bore(
            args,
            pitch_diameter=design.wheel.pitch_diameter_mm,
            root_diameter=design.wheel.root_diameter_mm,
            json_features=json_wheel_features,
            bore_cli_arg=args.wheel_bore,
            ddcut_depth_arg=args.wheel_ddcut_depth,
        )

        # Create hub feature: CLI arg > JSON features
        wheel_hub = None
        json_hub = json_wheel_features.hub if json_wheel_features else None
        hub_type = args.hub_type if args.hub_type != "flush" else (json_hub.type if json_hub else "flush")
        if hub_type != "flush":
            wheel_hub = HubFeature(
                hub_type=hub_type,
                length=args.hub_length or (json_hub.length_mm if json_hub else None),
                flange_diameter=args.flange_diameter or (json_hub.flange_diameter_mm if json_hub else None),
                flange_thickness=args.flange_thickness or (json_hub.flange_thickness_mm if json_hub else None),
                bolt_holes=args.flange_bolts or (json_hub.bolt_holes if json_hub else None),
                bolt_diameter=args.bolt_diameter or (json_hub.bolt_diameter_mm if json_hub else None)
            )

        # Build description
        wheel_type_desc = "hobbed (throated)" if args.hobbed else "helical"
        features_desc = _build_features_desc(wheel_bore, wheel_bore_diameter, wheel_keyway, wheel_ddcut, wheel_set_screw)

        if wheel_hub and wheel_hub.hub_type != "flush":
            hub_desc = f"{wheel_hub.hub_type} hub ({wheel_hub.length}mm"
            if wheel_hub.hub_type == "flanged":
                flange_dia = wheel_hub.flange_diameter if wheel_hub.flange_diameter else "auto"
                hub_desc += f", flange {flange_dia}mm"
                if wheel_hub.bolt_holes > 0:
                    hub_desc += f", {wheel_hub.bolt_holes} bolts"
            hub_desc += ")"
            features_desc += f", {hub_desc}"

        profile = use_profile
        if profile == WormProfile.ZK or profile == "ZK":
            profile_desc = "ZK/circular arc"
        elif profile == WormProfile.ZI or profile == "ZI":
            profile_desc = "ZI/involute"
        else:
            profile_desc = "ZA/straight"

        if use_virtual_hobbing:
            # EXPERIMENTAL: Virtual hobbing simulation
            # Pass the actual worm geometry as hob ONLY for globoid (important for accuracy)
            # For cylindrical, let VirtualHobbingWheelGeometry create a simpler hob internally
            hob_geo = worm if (worm is not None and use_globoid) else None
            hob_type = "globoid" if use_globoid else "cylindrical"
            print(f"\nGenerating wheel ({design.wheel.num_teeth} teeth, module {design.wheel.module_mm}mm, VIRTUAL HOBBING [EXPERIMENTAL], {profile_desc}{features_desc})...")
            print(f"  Using {use_hobbing_steps} hobbing steps, {hob_type} hob")
            hobbing_progress = CLIProgressReporter("Virtual hobbing")
            wheel_geo = VirtualHobbingWheelGeometry(
                params=design.wheel,
                worm_params=design.worm,
                assembly_params=design.assembly,
                face_width=use_wheel_width,
                hobbing_steps=use_hobbing_steps,
                bore=wheel_bore,
                keyway=wheel_keyway,
                ddcut=wheel_ddcut,
                set_screw=wheel_set_screw,
                hub=wheel_hub,
                profile=profile,
                hob_geometry=hob_geo,
                progress_callback=hobbing_progress,
                trim_to_min_engagement=use_trim_engagement
            )
        else:
            print(f"\nGenerating wheel ({design.wheel.num_teeth} teeth, module {design.wheel.module_mm}mm, {wheel_type_desc}, {profile_desc}{features_desc})...")
            wheel_geo = WheelGeometry(
                params=design.wheel,
                worm_params=design.worm,
                assembly_params=design.assembly,
                face_width=use_wheel_width,
                throated=args.hobbed,
                bore=wheel_bore,
                keyway=wheel_keyway,
                ddcut=wheel_ddcut,
                set_screw=wheel_set_screw,
                hub=wheel_hub,
                profile=profile,
                trim_to_min_engagement=use_trim_engagement
            )
        wheel = wheel_geo.build()
        print(f"  Volume: {wheel.volume:.2f} mm³")

    # Measure rim thickness
    if worm is not None and worm_bore_diameter is not None:
        worm_rim_result = measure_rim_thickness(
            part=worm,
            bore_diameter_mm=worm_bore_diameter,
            is_worm=True
        )
        if worm_rim_result is not None:
            print(f"\nWorm rim thickness: {worm_rim_result.minimum_thickness_mm:.2f} mm")
            if worm_rim_result.has_warning:
                print(f"  WARNING: Below {WORM_RIM_WARNING_THRESHOLD_MM}mm threshold")

    if wheel is not None and wheel_bore_diameter is not None:
        wheel_rim_result = measure_rim_thickness(
            part=wheel,
            bore_diameter_mm=wheel_bore_diameter,
            is_worm=False
        )
        if wheel_rim_result is not None:
            print(f"Wheel rim thickness: {wheel_rim_result.minimum_thickness_mm:.2f} mm")
            if wheel_rim_result.has_warning:
                print(f"  WARNING: Below {WHEEL_RIM_WARNING_THRESHOLD_MM}mm threshold")

    # Calculate mesh alignment (when both parts generated)
    mesh_alignment_result = None
    skip_mesh = args.skip_mesh_alignment

    # Check if wheel was pre-aligned during virtual hobbing
    wheel_pre_aligned = hasattr(wheel_geo, 'pre_alignment_deg') if wheel_geo else False

    if worm is not None and wheel is not None and not skip_mesh:
        if wheel_pre_aligned:
            # Virtual hobbing already aligned the wheel - use that result
            print(f"\nMesh alignment (from virtual hobbing):")
            print(f"  Pre-aligned rotation: {wheel_geo.pre_alignment_deg:.2f}°")
            print(f"  Interference volume: {wheel_geo.pre_alignment_interference_mm3:.4f} mm³")
            print(f"  Status: Pre-aligned during hobbing (no additional rotation needed)")
            # Create a result object for consistency
            mesh_alignment_result = MeshAlignmentResult(
                optimal_rotation_deg=0.0,  # Already applied
                interference_volume_mm3=wheel_geo.pre_alignment_interference_mm3,
                within_tolerance=wheel_geo.pre_alignment_interference_mm3 < 0.1,
                tooth_pitch_deg=360.0 / design.wheel.num_teeth,
                worm_position=(design.assembly.centre_distance_mm, 0.0, 0.0),
                message="Pre-aligned during virtual hobbing (hobbing uses +X approach)"
            )
        else:
            # Standard mesh alignment calculation
            print(f"\nCalculating mesh alignment...")
            mesh_alignment_result = find_optimal_mesh_rotation(
                wheel=wheel,
                worm=worm,
                centre_distance_mm=design.assembly.centre_distance_mm,
                num_teeth=design.wheel.num_teeth,
                module_mm=design.worm.module_mm,
            )
            print(f"  Optimal wheel rotation: {mesh_alignment_result.optimal_rotation_deg:.2f}°")
            print(f"  Interference volume: {mesh_alignment_result.interference_volume_mm3:.4f} mm³")
            print(f"  Tolerance: {mesh_alignment_result.tolerance_mm3:.1f} mm³ (module-scaled)")
            print(f"  Quality: {mesh_alignment_result.mesh_quality}")
            print(f"  Status: {mesh_alignment_result.message}")

    # Export all output files (STEP, 3MF, STL, design.json, design.md)
    pkg = None
    if not args.no_save:
        # Use pre-computed mesh rotation if available
        pre_rotation = None
        if mesh_alignment_result is not None:
            pre_rotation = mesh_alignment_result.optimal_rotation_deg

        print(f"\nExporting files...")
        pkg = generate_package(
            design=design,
            worm=worm,
            wheel=wheel,
            mesh_rotation_deg=pre_rotation,
            virtual_hobbing=use_virtual_hobbing,
            include_3mf=not args.no_3mf,
            include_stl=not args.no_stl,
            include_assembly=not args.no_3mf,
            log=print,
        )

    # Save geometry analysis JSON (mesh alignment + rim measurements)
    if not args.no_save:
        output_dir = Path(args.output_dir)
        output_dir.mkdir(parents=True, exist_ok=True)
        import json

        has_analysis = (mesh_alignment_result is not None or
                        worm_rim_result is not None or
                        wheel_rim_result is not None or
                        use_globoid)
        if has_analysis:
            from datetime import datetime
            analysis_file = output_dir / f"geometry_analysis_m{design.worm.module_mm}.json"

            analysis_data = {
                "design_info": {
                    "module_mm": design.worm.module_mm,
                    "ratio": design.assembly.ratio,
                    "centre_distance_mm": design.assembly.centre_distance_mm,
                    "worm_starts": design.worm.num_starts,
                    "wheel_teeth": design.wheel.num_teeth,
                },
                "analysis_timestamp": datetime.now().isoformat(),
            }

            # Add mesh alignment data
            if mesh_alignment_result is not None:
                analysis_data["mesh_alignment"] = mesh_alignment_to_dict(mesh_alignment_result)

            # Add rim thickness data
            if worm_rim_result is not None or wheel_rim_result is not None:
                analysis_data["rim_thickness"] = {}
                if worm_rim_result is not None:
                    analysis_data["rim_thickness"]["worm"] = rim_thickness_to_dict(worm_rim_result)
                if wheel_rim_result is not None:
                    analysis_data["rim_thickness"]["wheel"] = rim_thickness_to_dict(wheel_rim_result)

            # Add throat geometry for globoid worms
            if use_globoid and design.worm.throat_reduction_mm:
                arc_r = design.worm.tip_diameter_mm / 2 - design.worm.throat_reduction_mm
                cd = design.assembly.centre_distance_mm
                margin = design.worm.addendum_mm + 0.5 * design.wheel.addendum_mm
                min_blank_r = cd - arc_r + margin
                tip_r = design.wheel.tip_diameter_mm / 2
                throat_od = round(2 * min(tip_r, min_blank_r), 3)
                analysis_data["throat_geometry"] = {
                    "wheel_throat_od_mm": throat_od,
                    "wheel_tip_od_mm": design.wheel.tip_diameter_mm,
                    "od_difference_mm": round(design.wheel.tip_diameter_mm - throat_od, 3),
                    "note": "wheel_throat_od_mm is the minimum OD at the engagement zone; edges can be clipped to this without losing engagement"
                }
                print(f"\nWheel throat OD: {throat_od:.2f}mm (tip OD {design.wheel.tip_diameter_mm:.2f}mm, {design.wheel.tip_diameter_mm - throat_od:.2f}mm recoverable at edges)")

            if analysis_file.exists():
                analysis_file.unlink()
            with open(analysis_file, 'w') as f:
                json.dump(analysis_data, f, indent=2)
            print(f"  Saved: {analysis_file}")

    # Write package files to disk (or ZIP)
    if not args.no_save and pkg is not None:
        if args.zip_output:
            zip_path = Path(args.zip_output)
            zip_data = create_package_zip(pkg, design)
            zip_path.write_bytes(zip_data)
            print(f"\n  Saved ZIP: {zip_path} ({len(zip_data) / 1024:.1f} KB)")
        else:
            output_dir = Path(args.output_dir)
            written = save_package_to_dir(pkg, output_dir, design)
            for p in written:
                print(f"  Saved: {p}")

    # View in OCP viewer
    if args.view:
        try:
            from ocp_vscode import show

            # Position parts for mesh visualization
            if wheel is not None and worm is not None:
                # Use calculated mesh alignment for optimal positioning
                if mesh_alignment_result is not None:
                    wheel_positioned, worm_positioned = position_for_mesh(
                        wheel=wheel,
                        worm=worm,
                        centre_distance_mm=design.assembly.centre_distance_mm,
                        rotation_deg=mesh_alignment_result.optimal_rotation_deg,
                    )
                    print(f"\n  Wheel rotated {mesh_alignment_result.optimal_rotation_deg:.2f}° for optimal mesh")
                else:
                    # Fallback: position without alignment calculation
                    from build123d import Rot, Pos
                    wheel_positioned = wheel
                    worm_positioned = Pos(design.assembly.centre_distance_mm, 0, 0) * Rot(X=90) * worm

                show(wheel_positioned, worm_positioned, names=["wheel", "worm"], colors=["steelblue", "orange"])
            elif wheel is not None:
                show(wheel, names=["wheel"], colors=["steelblue"])
            elif worm is not None:
                show(worm, names=["worm"], colors=["orange"])
            print("Displayed in OCP viewer")

        except ImportError:
            print("\nWarning: ocp_vscode not available for viewing", file=sys.stderr)
            print("Install with: pip install ocp_vscode", file=sys.stderr)

    # Summary
    profile_upper = use_profile
    if profile_upper == WormProfile.ZK or profile_upper == "ZK":
        profile_name = "ZK (circular arc, 3D printing)"
    elif profile_upper == WormProfile.ZI or profile_upper == "ZI":
        profile_name = "ZI (involute, hobbing)"
    else:
        profile_name = "ZA (straight, CNC)"
    print(f"\nDesign summary:")
    print(f"  Ratio: {design.assembly.ratio}:1")
    print(f"  Centre distance: {design.assembly.centre_distance_mm:.2f} mm")
    print(f"  Pressure angle: {design.assembly.pressure_angle_deg}°")
    print(f"  Tooth profile: {profile_name} (DIN 3975)")

    # Bore/keyway summary with override hints
    if not args.no_bore:
        has_any_bore = (generate_worm and worm_bore_diameter) or (generate_wheel and wheel_bore_diameter)
        if has_any_bore:
            print(f"\nBore & keyway:")
            if generate_worm:
                if worm_bore_diameter:
                    keyway_dims = get_din_6885_keyway(worm_bore_diameter)
                    keyway_info = ""
                    if keyway_dims and not args.no_keyway:
                        keyway_info = f", keyway {keyway_dims[0]}x{keyway_dims[2]}mm (DIN 6885)"
                    elif worm_bore_diameter < 6.0:
                        keyway_info = " (no keyway - below DIN 6885 range)"
                    override_note = "" if args.worm_bore else " (override: --worm-bore X)"
                    worm_rim = (design.worm.root_diameter_mm - worm_bore_diameter) / 2
                    print(f"  Worm:  {worm_bore_diameter}mm{keyway_info}, rim {worm_rim:.2f}mm{override_note}")
                else:
                    print(f"  Worm:  solid (too small for bore)")

            if generate_wheel:
                if wheel_bore_diameter:
                    keyway_dims = get_din_6885_keyway(wheel_bore_diameter)
                    keyway_info = ""
                    if keyway_dims and not args.no_keyway:
                        keyway_info = f", keyway {keyway_dims[0]}x{keyway_dims[3]}mm (DIN 6885)"
                    elif wheel_bore_diameter < 6.0:
                        keyway_info = " (no keyway - below DIN 6885 range)"
                    override_note = "" if args.wheel_bore else " (override: --wheel-bore X)"
                    wheel_rim = (design.wheel.root_diameter_mm - wheel_bore_diameter) / 2
                    print(f"  Wheel: {wheel_bore_diameter}mm{keyway_info}, rim {wheel_rim:.2f}mm{override_note}")
                else:
                    print(f"  Wheel: solid (too small for bore)")

            # Print warnings for thin rims
            if (generate_worm and worm_thin_rim_warning) or (generate_wheel and wheel_thin_rim_warning):
                print(f"\n  Warning: thin rim on small bore - handle with care")

            print(f"\n  To generate solid parts: --no-bore")

    # Save extended JSON with manufacturing features
    if args.save_json:
        # Collect worm manufacturing features
        worm_features = None
        if worm_bore_diameter or worm_set_screw:
            keyway_dims = get_din_6885_keyway(worm_bore_diameter) if worm_bore_diameter else None
            worm_features = ManufacturingFeatures(
                bore_diameter=worm_bore_diameter,
                keyway_width=keyway_dims[0] if keyway_dims and worm_keyway else None,
                keyway_depth=keyway_dims[2] if keyway_dims and worm_keyway else None,
                set_screw_size=worm_set_screw.size if worm_set_screw else None,
                set_screw_count=worm_set_screw.count if worm_set_screw else None
            )

        # Collect wheel manufacturing features
        wheel_features = None
        if wheel_bore_diameter or wheel_set_screw or wheel_hub:
            keyway_dims = get_din_6885_keyway(wheel_bore_diameter) if wheel_bore_diameter else None
            wheel_features = ManufacturingFeatures(
                bore_diameter=wheel_bore_diameter,
                keyway_width=keyway_dims[0] if keyway_dims and wheel_keyway else None,
                keyway_depth=keyway_dims[3] if keyway_dims and wheel_keyway else None,
                set_screw_size=wheel_set_screw.size if wheel_set_screw else None,
                set_screw_count=wheel_set_screw.count if wheel_set_screw else None,
                hub_type=wheel_hub.hub_type if wheel_hub else None,
                hub_length=wheel_hub.length if wheel_hub else None,
                flange_diameter=wheel_hub.flange_diameter if wheel_hub else None,
                flange_thickness=wheel_hub.flange_thickness if wheel_hub else None,
                flange_bolts=wheel_hub.bolt_holes if wheel_hub else None,
                bolt_diameter=wheel_hub.bolt_diameter if wheel_hub else None
            )

        # Create manufacturing parameters
        manufacturing = ManufacturingParams(
            worm_type="globoid" if use_globoid else "cylindrical",
            worm_length=use_worm_length,
            wheel_width=use_wheel_width,
            wheel_throated=args.hobbed,
            profile=use_profile,
            worm_features=worm_features,
            wheel_features=wheel_features
        )

        # Create measured geometry from rim thickness results
        measured_geometry = None
        if worm_rim_result is not None or wheel_rim_result is not None:
            from datetime import datetime
            measured_geometry = MeasuredGeometry(
                wheel_rim_thickness_mm=wheel_rim_result.minimum_thickness_mm if wheel_rim_result else None,
                wheel_measurement_point=MeasurementPoint(
                    x_mm=wheel_rim_result.measurement_point_bore[0],
                    y_mm=wheel_rim_result.measurement_point_bore[1],
                    z_mm=wheel_rim_result.measurement_point_bore[2]
                ) if wheel_rim_result and wheel_rim_result.measurement_point_bore else None,
                wheel_rim_warning=wheel_rim_result.has_warning if wheel_rim_result else None,
                worm_rim_thickness_mm=worm_rim_result.minimum_thickness_mm if worm_rim_result else None,
                worm_measurement_point=MeasurementPoint(
                    x_mm=worm_rim_result.measurement_point_bore[0],
                    y_mm=worm_rim_result.measurement_point_bore[1],
                    z_mm=worm_rim_result.measurement_point_bore[2]
                ) if worm_rim_result and worm_rim_result.measurement_point_bore else None,
                worm_rim_warning=worm_rim_result.has_warning if worm_rim_result else None,
                measurement_timestamp=datetime.now().isoformat()
            )

        # Create complete design with manufacturing parameters
        complete_design = WormGearDesign(
            worm=design.worm,
            wheel=design.wheel,
            assembly=design.assembly,
            manufacturing=manufacturing,
            measured_geometry=measured_geometry
        )

        # Save to file
        output_path = Path(args.save_json)
        save_design_json(complete_design, output_path)
        print(f"\nSaved extended JSON: {output_path}")
        print(f"  This JSON includes all manufacturing features and can fully reproduce this design")

    return 0


if __name__ == '__main__':
    sys.exit(main())
