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

# === CONFIG ===
BROKER = "broker.emqx.io"  # Public test broker
TOPIC_TEMP = "iot-home/temperature"
TOPIC_CMD = "iot-home/command"

# === SETUP ===
led = machine.Pin("LED", machine.Pin.OUT)

def connect_wifi():
    wlan = network.WLAN(network.STA_IF)
    wlan.active(True)
    print("Connecting to WiFi...", end="")
    wlan.connect("Wokwi-GUEST", "")
    timeout = 30
    while not wlan.isconnected() and timeout > 0:
        print(".", end="")
        time.sleep(0.5)
        timeout -= 1
    print()
    if wlan.isconnected():
        print("WiFi OK | IP:", wlan.ifconfig()[0])
        return True
    else:
        print("WiFi FAILED")
        return False

# MQTT callback when a message arrives
def mqtt_callback(topic, msg):
    print("Received:", topic.decode(), "->", msg.decode())
    if msg.decode() == "ON":
        led.on()
    elif msg.decode() == "OFF":
        led.off()

print("\n=== MQTT Test ===")

if not connect_wifi():
    raise RuntimeError("No WiFi")

# Connect to MQTT broker
client = MQTTClient("pico_w_" + str(time.ticks_ms()), BROKER)
client.set_callback(mqtt_callback)
client.connect()
print("MQTT connected to", BROKER)

# Subscribe to commands
client.subscribe(TOPIC_CMD)
print("Subscribed to", TOPIC_CMD)

# Publish test data & check for commands
count = 0
while count < 10:
    count += 1
    msg = '{"temp": 25.5, "humidity": 60, "count": ' + str(count) + '}'
    client.publish(TOPIC_TEMP, msg)
    print("Published:", msg)

    # Check for incoming commands (non-blocking)
    client.check_msg()

    time.sleep(2)

client.disconnect()
print("\n=== Test Complete ===")