# main.py
# Pico W + DHT22 + SSD1306 + MQTT Example
# Features:
# 1. Reads temperature and humidity from DHT22 sensor.
# 2. Displays readings and system status on OLED.
# 3. Publishes readings to MQTT broker and handles commands.
# 4. Robust reconnect for Wi-Fi and MQTT disconnections.

import time
import json
import ubinascii
import machine
from machine import Pin, I2C
import network
import dht
import ssd1306_pico as ssd  # SSD1306 OLED driver for Pico W
from simple import MQTTClient

# CONFIG 
SSID = "saranya"               # Wi-Fi SSID
PASSWORD = "123456789"       # Wi-Fi password

# MQTT broker settings
MQTT_BROKER = "test.mosquitto.org"
MQTT_PORT = 1883
MQTT_PUB_TOPIC = b"picoW/dht/telemetry"  # Topic to publish sensor data
MQTT_SUB_TOPIC = b"picoW/dht/cmd"        # Topic to receive commands

PUBLISH_INTERVAL = 5  # Seconds between MQTT publishes

# PIN SETTINGS 
DHT_PIN = 22   # GPIO connected to DHT22 data pin
LED_PIN = 25   # Onboard LED to indicate commands
I2C_SDA = 14   # SDA pin for OLED
I2C_SCL = 15   # SCL pin for OLED

# OLED size
OLED_WIDTH = 128
OLED_HEIGHT = 64

#  PERIPHERALS 
# Initialize onboard LED
led = Pin(LED_PIN, Pin.OUT)
led.value(0)  # Turn off LED initially

# Initialize DHT22 sensor
d = dht.DHT22(Pin(DHT_PIN))

# Initialize I2C for OLED
i2c = I2C(1, scl=Pin(I2C_SCL), sda=Pin(I2C_SDA))
print("I2C devices found:", i2c.scan())  # Should print [60] for SSD1306

# Initialize OLED display
oled = ssd.SSD1306_I2C(OLED_WIDTH, OLED_HEIGHT, i2c)

def draw_display(line1="", line2="", line3="", line4=""):
    """
    Draw four lines of text on the OLED.
    Arguments:
        line1-4: Strings to display on OLED lines 0-3.
    """
    try:
        oled.fill(0)  # Clear display
        oled.text(str(line1), 0, 0)
        oled.text(str(line2), 0, 16)
        oled.text(str(line3), 0, 32)
        oled.text(str(line4), 0, 48)
        oled.show()
    except OSError as e:
        print("OLED I2C error:", e)

#  WIFI CONNECTION 
wlan = network.WLAN(network.STA_IF)
wlan.active(True)

def wifi_connect(ssid, password, timeout=15):
    """
    Connect to Wi-Fi network.
    Returns True if connected successfully, False on timeout.
    """
    if wlan.isconnected():
        return True
    wlan.connect(ssid, password)
    start = time.time()
    while not wlan.isconnected():
        if time.time() - start > timeout:
            return False
        time.sleep(0.5)
    return True

#  MQTT SETUP 
client = None
# Unique client ID
client_id = b"picoW-" + ubinascii.hexlify(machine.unique_id())  
def mqtt_cb(topic, msg):
    """
    MQTT callback: handle incoming messages from subscribed topic.
    Example commands sent via MQTT:
        - "LED ON"   -> turns onboard LED on
        - "LED OFF"  -> turns onboard LED off
        - "RESET"    -> clears OLED display
    """
    try:
        t = topic.decode()
        m = msg.decode().strip()
    except:
        return
    print("MQTT IN >", t, m)
    if m.upper() == "LED ON":
        led.value(1)
    elif m.upper() == "LED OFF":
        led.value(0)
    elif m.upper() == "RESET":
        oled.fill(0)
        oled.show()

def mqtt_connect_and_subscribe():
    """
    Connect to MQTT broker and subscribe to command topic.
    Returns MQTT client object.
    """
    global client
    client = MQTTClient(client_id, MQTT_BROKER, port=MQTT_PORT, keepalive=60)
    client.set_callback(mqtt_cb)
    client.connect()
    client.subscribe(MQTT_SUB_TOPIC)
    print("Connected to %s, subscribed to %s" % (MQTT_BROKER, MQTT_SUB_TOPIC.decode()))
    return client

#  DHT SENSOR READ 
def read_dht():
    """
    Read temperature (°C) and humidity (%) from DHT22.
    Returns (temperature, humidity) or (None, None) on error.
    """
    try:
        d.measure()
        return d.temperature(), d.humidity()
    except Exception as e:
        print("DHT read error:", e)
        return None, None

#  MAIN LOOP 
def main():
    print("Connecting to Wi-Fi...")
    wifi_ok = wifi_connect(SSID, PASSWORD)
    print("Wi-Fi OK:", wifi_ok, wlan.ifconfig() if wifi_ok else None)

    mqtt_ok = False
    if wifi_ok:
        try:
            mqtt_connect_and_subscribe()
            mqtt_ok = True
        except Exception as e:
            print("Initial MQTT connect failed:", e)
            mqtt_ok = False

    last_pub = 0

    while True:
        #  Handle incoming MQTT messages 
        if mqtt_ok:
            try:
                client.check_msg()  # Non-blocking check for new messages
            except Exception as e:
                print("MQTT check_msg error:", e)
                mqtt_ok = False

        #  Read DHT sensor 
        temp_c, hum = read_dht()

        #  Update OLED display 
        draw_display(
            "Temp: {}C".format(temp_c if temp_c is not None else "N/A"),
            "Hum: {}%".format(hum if hum is not None else "N/A"),
            "Wi-Fi: {}".format("OK" if wifi_ok else "ERR"),
            "MQTT: {}".format("OK" if mqtt_ok else "ERR")
        )

        #  Publish data via MQTT 
        now = time.time()
        if now - last_pub >= PUBLISH_INTERVAL:
            # Payload format: JSON with temperature, humidity, and timestamp
            payload = {
                "temp_c": None if temp_c is None else round(temp_c, 2),
                "humidity": None if hum is None else round(hum, 2),
                "ts": int(now)
            }

            if mqtt_ok:
                try:
                    client.publish(MQTT_PUB_TOPIC, json.dumps(payload))
                except Exception as e:
                    print("Publish failed:", e)
                    # Try reconnect immediately
                    try:
                        mqtt_connect_and_subscribe()
                        mqtt_ok = True
                        client.publish(MQTT_PUB_TOPIC, json.dumps(payload))
                    except Exception as e2:
                        print("Reconnect failed:", e2)
                        mqtt_ok = False
            else:
                # Attempt MQTT reconnect if Wi-Fi is still connected
                if wifi_ok:
                    try:
                        mqtt_connect_and_subscribe()
                        mqtt_ok = True
                        client.publish(MQTT_PUB_TOPIC, json.dumps(payload))
                    except Exception as e:
                        print("Reconnect failed:", e)
                        mqtt_ok = False

            last_pub = now

        # Small delay for loop stability
        time.sleep(1)

# RUN MAIN 
if __name__ == "__main__":
    main()