import time
from machine import Pin, I2C
from ssd1306 import SSD1306_I2C

# Initialize pins for buttons, LED, motors
btn_red = Pin(18, Pin.IN, Pin.PULL_UP)
btn_blue = Pin(19, Pin.IN, Pin.PULL_UP)
led_green = Pin(20, Pin.OUT)

# Initialize pins for stepper motor 1
Ap1 = Pin(10, Pin.OUT)
Am1 = Pin(11, Pin.OUT)
Bp1 = Pin(12, Pin.OUT)
Bm1 = Pin(13, Pin.OUT)

# Initialize pins for stepper motor 2
Ap2 = Pin(21, Pin.OUT)
Am2 = Pin(22, Pin.OUT)
Bp2 = Pin(23, Pin.OUT)
Bm2 = Pin(24, Pin.OUT)

# Initialize I2C and OLED display
i2c = I2C(0, scl=Pin(5), sda=Pin(4))
oled = SSD1306_I2C(128, 64, i2c)

state = False

def updateOLED(state):
    oled.fill(0)
    if state:
        oled.text("State: ON", 0, 0)
    else:
        oled.text("State: OFF", 0, 0)
    oled.show()

def engine(Ap, Am, Bm, Bp):
    # Simple stepper motor control logic
    Ap.value(1)
    Am.value(0)
    Bm.value(0)
    Bp.value(1)
    time.sleep(0.01)
    Ap.value(0)
    Am.value(1)
    Bm.value(1)
    Bp.value(0)
    time.sleep(0.01)

def loop():
    global state
    while True:
        if not btn_red.value():  # If red button is pressed (active low)
            print("Red button pressed")
            state = True
            led_green.value(1)
            updateOLED(state)
            time.sleep(0.1)
            while state:
                engine(Ap1, Am1, Bm1, Bp1)  # Run motor 1
                engine(Ap2, Am2, Bm2, Bp2)  # Run motor 2
                if not btn_blue.value():  # If blue button is pressed (active low)
                    print("Blue button pressed")
                    state = False
                    led_green.value(0)
                    updateOLED(state)
                    time.sleep(0.1)
                    break  # Exit the while loop
        time.sleep(0.1)  # Add a small delay to avoid button bounce

# Start the main loop
loop()