from machine import Pin, I2C
import time
import ssd1306

# ---------- OLED SETUP ----------
i2c = I2C(0, scl=Pin(1), sda=Pin(0))
oled = ssd1306.SSD1306_I2C(128, 64, i2c)

# ---------- RGB LED SETUP ----------
red = Pin(15, Pin.OUT)
green = Pin(14, Pin.OUT)

def leds_off():
    red.value(0)
    green.value(0)

def display(text, count):
    oled.fill(0)
    oled.text(text, 35, 20)
    oled.text("Time: {}".format(count), 25, 40)
    oled.show()

# ---------- MAIN LOOP ----------
while True:

    # 🔴 RED - STOP
    leds_off()
    red.value(1)
    for i in range(20, 0, -1):
        display("STOP", i)
        time.sleep(1)

    # 🟡 YELLOW - READY
    leds_off()
    red.value(1)
    green.value(1)
    for i in range(5, 0, -1):
        display("READY", i)
        time.sleep(1)

    # 🟢 GREEN - GO
    leds_off()
    green.value(1)
    for i in range(20, 0, -1):
        display("GO", i)
        time.sleep(1)

    # 🟡 YELLOW - READY
    leds_off()
    red.value(1)
    green.value(1)
    for i in range(5, 0, -1):
        display("READY", i)
        time.sleep(1)