from machine import Pin, I2C, PWM, ADC
from utime import sleep
import dht
import ssd1306
import ultrasonic_library

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

# ===== SENSOR SETUP =====
dht_sensor = dht.DHT22(Pin(15))
ultra = ultrasonic_library.HCSR04(trigger_pin=5, echo_pin=18)

button = Pin(4, Pin.IN, Pin.PULL_UP)
leak = Pin(2, Pin.IN, Pin.PULL_UP)

# ===== POTENTIOMETER =====
pot = ADC(Pin(34))
pot.atten(ADC.ATTN_11DB)

# ===== OUTPUT =====
led_green = Pin(12, Pin.OUT)
led_red = Pin(14, Pin.OUT)
buzzer = PWM(Pin(27))
relay = Pin(26, Pin.OUT)

# ===== STARTUP =====
def startup():
    oled.fill(0)
    oled.text("SMART WATER", 10, 10)
    oled.text("SYSTEM", 30, 25)
    oled.text("Starting...", 20, 45)
    oled.show()
    sleep(2)

# ===== LEAK FUNCTION =====
def is_leak_detected():
    return leak.value() == 0

# ===== SENSOR READ =====
def read_sensors():
    try:
        dht_sensor.measure()
        temp = dht_sensor.temperature()
        hum = dht_sensor.humidity()
    except:
        temp = 0
        hum = 0

    try:
        dist = ultra.distance_cm()
        if dist < 1 or dist > 400:
            dist = -1
    except:
        dist = -1

    pot_value = pot.read()

    return temp, hum, dist, pot_value

# ===== OUTPUT CONTROL =====
def set_output(green, red, buzz, rel):
    led_green.value(green)
    led_red.value(red)

    if buzz:
        buzzer.freq(2000)
        buzzer.duty(512)
    else:
        buzzer.duty(0)

    relay.value(rel)

# ===== DISPLAY =====
def display(temp, hum, dist, pot_value, status, valve):
    oled.fill(0)
    oled.text("Temp: {}C".format(temp), 0, 0)
    oled.text("Hum: {}%".format(hum), 0, 10)

    if dist == -1:
        oled.text("Dist: ERROR", 0, 20)
    else:
        oled.text("Dist: {:.1f}cm".format(dist), 0, 20)

    oled.text("Tank: {}".format(pot_value), 0, 30)

    oled.text(status, 0, 40)
    oled.text("Valve: {}".format("ON" if valve else "OFF"), 0, 55)
    oled.show()

# ===== MAIN =====
startup()

while True:
    temp, hum, dist, pot_value = read_sensors()
    btn = button.value()

    # ===== LOGIC =====
    if btn == 0:
        status = "SYSTEM OFF"
        set_output(0, 0, 0, 0)

    elif hum > 90:
        status = "LEAK DETECTED"
        set_output(0, 1, 1, 0)

    elif dist == -1:
        status = "SENSOR ERROR"
        set_output(0, 1, 0, 0)

    # 🔵 PIPE (ULTRASONIC)
    elif dist > 15:
        status = "LOW WATER IN PIPE"
        set_output(0, 1, 1, 1)

    # 🟢 TANK (POTENTIOMETER)
    elif pot_value < 1500:
        status = "LOW WATER IN TANK"
        set_output(0, 1, 1, 1)

    elif temp < 20:
        status = "LOW TEMP"
        set_output(0, 1, 0, 1)

    else:
        status = "NORMAL"
        set_output(1, 0, 0, 1)

    display(temp, hum, dist, pot_value, status, relay.value())

    sleep(2)