from machine import Pin, I2C
from time import *
from ssd1306 import SSD1306_I2C
from time import sleep, sleep_ms

# OLED Display setup
SCREEN_WIDTH = 128
SCREEN_HEIGHT = 64

i2c = I2C(0, scl=Pin(5), sda=Pin(4), freq=200000)  # I2C pins
oled = SSD1306_I2C(SCREEN_WIDTH, SCREEN_HEIGHT, i2c)

# Buttons
btn_red = Pin(9, Pin.IN, Pin.PULL_DOWN)
btn_blue = Pin(14, Pin.IN, Pin.PULL_DOWN)

# LEDs
led_red = Pin(7, Pin.OUT)

# Motors
Am1 = Pin(13, Pin.OUT)
Ap1 = Pin(12, Pin.OUT)
Bp1 = Pin(11, Pin.OUT)
Bm1 = Pin(10, Pin.OUT)

Am2 = Pin(21, Pin.OUT)
Ap2 = Pin(20, Pin.OUT)
Bp2 = Pin(18, Pin.OUT)
Bm2 = Pin(19, Pin.OUT)

# State variable
state = False

# Function to update OLED display
def updateOLED(state):
    oled.fill(0)
    if state:
        oled.text("Motors ON", 0, 10)
        oled.text("LED ON", 0, 20)
    else:
        oled.text("Motors OFF", 0, 10)
        oled.text("LED OFF", 0, 20)
    oled.show()

# Function to run motors
def engine(Ap, Am, Bm, Bp):
    sleep_ms(1)
    Bm.value(0)  # Bm off
    Ap.value(1)  # Ap on

    sleep_ms(1)
    Ap.value(0)
    Bp.value(1)

    sleep_ms(1)
    Bp.value(0)
    Am.value(1)

    sleep_ms(1)
    Am.value(0)
    Bm.value(1)

# Setup function
def setup():
    led_red.value(0)
    updateOLED(state)

# Main loop
def loop():
    global state
    if btn_red.value():  # If red button is pressed
        print("Red button pressed")
        state = True
        led_red.value(1)
        updateOLED(state)
        sleep(0.1)
        while state:
            engine(Ap1, Am1, Bm1, Bp1)  # Run motor 1
            engine(Ap2, Am2, Bm2, Bp2)  # Run motor 2
            if btn_blue.value():  # If blue button is pressed
                print("Blue button pressed")
                state = False
                led_red.value(0)
                updateOLED(state)
                sleep(0.1)

setup()

while True:
    loop()