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 = "stepper-<class>-<number in class>-<name>"
MQTT_TOPIC = "elsys/<class>/<number in class>/<name>"

# Motor constants
STEP_PIN = 12
DIR_PIN = 14

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

# Handle MQTT messages
def on_message(topic, message):
    topic = topic.decode("utf-8")
    temperature = float(message.decode("utf-8"))
    clear_display()
    if temperature <= 21:
        stepper.duty(0)
        fill_display(0, temperature)
    elif temperature <= 24:
        stepper.duty(512)
        stepper.freq(200)
        fill_display(60, temperature)
    else:
        stepper.duty(512)
        stepper.freq(400)
        fill_display(120, temperature)

# Fill display
def fill_display(rpm, temp):
    display.fill(0)
    display.text("Current RPM:", 0, 0)
    display.text("{}".format(rpm), 0, 15)
    display.text("Current Temp:", 0, 30)
    display.text("{:.2f}".format(temp), 0, 45)
    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 stepper motor
dir_pin = Pin(DIR_PIN, Pin.OUT)
dir_pin.value(1)
step_pin = Pin(STEP_PIN, Pin.OUT)
stepper = PWM(step_pin)
stepper.duty(0)

# 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:
        client.wait_msg()
finally:
    client.disconnect()