from utime import sleep, localtime
from hx711 import HX711
import network
import machine
import neopixel
from machine import Pin, I2C
from ssd1306 import SSD1306_I2C

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

# NeoPixel configuration
NEOPIXEL_PIN = 14  # Pin on ESP32
NUM_PIXELS = 16    # Number of NeoPixels (one for each chicken)
np = neopixel.NeoPixel(machine.Pin(NEOPIXEL_PIN), NUM_PIXELS)

# HX711 configuration
hx711 = HX711(d_out=4, pd_sck=5)
hx711.set_scale(1000)  # Adjust this scale value based on calibration
hx711.tare()

# WiFi configuration
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"

# 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")

connect_wifi()

# Helper functions
def update_neopixel(chickens):
    for i in range(NUM_PIXELS):
        if i < chickens:
            # Green color for chickens present
            np[i] = (0, 255, 0)
        else:
            # Red color for empty spaces
            np[i] = (255, 0, 0)
    np.write()

def update_oled(chickens):
    oled.fill(0)
    oled.text("Chickens in Coop:", 0, 0)
    oled.text(str(chickens), 0, 16)
    if chickens < 16:
        oled.text(f"{16 - chickens} IS MISSING!", 0, 32)
    else:
        oled.text("ALL CHICKENS PRESENT :)", 0, 32)
    oled.show()

# Main loop
while True:
    hx711.power_on()
    
    if hx711.is_ready():
        weight = hx711.get_weight()
        chickens = max(0, weight // 4000)  # Mass of 1 chicken is 4 kg, so divide by 4000 grams
        
        # Ensure the number of chickens does not exceed the number of NeoPixels
        if chickens > NUM_PIXELS:
            chickens = NUM_PIXELS

        update_neopixel(chickens)
        update_oled(chickens)

    sleep(1)