import network
import time
from machine import Pin
from umqtt.simple import MQTTClient
import ujson

# MQTT Server Parameters
MQTT_CLIENT_ID = "mqttx_d97fe9eb"
MQTT_BROKER = "broker.emqx.io"
MQTT_USER = ""
MQTT_PASSWORD = ""
MQTT_TOPIC = "BrainWave/actuate_lamp12"

# LED Setup
led = Pin(15, Pin.OUT)

# WiFi Connection
print("Connecting to WiFi", 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!")

# MQTT Connection
def connect_mqtt():
    global client
    print("Connecting to MQTT server... ", end="")
    try:
        client = MQTTClient(MQTT_CLIENT_ID, MQTT_BROKER, user=MQTT_USER, password=MQTT_PASSWORD)
        client.set_callback(message_callback)
        client.connect()
        client.subscribe(MQTT_TOPIC)
        print("Connected!")
    except Exception as e:
        print(f"Failed to connect to MQTT server: {e}. Retrying in 5 seconds...")
        time.sleep(5)
        connect_mqtt()

# MQTT Callback Function
def message_callback(topic, msg):
    print(f"Message received from topic {topic.decode()}: {msg.decode()}")
    try:
        # Decode JSON message
        data = ujson.loads(msg.decode())
        command = data.get("msg", "").lower()

        # Actuate LED based on the message
        if command == "on":
            led.on()
            print("LED Turned On")
        elif command == "off":
            led.off()
            print("LED Turned Off")
    except ValueError:
        print("Invalid message format. Expected JSON.")

# Initial MQTT Connection
connect_mqtt()

# Main Loop
try:
    while True:
        # Toggle LED state
        led.value(not led.value())  # Toggle LED: If ON, turn OFF, and vice versa

        # Publish LED status
        led_status = "on" if led.value() else "off"
        message = ujson.dumps({"status": led_status})
        print(f"Publishing to {MQTT_TOPIC}: {message}")
        try:
            client.publish(MQTT_TOPIC, message)
        except OSError as e:
            print(f"Error while publishing: {e}. Reconnecting MQTT...")
            connect_mqtt()  # Reconnect if publishing fails

        # Check for incoming messages
        try:
            client.check_msg()
        except OSError as e:
            print(f"Error while checking messages: {e}. Reconnecting MQTT...")
            connect_mqtt()

        # Sleep for 5 seconds before toggling again
        time.sleep(5)

except KeyboardInterrupt:
    print("Disconnecting...")
    try:
        client.disconnect()
    except OSError:
        pass
    print("Disconnected!")