import machine
import time

BUTTON_COUNT = 3
LED_COUNT = 9
BUTTON_START_ID = 16
LED_START_ID = 7

# global variables
last_button_time_stamp = 0
key_presses = []
passcode_entered = False

pin_to_index = {}

# Button presses
def button_handler(pin):
    global last_button_time_stamp, passcode_entered
    current_time = time.ticks_ms()
    if time.ticks_diff(current_time, last_button_time_stamp) > 200:
        last_button_time_stamp = current_time
        key_presses.append(pin)
        button_index = pin_to_index[pin]
        print(f'Button {button_index} pressed')
        if not passcode_entered:
            passcode_entered = True

# Set up buttons 
def setup_buttons():
    buttons = []
    for i in range(BUTTON_COUNT):
        pin = machine.Pin(BUTTON_START_ID + i, machine.Pin.IN, machine.Pin.PULL_DOWN)
        pin.irq(trigger=machine.Pin.IRQ_FALLING, handler=button_handler)
        buttons.append(pin)
        pin_to_index[pin] = i
    return buttons

# Set up LEDs
def setup_leds():
    return [machine.Pin(LED_START_ID + i, machine.Pin.OUT) for i in range(LED_COUNT)]

# Main function
def main():
    global key_presses, last_button_time_stamp, passcode_entered

    # Initialize the multiplexer
    s0 = machine.Pin(27, machine.Pin.OUT)
    s1 = machine.Pin(28, machine.Pin.OUT)
    mux_in = machine.Pin(26, machine.Pin.IN, machine.Pin.PULL_DOWN)

    # Set up buttons and LEDs
    buttons = setup_buttons()
    leds = setup_leds()

    # Correct passcode 
    PASS_CODE = [buttons[0], buttons[2], buttons[1]]

    last_mux_output = -1

    while True:
        # Read the multiplexer
        mux_output = 0
        for i in range(4):
            s0.value(i % 2)
            s1.value(i // 2)
            time.sleep(0.02)
            mux_output += mux_in.value() << i

        # Print new multiplexer output
        if mux_output != last_mux_output:
            last_mux_output = mux_output
            print(f'Multiplexer output: {mux_output}')
        
        time.sleep(0.1)

        # Check if 3 seconds have passed since the last button press
        if passcode_entered and time.ticks_diff(time.ticks_ms(), last_button_time_stamp) >= 3000:
            print("Time expired. Enter the passcode again.")
            key_presses.clear()
            passcode_entered = False

        # Handle passcode entry and LED toggling
        elif passcode_entered and len(key_presses) == len(PASS_CODE):
            if key_presses == PASS_CODE:
                if mux_output < LED_COUNT:
                    leds[mux_output].toggle()
                    print(f'Passcode Correct, Toggling LED {mux_output}')
                else:
                    print(f'Invalid mux output: {mux_output}')
            else:
                print('Incorrect passcode')
            key_presses.clear()
            passcode_entered = False

if __name__ == "__main__":
    main()