import dht
import network
import ujson
from time import sleep
from machine import Pin, I2C
from ssd1306 import SSD1306_I2C
from hcsr04 import HCSR04
from simple import MQTTClient
from servo import Servo

# MQTT settings
WIFI_SSID = "Wokwi-GUEST"
WIFI_PASSWORD = ""
MQTT_CLIENT_ID = ""
MQTT_BROKER = 'test.mosquitto.org'
MQTT_PORT = 1883
MQTT_TOPIC_DIST = b'poultry/distance'
MQTT_TOPIC_TEMP = b'poultry/temperature'
MQTT_TOPIC_HUM = b'poultry/humidity'
MQTT_TOPIC_ALERT = b'poultry/alert'
MQTT_TOPIC_PRESENCE = b'poultry/presence'
MQTT_TOPIC_DOOR = b'poultry/door'


# Initialize I2C and OLED display
i2c = I2C(1, scl=Pin(22), sda=Pin(21))
oled_width = 128
oled_height = 64
oled = SSD1306_I2C(oled_width, oled_height, i2c)

# Initialize the HC-SR04 ultrasonic sensor
ultrasonic = HCSR04(Pin(16), Pin(4))

# Initialize DHT22
sensor = dht.DHT22(Pin(5))

# Initialize servo motor
servo_motor = Servo(pin=17)

# Initialize LEDs
red = Pin(14, Pin.OUT)
white = Pin(27, Pin.OUT)

# Connect to Wi-Fi
def connect_wifi():
    sta_if = network.WLAN(network.STA_IF)
    if not sta_if.isconnected():
        print("Connecting to WiFi...")
        sta_if.active(True)
        sta_if.connect(WIFI_SSID, WIFI_PASSWORD)
        while not sta_if.isconnected():
            pass
        print("Connected to WiFi:", sta_if.ifconfig())

# Connect to MQTT broker
def connect_mqtt():
    client = MQTTClient('esp32_client', MQTT_BROKER, port=MQTT_PORT)
    client.set_callback(mqtt_callback)
    client.connect()
    print('Connected to MQTT broker')
    return client

# MQTT callback function
def mqtt_callback(topic, msg):
    print("Received message on topic {}: {}".format(topic, msg))

# Main program
def main():
    connect_wifi()
    client = connect_mqtt()
    client.set_callback(mqtt_callback)

    try:
        while True:
            # Measure distance
            distance_cm = ultrasonic.distance_cm()
            print("Distance:", distance_cm, "cm")
            
            # Measure temperature and humidity
            sensor.measure()
            temp = sensor.temperature()
            hum = sensor.humidity()
            print("Temperature:", temp, "C")
            print("Humidity:", hum, "%")

            # Display on OLED
            print("Updating OLED display...")
            oled.fill(0)
            oled.text(f"Temp: {temp} C", 0, 0)
            oled.text(f"Hum: {hum} %", 0, 10)
            oled.text(f"Dist: {distance_cm} cm", 0, 20)
            oled.show()

            # Publish data to MQTT and Node RED Dashboard
            client.publish(MQTT_TOPIC_DIST, str(distance_cm))
            client.publish(MQTT_TOPIC_TEMP, str(temp))
            client.publish(MQTT_TOPIC_HUM, str(hum))

            # Check temperature and humidity thresholds
            if temp > 30:
                client.publish(MQTT_TOPIC_ALERT, "High temperature alert: {} C Please observe the hens!".format(temp))
            elif temp < 18:
                client.publish(MQTT_TOPIC_ALERT, "Low temperature alert: {} C Please observe the hens!".format(temp))

            if hum > 70:
                client.publish(MQTT_TOPIC_ALERT, "High humidity alert: {} % Please observe the hens!".format(hum))
            elif hum < 25:
                client.publish(MQTT_TOPIC_ALERT, "Low humidity alert: {} % Please observe the hens!".format(hum))

            # Check for presence and update servo and LEDs
            if distance_cm < 150:
                client.publish(MQTT_TOPIC_DOOR, "Door open")
                servo_motor.move(75)
                white.value(1)
                red.value(0)
                print("White LED on")
                oled.text("Person detected", 0, 30)     
                oled.text("White light ON", 0, 40)
                oled.show()
                client.publish(MQTT_TOPIC_PRESENCE, "Workers are checking the hens, white light mode")
            else:
                client.publish(MQTT_TOPIC_DOOR, "Door closed")
                servo_motor.move(0)
                red.value(1)
                white.value(0)
                print("Red LED on")
                oled.text("No person detected", 0, 30)
                oled.text("Red light ON", 0, 40)
                oled.show()
                client.publish(MQTT_TOPIC_PRESENCE, "No worker detected inside the poultry, red light mode")
            
            sleep(3)

        client.check_msg()
        sleep(3)

    except KeyboardInterrupt:
        print('Program stopped')
    finally:
        client.disconnect()
        print('Disconnected from MQTT broker')      

if __name__ == '__main__':
    main()