from micropython import const
import asyncio
import aioble
import bluetooth
from machine import Pin, PWM
from machine import Pin, I2C
import time
import ssd1306
from hcsr04 import HCSR04

# OLED (change pins if needed)
i2c = I2C(0, scl=Pin(17), sda=Pin(18))
oled = ssd1306.SSD1306_I2C(128, 64, i2c)

# Two ultrasonic sensors
head = HCSR04(trigger_pin=10, echo_pin=9)
body  = HCSR04(trigger_pin=8, echo_pin=7)

async def ultrasound_info(oled, front, back):
    while True:
        d1 = head.distance_cm()
        d2 = body.distance_cm()

        oled.fill(0)
        oled.text("Ultrasonic", 0, 0)
        oled.text("Head: {} cm".format(int(d1)), 0, 20)
        oled.text("Body:  {} cm".format(int(d2)), 0, 35)
        oled.show()

        await asyncio.sleep(0.1)

# ============================
# BLE UUIDs 
# ============================
_UART_SERVICE_UUID = bluetooth.UUID("12345678-1234-5678-1234-56789abcdef0")
_UART_RX_CHAR_UUID = bluetooth.UUID("12345678-1234-5678-1234-56789abcdef1")
_UART_TX_CHAR_UUID = bluetooth.UUID("6e400003-b5a3-f393-e0a9-e50e24dcca9e")

uart_service = aioble.Service(_UART_SERVICE_UUID)
rx_char = aioble.Characteristic(
    uart_service, _UART_RX_CHAR_UUID, write=True, capture=True
)
tx_char = aioble.Characteristic(
    uart_service, _UART_TX_CHAR_UUID, notify=True
)
aioble.register_services(uart_service)

# ============================
# SERVO CONSTANTS
# ============================
SERVO_MIN = 0
SERVO_MAX = 180
SERVO_CENTER = 90
SERVO_STEP = 1
SERVO_UPDATE_MS = 1

SERVO_IDLE = 0
SERVO_MOVE = 1

# ============================
# SERVO CLASS
# ============================
class Servo:
    def __init__(self, pin):
        self.pwm = PWM(Pin(pin), freq=50)
        self.angle = SERVO_CENTER
        self.write(self.angle)

    def write(self, angle):
        angle = max(SERVO_MIN, min(SERVO_MAX, angle))
        self.angle = angle
        duty = int((angle / 180) * 102 + 26)
        self.pwm.duty(duty)

# ============================
# SERVOS
# ============================
servo_right_arm = Servo(6)
servo_head      = Servo(5)
servo_left_arm  = Servo(4)

servos = {
    "left_arm": servo_left_arm,
    "right_arm": servo_right_arm,
    "head": servo_head,
}

servo_dir = {
    "left_arm": 0,
    "right_arm": 0,
    "head": 0,
}

servo_state = SERVO_IDLE

# ============================
# MOTOR SETUP (UNCHANGED)
# ============================
IN1 = Pin(11, Pin.OUT)
IN2 = Pin(12, Pin.OUT)
ENA = PWM(Pin(15), freq=1000)

IN3 = Pin(13, Pin.OUT)
IN4 = Pin(14, Pin.OUT)
ENB = PWM(Pin(16), freq=1000)

def motors_forward(speed=900):
    IN1.value(1); IN2.value(0)
    IN3.value(1); IN4.value(0)
    ENA.duty(speed); ENB.duty(speed)

def motors_backward(speed=900):
    IN1.value(0); IN2.value(1)
    IN3.value(0); IN4.value(1)
    ENA.duty(speed); ENB.duty(speed)

def motors_left(speed=900):
    IN1.value(0); IN2.value(1)
    IN3.value(1); IN4.value(0)
    ENA.duty(speed); ENB.duty(speed)

def motors_right(speed=900):
    IN1.value(1); IN2.value(0)
    IN3.value(0); IN4.value(1)
    ENA.duty(speed); ENB.duty(speed)

def motors_stop():
    global servo_state

    # Stop motors
    IN1.value(0); IN2.value(0)
    IN3.value(0); IN4.value(0)
    ENA.duty(0); ENB.duty(0)

    # Stop servo motion
    for k in servo_dir:
        servo_dir[k] = 0

    servo_state = SERVO_IDLE

    # IMMEDIATE snap to center
    #servo_left_arm.write(SERVO_CENTER)
    #servo_right_arm.write(SERVO_CENTER)
    #servo_head.write(SERVO_CENTER)

# ============================
# SERVO TASK (SINGLE WRITER)
# ============================
async def servo_task():
    while True:
        if servo_state == SERVO_MOVE:
            for name, servo in servos.items():
                d = servo_dir[name]
                if d:
                    servo.write(servo.angle + d * SERVO_STEP)

        await asyncio.sleep_ms(SERVO_UPDATE_MS)

# ============================
# BLE COMMAND HANDLER (UNCHANGED)
# ============================
async def uart_rx_handler():
    global servo_state

    while True:
        conn, data = await rx_char.written()
        try:
            msg = data.decode().strip()
        except:
            msg = str(data)

        print("Received:", msg)

        # MOVEMENT
        if msg == "f":
            motors_forward()

        elif msg == "b":
            motors_backward()

        elif msg == "r":
            motors_right()

        elif msg == "l":
            motors_left()

        elif msg == "s":
            motors_stop()

        # SERVOS (CONTINUOUS)
        elif msg == "1":
            servo_dir["left_arm"] = -1
            servo_state = SERVO_MOVE

        elif msg == "2":
            servo_dir["left_arm"] = 1
            servo_state = SERVO_MOVE

        elif msg == "3":
            servo_dir["right_arm"] = 1
            servo_state = SERVO_MOVE

        elif msg == "4":
            servo_dir["right_arm"] = -1
            servo_state = SERVO_MOVE

        elif msg == "5":
            servo_dir["head"] = 1
            servo_state = SERVO_MOVE

        elif msg == "6":
            servo_dir["head"] = -1
            servo_state = SERVO_MOVE

        else:
            print("Wut? Plz map command")

# ============================
# BLE ADVERTISING (UNCHANGED)
# ============================
async def uart_advertise():
    while True:
        print("Advertising BLE UART...")
        async with await aioble.advertise(
            250_000,
            name="ESP32-BLE-UART",
            services=[_UART_SERVICE_UUID]
        ) as conn:
            print("Connected:", conn.device)
            await conn.disconnected()
            print("Disconnected")

# ============================
# MAIN
# ============================
async def main():
    asyncio.create_task(servo_task())
    asyncio.create_task(ultrasound_info(oled, head, body))
    asyncio.create_task(uart_rx_handler())
    asyncio.create_task(uart_advertise())
    
    while True:
        await asyncio.sleep(1)

asyncio.run(main())