import math
import argparse
import sys

# --- config defaults ---
# Cutting on the RIGHT side of the blank:
# - X0 is "just touching" the OD on the right side.
# - Deeper cuts go toward NEGATIVE X.
X_SAFE = 3.0           # mm, retract away from blank (positive X)
X_RELIEF = 1.0         # mm, tiny retract away from walls between passes (+X direction)
F_RAPID = 1000.0       # mm/min, retract and non-cut moves
SPINDLE_RPM = 8000     # rpm
FINAL_STEP = 0.05      # mm, final skim step size target
B_DEG_PER_SEC = 24.0   # deg/sec, B axis rotation speed (for time estimate only)

def pick_cutter_number(z_teeth):
    # Your ranges:
    # #1 12-13, #2 14-16, #3 17-20, #4 21-25, #5 26-34, #6 35-54, #7 55-134, #8 135+
    if 12 <= z_teeth <= 13:
        return 1
    if 14 <= z_teeth <= 16:
        return 2
    if 17 <= z_teeth <= 20:
        return 3
    if 21 <= z_teeth <= 25:
        return 4
    if 26 <= z_teeth <= 34:
        return 5
    if 35 <= z_teeth <= 54:
        return 6
    if 55 <= z_teeth <= 134:
        return 7
    if z_teeth >= 135:
        return 8
    raise ValueError("Tooth count Z is outside the cutter set ranges (Z must be >= 12).")

def build_infeed_positions(x_final, max_step, final_step=FINAL_STEP):
    # Returns a list of positive infeed depths (> 0), ending at x_final.
    # Last increment is exactly final_step (if possible).
    # All increments before the last are exactly max_step, except possibly the very first.
    if x_final <= 0.0:
        return []
    if max_step <= 0.0:
        raise ValueError("max_step must be > 0.")
    if final_step <= 0.0:
        raise ValueError("final_step must be > 0.")

    if x_final <= final_step:
        return [x_final]

    target = x_final - final_step
    k = int(math.floor(target / max_step))
    remainder = target - k * max_step

    xs = []
    eps = 1e-9

    if remainder > eps:
        xs.append(remainder)
        for _ in range(k):
            xs.append(xs[-1] + max_step)
    else:
        if target > eps:
            steps = int(round(target / max_step))
            xs.append(max_step)
            for _ in range(steps - 1):
                xs.append(xs[-1] + max_step)

    # Numerical safety to land on target
    if target > eps:
        if not xs:
            xs = [target]
        elif abs(xs[-1] - target) > 1e-6:
            xs.append(target)

    # Final skim to x_final
    xs.append(x_final)

    # Clean duplicates and non-positive
    cleaned = []
    for v in xs:
        if v <= 0.0:
            continue
        if cleaned and abs(cleaned[-1] - v) < 1e-6:
            continue
        cleaned.append(v)

    return cleaned

def g1(axis_words, feed=None, comment=None):
    line = "G1 " + " ".join(axis_words)
    if feed is not None:
        line += " F{:.3f}".format(feed)
    if comment:
        line += " ; " + comment
    return line

def fmt(val):
    return "{:.4f}".format(val)

def die(msg):
    print("Error: " + msg, file=sys.stderr)
    sys.exit(2)

def parse_args():
    p = argparse.ArgumentParser(
        description="Gear cutting G-code generator for Snapmaker 2.0, by Illusionmanager 2026"
    )

    # Geometry
    p.add_argument("-m", "-M", dest="m", type=float, required=True,
                   help="Module (required), e.g. 0.5")
    p.add_argument("-z", "-Z", dest="z", type=int,
                   help="Number of teeth (Z)")
    p.add_argument("-r", "-R", dest="r", type=float,
                   help="Pitch radius in mm")
    p.add_argument("-o", "-O", dest="o", type=float,
                   help="Blank outside diameter (OD) in mm")

    # Process
    p.add_argument("-w", "-W", dest="w", type=float, required=True,
                   help="Face width W in mm")
    p.add_argument("-f", "-F", dest="f", type=float, required=True,
                   help="Cutting feed in mm/min (Y stroke)")
    p.add_argument("-s", "-S", dest="s", type=float, required=True,
                   help="Max step size in X in mm")

    # Y convention:
    # Y0 is at the FRONT face. Front is toward negative Y.
    # Cutting stroke goes from front to back, so:
    #   y_start = -y_extra
    #   y_end   = W + y_extra
    p.add_argument("-e", "-E", dest="y_extra", type=float, default=0.5,
                   help="Extra Y overshoot past each face edge in mm (default 0.5)")

    # Output / mode
    p.add_argument("--cut", "--CUT", dest="cut", action="store_true",
                   help="Enable real cutting (spindle on, real X infeeds). Default is DRY_RUN.")
    p.add_argument("--rpm", "--RPM", dest="rpm", type=int, default=SPINDLE_RPM,
                   help="Spindle RPM (default 8000)")
    p.add_argument("--out", "--OUT", dest="out", type=str, default="",
                   help="Output filename (.cnc). Default auto name.")
    p.add_argument("--cal", "--CAL", dest="cal", action="store_true",
                   help="Calibration cut: stop after the very first cut stroke")

    return p.parse_args()

def derive_geometry(module, z_teeth, pitch_radius, blank_od):
    provided = [z_teeth is not None, pitch_radius is not None, blank_od is not None]
    if sum(1 for v in provided if v) != 1:
        die("Provide exactly one of -z (teeth), -r (pitch radius), or -o (OD).")

    m = module
    if m <= 0.0:
        die("Module must be > 0.")

    if z_teeth is None:
        if pitch_radius is not None:
            z_float = (2.0 * pitch_radius) / m
        else:
            z_float = (blank_od / m) - 2.0

        z_round = int(round(z_float))
        if abs(z_float - z_round) > 1e-6:
            z_floor = int(math.floor(z_float))
            z_ceil = int(math.ceil(z_float))

            def pitch_r_for_z(z):
                return 0.5 * m * float(z)

            def od_for_z(z):
                return m * (float(z) + 2.0)

            r_floor = pitch_r_for_z(z_floor)
            r_ceil = pitch_r_for_z(z_ceil)
            od_floor = od_for_z(z_floor)
            od_ceil = od_for_z(z_ceil)

            msg = []
            msg.append("Computed Z is not an integer (got {:.6f}).".format(z_float))
            msg.append("Your inputs imply Z = (2*r)/m (or Z = OD/m - 2), but tooth count must be an integer.")
            if pitch_radius is not None:
                msg.append("Given m={:.6f} and r={:.6f}, Z would be {:.6f}.".format(m, pitch_radius, z_float))
            else:
                msg.append("Given m={:.6f} and OD={:.6f}, Z would be {:.6f}.".format(m, blank_od, z_float))
            msg.append("")
            msg.append("Fix options:")
            msg.append("1) Use an integer tooth count near that value:")
            msg.append("   - Z={} -> pitch radius r={:.6f} mm, blank OD={:.6f} mm".format(z_floor, r_floor, od_floor))
            msg.append("   - Z={} -> pitch radius r={:.6f} mm, blank OD={:.6f} mm".format(z_ceil,  r_ceil,  od_ceil))
            if pitch_radius is not None and z_floor > 0 and z_ceil > 0:
                msg.append("2) If r={:.6f} mm is correct, adjust module to make Z an integer:".format(pitch_radius))
                msg.append("   - For Z={} -> module m = (2*r)/Z = {:.6f}".format(z_floor, (2.0 * pitch_radius) / float(z_floor)))
                msg.append("   - For Z={} -> module m = (2*r)/Z = {:.6f}".format(z_ceil,  (2.0 * pitch_radius) / float(z_ceil)))
            msg.append("3) Or provide -z directly if you already know the tooth count, and ignore -r/-o.")
            die("\n".join(msg))

        z_teeth = z_round

    if z_teeth < 1:
        die("Z must be >= 1.")

    pitch_radius = 0.5 * m * float(z_teeth)
    blank_od = m * (float(z_teeth) + 2.0)

    return z_teeth, pitch_radius, blank_od

def format_seconds(seconds):
    if seconds < 0:
        seconds = 0
    s = int(round(seconds))
    h = s // 3600
    s -= h * 3600
    m = s // 60
    s -= m * 60
    return "{:02d}:{:02d}:{:02d}".format(h, m, s)

def estimate_total_time_seconds(z_teeth, xs, y_start, y_end, feed_cut, dry_run, spindle_dwell_s=1.0):
    # Estimate based on the same sequence of moves we generate.
    x = 0.0
    y = 0.0
    b = 0.0
    total = 0.0

    def move_x(x_new, feed_mm_min):
        nonlocal x, total
        dist = abs(x_new - x)
        total += (dist / feed_mm_min) * 60.0
        x = x_new

    def move_y(y_new, feed_mm_min):
        nonlocal y, total
        dist = abs(y_new - y)
        total += (dist / feed_mm_min) * 60.0
        y = y_new

    def move_b(b_new):
        nonlocal b, total
        dist = abs(b_new - b)
        total += dist / B_DEG_PER_SEC
        b = b_new

    # Retract X before anything else
    if dry_run:
        move_x(X_SAFE + X_RELIEF, F_RAPID)
    else: 
        move_x(X_SAFE, F_RAPID)

    # Dwell after spindle start exists only in CUT mode (matches generator)
    if not dry_run:
        total += spindle_dwell_s

    # Move to Y start
    move_y(y_start, F_RAPID)

    b_step = 360.0 / float(z_teeth)

    # Backlash take-up: move one tooth forward, no cutting
    move_b(b + b_step)

    # For each tooth: cut at current B, then step B forward for next tooth
    for i in range(z_teeth):
        # Ensure safe retract and Y start at each tooth (matches generator)
        move_x(X_SAFE, F_RAPID)
        move_y(y_start, F_RAPID)

        for x_pass in xs:
            x_cmd = X_SAFE if dry_run else -x_pass
            move_x(x_cmd, feed_cut)

            # Cut stroke
            move_y(y_end, feed_cut)

            # Relief
            if dry_run:
                x_rel = X_SAFE + X_RELIEF
            else:
                x_rel = -x_pass + X_RELIEF
                if x_rel > X_SAFE:
                    x_rel = X_SAFE

            move_x(x_rel, F_RAPID)

            # Return
            move_y(y_start, F_RAPID)

        # Index to next tooth (forward only), except after last
        if i != z_teeth - 1:
            move_x(X_SAFE, F_RAPID)
            move_b(b + b_step)

    # End: retract and optionally park Y to 0 (we do NOT command B back to 0)
    move_x(X_SAFE, F_RAPID)
    move_y(0.0, F_RAPID)

    return total

def main():
    args = parse_args()
    dry_run = (not args.cut)
    cal_mode = bool(args.cal)

    z_teeth, pitch_radius, blank_od = derive_geometry(
        module=args.m,
        z_teeth=args.z,
        pitch_radius=args.r,
        blank_od=args.o
    )

    module = args.m
    face_w = args.w
    feed_cut = args.f
    max_step = args.s
    spindle_rpm = int(args.rpm)
    y_extra = float(args.y_extra)

    cutter_no = pick_cutter_number(z_teeth)

    # Standard full-depth whole depth (no profile shift)
    x_final = 2.25 * module

    # Y0 is at front face, and front is toward negative Y.
    # Cut only in the "good" direction: front -> back (Y increasing).
    y_start = -y_extra
    y_end = face_w + y_extra

    xs = build_infeed_positions(x_final, max_step, final_step=FINAL_STEP)

    outside_radius = blank_od * 0.5
    root_radius = pitch_radius - 1.25 * module  # standard dedendum

    # Only estimate time if not in calibration mode
    est_s = 0.0
    if not cal_mode:
        est_s = estimate_total_time_seconds(
            z_teeth=z_teeth,
            xs=xs,
            y_start=y_start,
            y_end=y_end,
            feed_cut=feed_cut,
            dry_run=dry_run,
            spindle_dwell_s=1.0
        )

    print("")
    print("Inputs / derived geometry")
    print("  module m (mm): {:.4f}".format(module))
    print("  teeth Z: {}".format(z_teeth))
    print("  pitch radius (mm): {:.4f}".format(pitch_radius))
    print("  outside diameter OD (mm): {:.4f}".format(blank_od))
    print("  outside radius (mm): {:.4f}".format(outside_radius))
    print("  root (valley) radius (mm): {:.4f}".format(root_radius))
    print("")
    print("Cutter")
    print("  recommended cutter number: #{}".format(cutter_no))
    print("")
    print("Cut plan")
    print("  face width W (mm): {:.4f}".format(face_w))
    print("  y_extra (mm): {:.4f}".format(y_extra))
    print("  Y stroke: start={:.4f}, end={:.4f} (front -> back)".format(y_start, y_end))
    print("  cutting feed (mm/min): {:.3f}".format(feed_cut))
    print("  max step size X (mm): {:.4f}".format(max_step))
    print("  final skim (mm): {:.4f}".format(FINAL_STEP))
    print("  x_final (whole depth) (mm): {:.4f}".format(x_final))
    print("  X_SAFE (mm): {:.4f}".format(X_SAFE))
    print("  X_RELIEF (mm): {:.4f}".format(X_RELIEF))
    print("  rapids (mm/min): {:.3f}".format(F_RAPID))
    print("  spindle rpm: {}".format(spindle_rpm))
    print("  B speed (deg/sec): {:.3f}".format(B_DEG_PER_SEC))
    print("  mode: {}".format("DRY_RUN (spindle off, X forced to X_SAFE)" if dry_run else "CUT (spindle on, real X infeeds)"))
    print("  passes per tooth: {}".format(len(xs)))
    print("")
    if cal_mode:
        print("  calibration mode: enabled (will stop after first cut stroke)")
    else:
        print("Estimated total time: {}".format(format_seconds(est_s)))
    print("")

    out_name = args.out.strip()
    if not out_name:
        out_name = "gear_Z{}_m{:.3f}.cnc".format(z_teeth, module)
    if not out_name.lower().endswith(".cnc"):
        out_name += ".cnc"

    lines = []

    lines.append("; Gear cutting with form cutter (module {:.3f}, Z {})".format(module, z_teeth))
    lines.append("; Pitch radius: {:.4f} mm".format(pitch_radius))
    lines.append("; Outside diameter OD: {:.4f} mm".format(blank_od))
    lines.append("; Outside radius: {:.4f} mm".format(outside_radius))
    lines.append("; Root (valley) radius: {:.4f} mm (pitch_radius - 1.25*m)".format(root_radius))
    lines.append("; Recommended cutter number: #{}".format(cutter_no))
    lines.append("; Whole depth (x_final): {:.4f} mm (2.25*m)".format(x_final))
    lines.append("; Y0 is front face; front is toward negative Y; cut stroke is front -> back")
    lines.append("; Y stroke: start={:.4f}, end={:.4f}, y_extra={:.4f}".format(y_start, y_end, y_extra))
    lines.append("; Cutting on RIGHT side: X0 touch on right OD; deeper cuts go toward negative X")
    lines.append("; Backlash take-up: one dummy tooth step before cutting; then forward-only B steps")
    if cal_mode:
        lines.append("; CAL mode: stop after first cut stroke, then park and end")
    else:
        lines.append("; Estimated total time: {} (B speed {:.2f} deg/sec)".format(format_seconds(est_s), B_DEG_PER_SEC))
    if dry_run:
        lines.append("; DRY_RUN enabled: spindle off; X toggles between X_SAFE+X_RELIEF and X_SAFE (air cut)")

    lines.append("G21 ; mm")
    lines.append("G90 ; absolute")
    lines.append("")

    # Start behavior
    if not dry_run:
        lines.append(g1(["X" + fmt(X_SAFE)], feed=F_RAPID, comment="retract before spindle start"))
        lines.append("M3 S{} ; spindle on clockwise".format(spindle_rpm))
        lines.append("G4 P1 ; dwell 1s")
    else:
        lines.append("; DRY_RUN enabled: spindle remains off, X toggles between X_SAFE+X_RELIEF and X_SAFE")
        lines.append(g1(["X" + fmt(X_SAFE + X_RELIEF)], feed=F_RAPID, comment="dry run start at relief"))
    lines.append("")

    # Move to Y start
    lines.append(g1(["Y" + fmt(y_start)], feed=F_RAPID, comment="go to Y start"))
    lines.append("")

    b_step = 360.0 / float(z_teeth)

    # Backlash take-up: move one tooth forward before any cutting
    b_pos = 0.0
    b_pos += b_step
    lines.append(g1(["B" + fmt(b_pos)], feed=F_RAPID, comment="backlash take-up (one tooth)"))
    lines.append("")

    for i in range(z_teeth):
        lines.append("; --- tooth index {} / {} ---".format(i + 1, z_teeth))

        # Always retract and go to Y start before cutting
        lines.append(g1(["X" + fmt(X_SAFE)], feed=F_RAPID, comment="safe retract"))
        lines.append(g1(["Y" + fmt(y_start)], feed=F_RAPID, comment="ensure at Y start"))

        for pass_idx, x_pass in enumerate(xs):
            # Right-side cutting: machine X is negative for cutting depths
            x_cmd = X_SAFE if dry_run else -x_pass

            # Infeed (relief -> cut) at cutting speed
            lines.append(g1(
                ["X" + fmt(x_cmd)],
                feed=feed_cut,
                comment="infeed pass {} x={}".format(pass_idx + 1, fmt(x_cmd))
            ))

            # Cut stroke
            lines.append(g1(
                ["Y" + fmt(y_end)],
                feed=feed_cut,
                comment="cut stroke (front -> back)"
            ))

            if cal_mode and i == 0 and pass_idx == 0:
                lines.append(g1(["X" + fmt(X_SAFE)], feed=F_RAPID, comment="CAL: retract X"))
                lines.append(g1(["Y" + fmt(0.0)], feed=F_RAPID, comment="CAL: park Y0"))
                if not dry_run:
                    lines.append("M5 ; spindle off")
                lines.append("; CAL mode end")
                lines.append("M2")

                with open(out_name, "w", encoding="utf-8") as f:
                    f.write("\n".join(lines) + "\n")

                print("Wrote: {}".format(out_name))
                print("Reminder: Set Z height manually (this program does not change Z).")
                print("Reminder: Set X0 by touching the right side of a blank with a diameter of {:.2f} mm".format(blank_od))
                print("Reminder: Set Y0 at the front face.")
                if dry_run:
                    print("DRY_RUN reminder: No spindle commands. X toggles between X_SAFE+X_RELIEF and X_SAFE (air cut).")
                else:
                    print("CUT reminder: Spindle will start at {} rpm. Verify clearances first.".format(spindle_rpm))
                return

            # Relief (cut -> relief) at rapid
            if dry_run:
                x_rel = X_SAFE + X_RELIEF
            else:
                x_rel = -x_pass + X_RELIEF
                if x_rel > X_SAFE:
                    x_rel = X_SAFE

            lines.append(g1(
                ["X" + fmt(x_rel)],
                feed=F_RAPID,
                comment="relief retract"
            ))

            # Return
            lines.append(g1(
                ["Y" + fmt(y_start)],
                feed=F_RAPID,
                comment="return"
            ))

        # Index to next tooth (forward only), except after last tooth
        if i != z_teeth - 1:
            lines.append(g1(["X" + fmt(X_SAFE)], feed=F_RAPID, comment="safe retract before B step"))
            b_pos += b_step
            lines.append(g1(["B" + fmt(b_pos)], feed=F_RAPID, comment="index B +1 tooth"))

        lines.append("")

    # End: retract and park Y to 0. Do NOT command B back to 0.
    lines.append(g1(["X" + fmt(X_SAFE)], feed=F_RAPID, comment="final retract"))
    if not dry_run:
        lines.append("M5 ; spindle off")
    else:
        lines.append("; DRY_RUN enabled: spindle was not started")
    lines.append(g1(["Y" + fmt(0.0)], feed=F_RAPID, comment="park Y0 (front face reference)"))
    lines.append("; End")
    lines.append("M2")

    with open(out_name, "w", encoding="utf-8") as f:
        f.write("\n".join(lines) + "\n")

    print("Wrote: {}".format(out_name))
    print("Reminder: Set Z height manually (this program does not change Z).")
    print("Reminder: Set X0 by touching the right side of a blank with a diameter of {:.2f} mm".format(blank_od))
    print("Reminder: Set Y0 at the front face.")
    if dry_run:
        print("DRY_RUN reminder: No spindle commands. X toggles between X_SAFE+X_RELIEF and X_SAFE (air cut).")
    else:
        print("CUT reminder: Spindle will start at {} rpm. Verify clearances first.".format(spindle_rpm))

if __name__ == "__main__":
    main()