"""Set the current physical position of each motor as its middle position (2048).

Usage: python calibrate_motors.py 1 2 3
  - Manually position each motor where you want its center to be
  - Run this script with the motor IDs to calibrate
  - Each motor's current position will become 2048 (middle)
"""
import sys, time

from scservo_sdk import PortHandler
import scservo_patch
from scservo_sdk.sms_sts import sms_sts

USB_PORT = "/dev/tty.usbserial-0001"
BAUD_RATE = 115200
INST_OFSCAL = 11  # offset calibration instruction

PKT_ID = 2
PKT_LENGTH = 3
PKT_INSTRUCTION = 4


def offset_calibrate(sts, motor_id):
    """Send INST_OFSCAL: sets current position as middle (2048)."""
    txpacket = [0] * 6
    txpacket[PKT_ID] = motor_id
    txpacket[PKT_LENGTH] = 2
    txpacket[PKT_INSTRUCTION] = INST_OFSCAL
    _, result, error = sts.txRxPacket(sts.portHandler, txpacket)
    return result, error


def main():
    if len(sys.argv) < 2:
        sys.exit("Usage: python calibrate_motors.py <motor_id> [motor_id ...]")

    motor_ids = [int(x) for x in sys.argv[1:]]

    port = PortHandler(USB_PORT)
    if not port.openPort():
        sys.exit("Failed to open port")
    if not port.setBaudRate(BAUD_RATE):
        sys.exit("Failed to set baud rate")

    sts = sms_sts(port)

    try:
        # Verify all motors are reachable
        for mid in motor_ids:
            _, comm, err = sts.ping(mid)
            if comm != 0 or err != 0:
                sys.exit(f"Motor {mid} not responding (comm={comm}, err={err})")
            pos, _, _ = sts.ReadPos(mid)
            print(f"Motor {mid}: current position = {pos}")

        print(f"\nThis will set the current position of motors {motor_ids} as middle (2048).")
        confirm = input("Proceed? [y/N] ").strip().lower()
        if confirm != "y":
            print("Aborted.")
            return

        for mid in motor_ids:
            for attempt in range(3):
                result, error = offset_calibrate(sts, mid)
                if result == 0:
                    break
                time.sleep(0.1)
            if result != 0:
                print(f"Motor {mid}: FAILED after 3 attempts (comm={result}, err={error})")
                continue

            # After OFSCAL the motor may take a moment to respond; retry the read
            # and tolerate the SDK's IndexError on truncated replies.
            pos = None
            for read_attempt in range(5):
                time.sleep(0.1)
                try:
                    p, comm, err = sts.ReadPos(mid)
                    if comm == 0 and err == 0:
                        pos = p
                        break
                except IndexError:
                    pass
            retry_note = f" (after {attempt + 1} attempts)" if attempt > 0 else ""
            if pos is None:
                print(f"Motor {mid}: calibrated, but readback failed{retry_note}")
            else:
                print(f"Motor {mid}: calibrated, position now reads {pos}{retry_note}")

        print("\nDone. Middle position set for all motors.")
    finally:
        port.closePort()


if __name__ == "__main__":
    main()