import time
import network
import ssd1306

from machine import Pin, PWM, I2C
from umqtt.simple import MQTTClient

# Modify wherever the values contain <>
MQTT_BROKER = "broker.hivemq.com"
MQTT_CLIENT_ID = "button-<class>-<number in class>-<name>"
MQTT_TOPIC = "elsys/<class>/<number in class>/<name>"

# Button constants
BUTTON_PINS = [26, 25, 33, 32] 
BUTTONS = []

# Display constants
DISPLAY_SCL_PIN = 21
DISPLAY_SDA_PIN = 22
DISPLAY_WIDTH = 128
DISPLAY_HEIGHT = 64

# Connect to Wi-Fi
def connect_to_wifi():
    print("Connecting to Wi-Fi", end="")
    sta_if = network.WLAN(network.STA_IF)
    sta_if.active(True)
    sta_if.connect('Wokwi-GUEST', '')
    while not sta_if.isconnected():
        print(".", end="")
        time.sleep(0.1)
    print(" Connected!")

# Connect to MQTT Broker
def connect_to_mqtt_broker():
    connect_to_wifi()
    print("Connecting to MQTT server... ", end="")
    client.connect()
    print("Connected!")

# Publish button when pressed
def publish_button(pin):
    clear_display()
    button = str(BUTTONS.index(pin))
    client.publish(MQTT_TOPIC, button)
    fill_display(button)

# Fill display
def fill_display(button):
    display.fill(0)
    display.text("Button: " + button, 0, 0)
    display.show()

# Clear display
def clear_display():
    display.fill(0)
    display.show()

# Init MQTT client
client = MQTTClient(MQTT_CLIENT_ID, MQTT_BROKER)
connect_to_mqtt_broker()

# Init buttons
BUTTONS = [Pin(pin, Pin.IN, Pin.PULL_DOWN) for pin in BUTTON_PINS]
for button_pin in BUTTONS:
    button_pin.irq(trigger=Pin.IRQ_RISING, handler=publish_button)

# Init OLED display
i2c = I2C(0, scl=Pin(DISPLAY_SCL_PIN), sda=Pin(DISPLAY_SDA_PIN))
display = ssd1306.SSD1306_I2C(DISPLAY_WIDTH, DISPLAY_HEIGHT, i2c)

# Loop
try:
    while 1:
        pass
finally:
    client.disconnect()