import time
import network
import ssd1306

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

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

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

# Relay Constants
RELAY_IN_PIN = 23

# Roles
ROLES = {
    "123#": "Janitor",
    "456#": "Security Guard"
}

# 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!")
    client.subscribe(MQTT_TOPIC)
    print("Subscribed to {}".format(MQTT_TOPIC))

# Handle MQTT messages
def on_message(topic, message):
    topic = topic.decode("utf-8")
    code = message.decode("utf-8")
    fill_display(code, ROLES[code] if code in ROLES.keys() else "Unknown")
    if code in ROLES.keys():
        relay.value(1) 
    time.sleep(5)
    relay.value(0)
    clear_display()

# Fill display
def fill_display(code, role):
    display.text(code, 0, 5)
    display.text(role, 0, 30)
    display.show()

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

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

# 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)

# Init relay
relay = Pin(RELAY_IN_PIN, Pin.OUT)
relay.value(0)

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