from machine import Pin, I2C
import ssd1306
import time

# OLED display settings (SSD1306 I2C)
i2c = I2C(-1, scl=Pin(22), sda=Pin(21))
oled = ssd1306.SSD1306_I2C(128, 64, i2c)

# Button pins
plus_button = Pin(14, Pin.IN, Pin.PULL_UP)
minus_button = Pin(27, Pin.IN, Pin.PULL_UP)
multiply_button = Pin(26, Pin.IN, Pin.PULL_UP)
divide_button = Pin(25, Pin.IN, Pin.PULL_UP)
equal_button = Pin(33, Pin.IN, Pin.PULL_UP)

# Variables to store the calculation state
first_number = ''
second_number = ''
operation = ''
result = ''

def update_display():
    oled.fill(0)  # Clear display
    if result == '':
        # Display the operation as user inputs it
        oled.text(f'{first_number} {operation} {second_number}', 0, 0)
    else:
        # Display the final result after "=" is pressed
        oled.text(f'{first_number} {operation} {second_number} =', 0, 0)
        oled.text(str(result), 0, 10)
    oled.show()

def calculate():
    global first_number, second_number, operation, result
    try:
        num1 = float(first_number)
        num2 = float(second_number)
        if operation == '+':
            result = num1 + num2
        elif operation == '-':
            result = num1 - num2
        elif operation == '*':
            result = num1 * num2
        elif operation == '/':
            result = num1 / num2 if num2 != 0 else 'Error'
    except:
        result = 'Error'
    update_display()

# Interrupt service routines for buttons
def handle_plus(pin):
    global operation
    if first_number != '' and result == '':
        operation = '+'
    update_display()

def handle_minus(pin):
    global operation
    if first_number != '' and result == '':
        operation = '-'
    update_display()

def handle_multiply(pin):
    global operation
    if first_number != '' and result == '':
        operation = '*'
    update_display()

def handle_divide(pin):
    global operation
    if first_number != '' and result == '':
        operation = '/'
    update_display()

def handle_equal(pin):
    global result
    if first_number != '' and second_number != '' and operation != '':
        calculate()

# Attach button interrupts
plus_button.irq(trigger=Pin.IRQ_FALLING, handler=handle_plus)
minus_button.irq(trigger=Pin.IRQ_FALLING, handler=handle_minus)
multiply_button.irq(trigger=Pin.IRQ_FALLING, handler=handle_multiply)
divide_button.irq(trigger=Pin.IRQ_FALLING, handler=handle_divide)
equal_button.irq(trigger=Pin.IRQ_FALLING, handler=handle_equal)

# Main loop to input numbers and update the display
while True:
    if operation == '':
        # Read first number
        if plus_button.value() == 0 or minus_button.value() == 0:
            pass  # Skip if operation button is pressed
        else:
            first_number += input("Enter first number: ")
    elif result == '':
        # Read second number
        if equal_button.value() == 0:
            pass  # Skip if equal button is pressed
        else:
            second_number += input("Enter second number: ")
    update_display()
    time.sleep(0.1)