import machine
from machine import Pin, PWM
from time import sleep_ms, ticks_ms, ticks_diff

buzz = PWM(Pin(13))
button = Pin(16, Pin.IN)

lcd_rs = Pin(12, Pin.OUT)
lcd_e = Pin(11, Pin.OUT)
lcd_d4 = Pin(10, Pin.OUT)
lcd_d5 = Pin(9, Pin.OUT)
lcd_d6 = Pin(8, Pin.OUT)
lcd_d7 = Pin(7, Pin.OUT)

import lcd
lcd.init(lcd_rs, lcd_e, lcd_d4, lcd_d5, lcd_d6, lcd_d7)
lcd.clear()

start_time = ticks_ms()
buzzer_active = False
previous_button = 0

# MAIN LOOP
while True:
    elapsed = ticks_diff(ticks_ms(), start_time)
    
    # COUNTDOWN LOGIC
    if elapsed < 5000:
        if elapsed < 1000:
            lcd.clear()
            lcd.setCursor(0, 0)
            lcd.print("Next reminder in:")
            lcd.setCursor(7, 1)
            lcd.print("5")
        elif elapsed < 2000:
            lcd.clear()
            lcd.setCursor(0, 0)
            lcd.print("Next reminder in:")
            lcd.setCursor(7, 1)
            lcd.print("4")
        elif elapsed < 3000:
            lcd.clear()
            lcd.setCursor(0, 0)
            lcd.print("Next reminder in:")
            lcd.setCursor(7, 1)
            lcd.print("3")
        elif elapsed < 4000:
            lcd.clear()
            lcd.setCursor(0, 0)
            lcd.print("Next reminder in:")
            lcd.setCursor(7, 1)
            lcd.print("2")
        elif elapsed < 5000:
            lcd.clear()
            lcd.setCursor(0, 0)
            lcd.print("Next reminder in:")
            lcd.setCursor(7, 1)
            lcd.print("1")
    elif elapsed >= 5000:
        # Time to activate buzzer and message
        if not buzzer_active:
            buzzer_active = True
            buzz.freq(523)
            buzz.duty_u16(32768)
            
            lcd.clear()
            lcd.setCursor(0, 0)
            lcd.print("Time to drink")
            lcd.setCursor(0, 1)
            lcd.print("water! Press btn")
    
    # BUTTON LOGIC
    current_button = button.value()
    
    if current_button == 1 and previous_button == 0:
        if buzzer_active:
            # Stop buzzer and reset everything
            buzz.duty_u16(0)
            buzzer_active = False
            start_time = ticks_ms()  # Reset timer
            
            # Show countdown starting again
            lcd.clear()
            lcd.setCursor(0, 0)
            lcd.print("Next reminder in:")
            lcd.setCursor(7, 1)
            lcd.print("5")
    
    previous_button = current_button
    
    sleep_ms(100)