from machine import Pin
import time

# Define GPIO pins
green_led_pin = 21
blue_led_pin = 20
red_led_pin = 19
button_pin = 22

# Initialize LED and Button
green_led = Pin(green_led_pin, Pin.OUT)
blue_led = Pin(blue_led_pin, Pin.OUT)
red_led = Pin(red_led_pin, Pin.OUT)
button = Pin(button_pin, Pin.IN, Pin.PULL_UP)

def blink_led(led, nb_time:int=None, speed:float=0.5):
    i = 0
    if nb_time is None:
        while True:
            if button.value():
                led.off()
                return False
            led.on()
            time.sleep(speed)
            led.off()
            time.sleep(speed)
            print("blue led blinking")
    else:
        while i < nb_time:
            if button.value():
                led.off()
                return False
            green_led.on()
            time.sleep(speed)
            green_led.off()
            time.sleep(speed)
            i += 1
            print(f"green led blinked {i} times")
        return True
    

def turn_off_all():
    green_led.off()
    blue_led.off()
    red_led.off()

while True:
  # If Button pressed
  if not button.value():
    red_led.on()
    if blink_led(green_led, nb_time=20):
        blink_led(blue_led)
    # Else
  else:
    turn_off_all()