from utime import sleep
from hx711 import HX711
import network
import machine
import ujson
from umqtt.simple import MQTTClient
import neopixel
from machine import Pin, I2C, PWM
from ssd1306 import SSD1306_I2C

# WiFi settings
WIFI_SSID = "Wokwi-GUEST"
WIFI_PASSWORD = ""

# MQTT settings
MQTT_BROKER = "broker.mqttdashboard.com"
CLIENT_ID = "esp32_client"
NEOPIXEL_TOPIC = b"neopixel/control/shaiful/ps4"
SERVO_TOPIC = b"servo/control/shaiful/ps4"
WEIGHT_TOPIC = b"coop/weight/shaiful/ps4"

# NeoPixel configuration
NEOPIXEL_PIN = 14
NUM_PIXELS = 10
np = neopixel.NeoPixel(machine.Pin(NEOPIXEL_PIN), NUM_PIXELS)

# Servo configuration
SERVO_PIN = 15
servo = PWM(machine.Pin(SERVO_PIN), freq=50)
SERVO_MIN_ANGLE = 0
SERVO_MAX_ANGLE = 180

# HX711 Load Cell configuration
HX711_DOUT = 4
HX711_SCK = 5
hx711 = HX711(d_out=HX711_DOUT, pd_sck=HX711_SCK)

# OLED display configuration
WIDTH = 128
HEIGHT = 64
i2c = I2C(0, scl=Pin(22), sda=Pin(21), freq=200000)
oled = SSD1306_I2C(WIDTH, HEIGHT, i2c)

# Function to connect to WiFi
def connect_wifi():
    wlan = network.WLAN(network.STA_IF)
    wlan.active(True)
    wlan.connect(WIFI_SSID, WIFI_PASSWORD)
    while not wlan.isconnected():
        pass
    print("Connected to WiFi:", wlan.ifconfig())

# Function to connect to MQTT broker
def connect_mqtt():
    client = MQTTClient(CLIENT_ID, MQTT_BROKER)
    client.set_callback(mqtt_callback)
    client.connect()
    client.subscribe(NEOPIXEL_TOPIC)
    client.subscribe(SERVO_TOPIC)
    print("Connected to MQTT Broker")
    return client

# MQTT callback function
def mqtt_callback(topic, msg):
    print("Received MQTT message on topic:", topic.decode())
    try:
        payload = ujson.loads(msg)
        if topic == NEOPIXEL_TOPIC:
            num_chickens = payload.get('chickens', 0)
            update_neopixel(num_chickens)
        elif topic == SERVO_TOPIC:
            angle = payload.get('angle', SERVO_MIN_ANGLE)
            set_servo_angle(angle)
    except Exception as e:
        print("Error processing MQTT message:", e)

# Function to set servo angle
def set_servo_angle(angle):
    if angle < SERVO_MIN_ANGLE:
        angle = SERVO_MIN_ANGLE
    elif angle > SERVO_MAX_ANGLE:
        angle = SERVO_MAX_ANGLE
    duty = int((angle / 180) * 1023)
    servo.duty(duty)
    print("Set servo angle:", angle)

# Function to update NeoPixel based on number of chickens
def update_neopixel(num_chickens):
    for i in range(NUM_PIXELS):
        if i < num_chickens:
            np[i] = (0, 255, 0)  # Green
        else:
            np[i] = (0, 0, 0)  # Off
    np.write()
    print("Updated NeoPixel with", num_chickens, "chickens")

# Function to calculate number of chickens
def calculate_chickens(weight):
    chicken_weight = 4.0  # weight of one chicken in kg
    return int(weight / chicken_weight)

# Main function
def main():
    connect_wifi()
    client = connect_mqtt()

    while True:
        hx711.tare()
        weight = hx711.get_units(10)
        num_chickens = calculate_chickens(weight)

        # Send weight data to MQTT
        weight_data = ujson.dumps({"weight": weight, "chickens": num_chickens})
        client.publish(WEIGHT_TOPIC, weight_data)
        print("Published weight data to MQTT:", weight_data)

        # Update NeoPixel
        update_neopixel(num_chickens)

        # Update OLED display
        oled.fill(0)
        oled.text(f"Weight: {weight:.2f} kg", 0, 0)
        oled.text(f"Chickens: {num_chickens}", 0, 10)

        if num_chickens < 16:
            oled.text(f"{16 - num_chickens} IS MISSING!", 0, 20)
        else:
            oled.text("ALL CHICKENS:)", 0, 20)

        oled.show()
        
        client.check_msg()
        sleep(10)

if __name__ == "__main__":
    main()