from machine import Pin
import time

# Define GPIO pins for ESP32
redPin = Pin(16, Pin.OUT)
yellowPin = Pin(17, Pin.OUT)
greenPin = Pin(18, Pin.OUT)
buttonPin = Pin(19, Pin.IN, Pin.PULL_UP)

previousTime = time.ticks_ms()
state = 0  # 0: Green, 1: Yellow, 2: Red
buttonPressed = False

def update_lights():
    global previousTime, state, buttonPressed
    currentTime = time.ticks_ms()

    # Check button
    if buttonPin.value() == 0:  # Button pressed (LOW)
        if state != 2:  # If not in the red state
            buttonPressed = True

    # Traffic light logic
    if state == 0:  # Green light for 3 seconds
        greenPin.on()
        if time.ticks_diff(currentTime, previousTime) >= 3000:
            previousTime = currentTime
            greenPin.off()
            state = 1  # Switch to yellow

    elif state == 1:  # Yellow light for 1 second
        yellowPin.on()
        if time.ticks_diff(currentTime, previousTime) >= 1000:
            previousTime = currentTime
            yellowPin.off()
            state = 2  # Switch to red

    elif state == 2:  # Red light for 3 or 5 seconds
        redPin.on()
        delay_time = 5000 if buttonPressed else 3000
        if time.ticks_diff(currentTime, previousTime) >= delay_time:
            previousTime = currentTime
            redPin.off()
            state = 0  # Back to green
            buttonPressed = False  # Reset button

# Main loop
while True:
    update_lights()
    time.sleep(0.01)