from machine import Pin, PWM
from neopixel import NeoPixel
import time
from time import sleep
import ujson
from umqtt.simple import MQTTClient
import network

# --- Configuration ---
SSID = "Wokwi-GUEST"
PASSWORD = ""
MQTT_BROKER = "mqtt.netpie.io"
MQTT_PORT = 1883  # Added this back in case it was missing
MQTT_CLIENT_ID = "5438c716-4b0e-49f5-b8cf-5f8e88e26fb2"
MQTT_USER = "AfpS2KHTJphuJbtUmwxwi6shcSPY6kCx"
MQTT_PASSWORD = "2dRrD9PUy9QNwQoYLizvgKLZvcPfX4km"

# Ensure these match exactly
MQTT_TOPIC_SHADOW = "@shadow/data/update"
MQTT_TOPIC_CONTROL = "@msg/operator"

# Hardware Setup
red = Pin(4, Pin.OUT)
green = Pin(0, Pin.OUT)
blue = Pin(2, Pin.OUT)
buzzer = PWM(Pin(14), Pin.OUT)
buzzer.duty(0)
npPin = Pin(23, Pin.OUT)
np = NeoPixel(npPin, 1)

alarm_active = False
current_led_state = "OFF"

def connect_wifi():
    wlan = network.WLAN(network.STA_IF)
    wlan.active(True)
    wlan.connect(SSID, PASSWORD)
    print("Connecting to WiFi", end="")
    while not wlan.isconnected():
        print(".", end="")
        sleep(0.5)
    print("\nWiFi Connected")

def update_shadow(data_dict):
    """Sends data to NETPIE Shadow"""
    try:
        # Use the global variable explicitly
        payload = ujson.dumps({"data": data_dict})
        client.publish(MQTT_TOPIC_SHADOW, payload)
    except Exception as e:
        print("Shadow Update Error:", e)

def mqtt_callback(topic, msg):
    global alarm_active
    command = msg.decode()
    print(f"Command received: {command}")
    # This responds to the button publish from NETPIE
    if command == "OFF" or command == "0":
        alarm_active = False

def set_light(color):
    global current_led_state
    red.value(0); green.value(0); blue.value(0)
    
    if color == 'red':
        red.value(1); np[0] = (255, 0, 0); current_led_state = "RED"
    elif color == 'green':
        green.value(1); np[0] = (0, 255, 0); current_led_state = "GREEN"
    elif color == 'blue':
        blue.value(1); np[0] = (0, 0, 255); current_led_state = "BLUE"
    else:
        np[0] = (0, 0, 0); current_led_state = "OFF"
    
    np.write()
    update_shadow({"led_status": current_led_state})

# --- Initialize ---
connect_wifi()
client = MQTTClient(MQTT_CLIENT_ID, MQTT_BROKER, port=MQTT_PORT, user=MQTT_USER, password=MQTT_PASSWORD)
client.set_callback(mqtt_callback)
client.connect()
client.subscribe(MQTT_TOPIC_CONTROL)

durations = []
channels = ['red', 'green', 'blue']
for i in range(3):
    # Wrap in try-except to prevent crash on empty input
    try:
        m = float(input(f"Minutes for {channels[i]}: "))
    except:
        m = 0.1 # Default to 6 seconds if input fails
    durations.append([time.time() + (m * 60), channels[i]])

print("Monitoring and updating Shadow...")



while True:
    client.check_msg() # Listen for NETPIE button clicks
    now = time.time()
    
    for i in range(len(durations)):
        target_time, color = durations[i]
        if target_time == -1: continue

        remaining = int(target_time - now)
        
        # Update Shadow every 5 seconds
        if remaining > 0:
            if remaining % 5 == 0:
                print(f"{color} countdown: {remaining}s")
                update_shadow({f"timer_{color}": remaining})
        else:
            print(f"ALARM: {color} - Check Dashboard to Turn Off")
            alarm_active = True
            
            # Loop until 'OFF' message received from MQTT
            while alarm_active:
                client.check_msg() 
                set_light(color)
                buzzer.freq(1000); buzzer.duty(512)
                sleep(0.5)
                # Quick check inside blink
                client.check_msg()
                buzzer.duty(0); sleep(0.5)
            
            # Reset after MQTT button pressed
            print("Alarm deactivated via NETPIE")
            buzzer.duty(0)
            set_light("OFF")
            durations[i][0] = -1 
            update_shadow({f"timer_{color}": 0, "status": f"{color}_done"})

    sleep(0.1) # Faster response time