from machine import Pin, I2C
import time

# =============---==== PCA9685 =================
class PCA9685:
    def __init__(self, i2c, addr=0x40):
        self.i2c = i2c
        self.addr = addr
        self.i2c.writeto_mem(self.addr, 0x00, b'\x00')
        self.set_pwm_freq(50)

    def set_pwm_freq(self, freq):
        prescale = int(25000000.0 / (4096 * freq) - 1)
        old_mode = self.i2c.readfrom_mem(self.addr, 0x00, 1)[0]
        self.i2c.writeto_mem(self.addr, 0x00, bytes([old_mode | 0x10]))
        self.i2c.writeto_mem(self.addr, 0xFE, bytes([prescale]))
        self.i2c.writeto_mem(self.addr, 0x00, bytes([old_mode]))
        time.sleep(0.005)
        self.i2c.writeto_mem(self.addr, 0x00, bytes([old_mode | 0xa1]))

    def set_pwm(self, ch, on, off):
        data = bytearray([on & 0xFF, on >> 8, off & 0xFF, off >> 8])
        self.i2c.writeto_mem(self.addr, 0x06 + 4 * ch, data)

# ================= INIT =================
i2c = I2C(0, scl=Pin(1), sda=Pin(0))
pca = PCA9685(i2c)

# ================= SERVO =================
def set_servo_angle(ch, angle):
    angle = max(0, min(180, angle))
    pulse = int(120 + (angle / 180) * 400)
    pca.set_pwm(ch, 0, pulse)

def move_smooth(ch, current, target, step=1, delay=0.015):
    if current == target:
        set_servo_angle(ch, target)
        return target

    direction = 1 if target > current else -1

    for a in range(current, target + direction, direction * step):
        angle = max(0, min(180, a))
        set_servo_angle(ch, angle)
        time.sleep(delay)

    return target

# ================= CONFIG =================
servo_list = [
    ("BASE", 0),
    ("SHOULDER", 4),
    ("ELBOW", 8),
    ("CLAW", 12)
]

angles = [90, 90, 90, 0]
current_idx = 0

# ================= INIT POSITION =================
print("Initializing servos...")
for i in range(len(servo_list)):
    set_servo_angle(servo_list[i][1], angles[i])
    time.sleep(0.2)

print("\n=== STEP CONTROL MODE ===")
print("n: Next | +: +1° | -: -1° | s: Status | number: jump | reset\n")

# ================= LOOP =================
while True:
    name, ch = servo_list[current_idx]
    cmd = input(f"[{name} @ {angles[current_idx]}°] >> ").strip().lower()

    if not cmd:
        continue

    if cmd == "n":
        current_idx = (current_idx + 1) % len(servo_list)

    elif cmd == "+":
        if angles[current_idx] < 180:
            angles[current_idx] += 1
            set_servo_angle(ch, angles[current_idx])
            time.sleep(0.01)

    elif cmd == "-":
        if angles[current_idx] > 0:
            angles[current_idx] -= 1
            set_servo_angle(ch, angles[current_idx])
            time.sleep(0.01)

    elif cmd.isdigit():
        target = int(cmd)
        target = max(0, min(180, target))
        print(f"Moving {name} to {target}°...")
        angles[current_idx] = move_smooth(ch, angles[current_idx], target)

    elif cmd == "s":
        print("\n--- Current State ---")
        for i, s in enumerate(servo_list):
            print(f"{s[0]}: {angles[i]}°")

    elif cmd == "reset":
        print("Executing Safety Reset...")
        defaults = [90, 90, 90, 0]
        for i in range(len(servo_list)-1, -1, -1):
            if angles[i] != defaults[i]:
                print(f"Resetting {servo_list[i][0]}...")
                angles[i] = move_smooth(servo_list[i][1], angles[i], defaults[i])

    else:
        print(f"Unknown command: '{cmd}'")