from machine import Pin, I2C, time_pulse_us
import time, network, machine
from simple import MQTTClient
import ssd1306

# === Pin Definitions ===
TRIG_PIN    = Pin(5, Pin.OUT)
ECHO_PIN    = Pin(18, Pin.IN)
BUZZER_PIN  = Pin(14, Pin.OUT)
GREEN_LED   = Pin(2, Pin.OUT)   # onboard LED
RED_LED     = Pin(4, Pin.OUT)   # external LED

# === OLED Setup ===
i2c  = I2C(scl=Pin(22), sda=Pin(21))
oled = ssd1306.SSD1306_I2C(128, 64, i2c)

# === WiFi/MQTT ===
SSID               = "Wokwi-GUEST"
PASSWORD           = ""
MQTT_SERVER        = "broker.hivemq.com"
MQTT_CLIENT_ID     = "esp32-bus-alert"
MQTT_TOPIC_DISTANCE= b"teddy/distance"
MQTT_TOPIC_STATUS  = b"teddy/status"

ARRIVAL_THRESHOLD  = 100.0
bus_previously_arrived = False

def connect_wifi():
    wlan = network.WLAN(network.STA_IF)
    wlan.active(True)
    wlan.connect(SSID, PASSWORD)
    print("Connecting to WiFi", end="")
    while not wlan.isconnected():
        print(".", end="")
        time.sleep(0.3)
    print(" connected!")

client = MQTTClient(MQTT_CLIENT_ID, MQTT_SERVER)

def connect_mqtt():
    try:
        client.connect()
        print("Connected to MQTT!")
    except Exception as e:
        print("Failed to connect to MQTT:", e)
        time.sleep(2)
        machine.reset()

def read_distance_cm():
    TRIG_PIN.off()
    time.sleep_us(2)
    TRIG_PIN.on()
    time.sleep_us(10)
    TRIG_PIN.off()
    try:
        duration = time_pulse_us(ECHO_PIN, 1, 30000)
        return duration * 0.034 / 2
    except OSError:
        return 999

def main():
    global bus_previously_arrived
    # Default states:
    BUZZER_PIN.off()
    GREEN_LED.off()
    RED_LED.off()

    connect_wifi()
    connect_mqtt()

    while True:
        distance = read_distance_cm()
        print("Distance:", distance, "cm")

        # Publish distance (optional: retain too)
        client.publish(MQTT_TOPIC_DISTANCE, str(round(distance,2)), retain=True)

        oled.fill(0)
        oled.text("Smart Bus Alert", 0, 0)
        oled.text("Distance: {:.1f}cm".format(distance), 0, 16)

        if distance <= ARRIVAL_THRESHOLD:
            # Arriving
            status = "Arriving"
            print(">> Sending status:", status)
            client.publish(MQTT_TOPIC_STATUS, status, retain=True)
            GREEN_LED.on()
            RED_LED.off()
            # Beep
            BUZZER_PIN.on()
            time.sleep_ms(500)
            BUZZER_PIN.off()
            bus_previously_arrived = True
            oled.text("Status: " + status, 0, 40)

        else:
            # Not here yet
            status = "Not Yet"
            print(">> Sending status:", status)
            client.publish(MQTT_TOPIC_STATUS, status, retain=True)
            GREEN_LED.off()
            RED_LED.on()
            BUZZER_PIN.off()
            bus_previously_arrived = False
            oled.text("Status: " + status, 0, 40)

        oled.show()
        time.sleep(1)

main()