"""
COE 550 - HW2 : ESP32 IoT client - WOKWI SIMULATOR + MQTT VERSION
=================================================================
Paste this file into the main.py tab of a Wokwi "MicroPython on ESP32"
project (classic ESP32 board). It is a learning version; the hardware
version for the real ESP32-S3 is HW2_main_code.py.

Why the classic ESP32 here? MicroPython on Wokwi's simulated ESP32-S3 got
stuck in a watchdog-reset loop, while the classic ESP32 is Wokwi's most
mature MicroPython board. The code is the same for both chips; only the
pin numbers differ.

Why MQTT here?
  The Wokwi simulator runs in the cloud, so it cannot open a connection to
  your PC. Instead, the ESP32 and the PC server BOTH connect out to a free
  public MQTT broker, which passes messages between them:

      ESP32 (Wokwi) --publish DATA--> broker --> PC server (dashboard)
      ESP32 (Wokwi) <--CMD-- broker <--publish CMD-- PC server

  MQTT itself runs over a TCP socket, and the messages are the same as in
  the TCP version:
      topic <TOPIC_PREFIX>/data : DATA,<seq>,<temperature_C>,<humidity_pct>
      topic <TOPIC_PREFIX>/cmd  : CMD,<seq>,<RED|GREEN|BLUE>
  (No newline is needed: unlike a raw TCP stream, MQTT keeps each message separate.)

Other differences from the hardware version:
  * Wi-Fi is Wokwi's free virtual network "Wokwi-GUEST" (no password).
  * Wokwi has no DHT11 part, so a DHT22 is used (same library, same wiring).
  * The classic ESP32 board has no built-in RGB LED, so the same kind of LED
    (a WS2812 "NeoPixel") is wired to GPIO18. The real ESP32-S3 has it on board.

Local LED status color: YELLOW = connecting, or no reply from the server.
"""

import time
import network
import machine
import dht
import neopixel
import ubinascii
from umqtt.simple import MQTTClient, MQTTException

# ============================ USER CONFIGURATION ============================
WIFI_SSID = "Wokwi-GUEST"              # Wokwi's simulated Wi-Fi network
WIFI_PASSWORD = ""                     # it has no password

MQTT_BROKER = "broker.hivemq.com"      # free public broker (must match the server)
MQTT_PORT = 1883
TOPIC_PREFIX = "coe550/hw2/202611380"  # something unique, e.g. your student ID;
                                       # use the SAME value in the server
DATA_TOPIC = TOPIC_PREFIX + "/data"    # ESP32 -> server
CMD_TOPIC = TOPIC_PREFIX + "/cmd"      # server -> ESP32

DHT_PIN = 4             # DHT22 SDA (data) pin
RGB_LED_PINS = (18,)    # NeoPixel data pin (on the real ESP32-S3: built-in LED, GPIO48/38)
LED_BRIGHTNESS = 150    # 0-255
SWAP_RED_GREEN = False  # set True if GREEN shows when it should be RED

SEND_PERIOD_MS = 1000   # one reading per second
REPLY_TIMEOUT_MS = 3000 # longest wait for the server's CMD reply
WIFI_TIMEOUT_S = 15     # longest wait for the Wi-Fi connection
RETRY_DELAY_S = 2       # pause before reconnecting after an error
# ============================================================================

# ---------------------------- hardware objects ------------------------------
wlan = network.WLAN(network.STA_IF)
dht_sensor = dht.DHT22(machine.Pin(DHT_PIN, machine.Pin.IN, machine.Pin.PULL_UP))
# one single-pixel NeoPixel object per LED; set_led() writes the same color to all
rgb_leds = [neopixel.NeoPixel(machine.Pin(pin, machine.Pin.OUT), 1) for pin in RGB_LED_PINS]

LED_COLORS = {
    "RED": (LED_BRIGHTNESS, 0, 0),
    "GREEN": (0, LED_BRIGHTNESS, 0),
    "BLUE": (0, 0, LED_BRIGHTNESS),
    "YELLOW": (LED_BRIGHTNESS, LED_BRIGHTNESS // 2, 0),
    "OFF": (0, 0, 0),
}
SERVER_COLORS = ("RED", "GREEN", "BLUE")   # the only colors the server may command

latest_reply = None     # filled in by on_mqtt_message() when a CMD arrives


def set_led(color_name):
    """Show one of the named colors on every RGB LED."""
    red, green, blue = LED_COLORS[color_name]
    if SWAP_RED_GREEN:
        red, green = green, red
    for led in rgb_leds:
        led[0] = (red, green, blue)
        led.write()


# --------------------------------- Wi-Fi ------------------------------------
def connect_wifi():
    """Connect to the Wi-Fi network. Raises OSError if it takes too long."""
    wlan.active(True)
    if wlan.isconnected():
        return
    print("Connecting to Wi-Fi '{}' ...".format(WIFI_SSID))
    wlan.connect(WIFI_SSID, WIFI_PASSWORD)
    start_time = time.time()
    led_on = False
    while not wlan.isconnected():
        if time.time() - start_time > WIFI_TIMEOUT_S:
            try:
                wlan.disconnect()          # cancel the attempt before retrying
            except OSError:
                pass
            raise OSError("Wi-Fi connection timed out")
        led_on = not led_on
        set_led("YELLOW" if led_on else "OFF")   # blink while connecting
        time.sleep(0.5)
    print("Wi-Fi connected. ESP32 IP address:", wlan.ifconfig()[0])
    set_led("YELLOW")


# ---------------------------------- MQTT ------------------------------------
def on_mqtt_message(topic, message):
    """Called by the MQTT client whenever a message arrives on CMD_TOPIC."""
    global latest_reply
    latest_reply = message.decode().strip()


def connect_mqtt():
    """Connect to the broker, subscribe to the command topic, return the client."""
    client_id = b"esp32s3-" + ubinascii.hexlify(machine.unique_id())   # unique per board
    client = MQTTClient(client_id, MQTT_BROKER, port=MQTT_PORT, keepalive=60)
    client.set_callback(on_mqtt_message)
    client.connect()
    client.subscribe(CMD_TOPIC)
    print("Connected to MQTT broker {} - listening on '{}'".format(MQTT_BROKER, CMD_TOPIC))
    return client


def close_mqtt(client):
    """Disconnect from the broker, ignoring errors (the link may already be broken)."""
    if client is not None:
        try:
            client.disconnect()
        except OSError:
            pass


def wait_for_reply(client, timeout_ms):
    """Poll the broker until the server's reply arrives; None if it never does."""
    global latest_reply
    start = time.ticks_ms()
    while time.ticks_diff(time.ticks_ms(), start) < timeout_ms:
        client.check_msg()            # runs on_mqtt_message() if a message is waiting
        if latest_reply is not None:
            reply = latest_reply
            latest_reply = None
            return reply
        time.sleep_ms(20)
    return None


# --------------------------------- sensor -----------------------------------
def read_sensor():
    """Return (temperature_C, humidity_pct), or None if the sensor read failed."""
    try:
        dht_sensor.measure()
        return dht_sensor.temperature(), dht_sensor.humidity()
    except OSError as error:             # timeout or checksum error from the sensor
        print("DHT22 read failed:", error)
        return None


# -------------------------------- protocol ----------------------------------
def handle_server_reply(reply, sent_sequence):
    """Parse 'CMD,<seq>,<color>', update the LED and return the color (or None)."""
    fields = reply.split(",")
    if len(fields) == 3 and fields[0] == "CMD" and fields[2] in SERVER_COLORS:
        if fields[1] != str(sent_sequence):
            print("Warning: reply is for message", fields[1], "but we sent", sent_sequence)
        set_led(fields[2])
        return fields[2]
    print("Unexpected reply from server:", reply)
    return None


# ------------------------------- main loop ----------------------------------
def main():
    global latest_reply
    set_led("YELLOW")
    time.sleep(1)                # the DHT sensor needs about 1 s after power-up
    mqtt_client = None
    sequence_number = 0

    try:
        while True:
            loop_start = time.ticks_ms()
            try:
                # (re)connect Wi-Fi and the broker whenever needed
                if not wlan.isconnected():
                    close_mqtt(mqtt_client)
                    mqtt_client = None
                    connect_wifi()
                if mqtt_client is None:
                    mqtt_client = connect_mqtt()

                reading = read_sensor()
                if reading is not None:
                    temperature, humidity = reading
                    message = "DATA,{},{},{}".format(sequence_number, temperature, humidity)
                    print("Sending  :", message)
                    latest_reply = None                 # forget any late, old reply
                    mqtt_client.publish(DATA_TOPIC, message)

                    reply = wait_for_reply(mqtt_client, REPLY_TIMEOUT_MS)
                    if reply is None:
                        print("No reply - is the PC server running with the same TOPIC_PREFIX?")
                        set_led("YELLOW")
                    else:
                        print("Received :", reply)
                        led_color = handle_server_reply(reply, sequence_number)
                        if led_color is not None:
                            print("LED      :", led_color)
                    sequence_number += 1

            except (OSError, MQTTException) as error:   # Wi-Fi or broker failure
                print("Connection error:", error, "- retrying in", RETRY_DELAY_S, "s")
                close_mqtt(mqtt_client)
                mqtt_client = None
                set_led("YELLOW")
                time.sleep(RETRY_DELAY_S)

            # sleep for whatever is left of the 1-second period
            elapsed_ms = time.ticks_diff(time.ticks_ms(), loop_start)
            if elapsed_ms < SEND_PERIOD_MS:
                time.sleep_ms(SEND_PERIOD_MS - elapsed_ms)

    except KeyboardInterrupt:            # Stop button / Ctrl+C
        print("Stopped by user")
    finally:
        close_mqtt(mqtt_client)
        set_led("OFF")


main()