from machine import Pin, I2C
import ssd1306
import time

# LED pins
red = Pin(10, Pin.OUT)
yellow = Pin(9, Pin.OUT)
green = Pin(8, Pin.OUT)

# I2C for OLED
i2c = I2C(1, scl=Pin(7), sda=Pin(6), freq=400000)
oled = ssd1306.SSD1306_I2C(128, 64, i2c)

def display_message(state, seconds):
    oled.fill(0)
    oled.text("Traffic Light", 10, 5)
    oled.text(state, 45, 25)
    oled.text("Time:", 35, 45)
    oled.text(str(seconds), 75, 45)
    oled.show()

def countdown(state, led_on, duration):
    # Turn all OFF first
    red.off()
    yellow.off()
    green.off()

    # Turn ON required LED
    led_on.on()

    for t in range(duration, 0, -1):
        display_message(state, t)
        time.sleep(1)

while True:
    # STOP - Red (5 sec)
    countdown("STOP", red, 5)

    # WAIT - Yellow (2 sec)
    countdown("WAIT", yellow, 2)

    # GO - Green (5 sec)
    countdown("GO", green, 5)