import network
import time
import ubinascii # For generating a unique client ID
import machine
from umqtt.simple import MQTTClient
import uasyncio # For non-blocking operation

# --- Configuration ---
WIFI_SSID = "YOUR_WIFI_SSID"
WIFI_PASSWORD = "YOUR_WIFI_PASSWORD"

MQTT_BROKER = "YOUR_MQTT_BROKER_IP_OR_HOSTNAME" # e.g., "192.168.1.100" or "broker.hivemq.com"
MQTT_PORT = 1883 # Default MQTT port
MQTT_USER = None # Set to your MQTT username if required, otherwise None
MQTT_PASSWORD = None # Set to your MQTT password if required, otherwise None
MQTT_TOPIC = b"pico/test" # The topic to subscribe to (bytes format)
# Generate a unique client ID based on the Pico's MAC address
CLIENT_ID = b"pico_w_" + ubinascii.hexlify(machine.unique_id())

# LED Pin (Onboard LED on Pico W is controlled differently)
# It's usually accessed via the name "LED"
led = machine.Pin("LED", machine.Pin.OUT)

# Global flag/state for LED control
message_received_flag = False
led_solid_task = None # To hold the asyncio task for solid LED

# --- WiFi Connection ---
def connect_wifi(ssid, password):
    wlan = network.WLAN(network.STA_IF)
    wlan.active(True)
    if not wlan.isconnected():
        print('Connecting to WiFi...')
        wlan.connect(ssid, password)
        # Wait for connection with a timeout
        max_wait = 15
        while max_wait > 0:
            if wlan.status() < 0 or wlan.status() >= 3:
                break
            max_wait -= 1
            print('.')
            time.sleep(1)

    if wlan.isconnected():
        print('WiFi Connected!')
        print('IP Address:', wlan.ifconfig()[0])
        return wlan
    else:
        print('WiFi Connection Failed!')
        # Optional: Blink LED rapidly to indicate failure
        for _ in range(10):
            led.toggle()
            time.sleep(0.1)
        led.off() # Ensure LED is off after failure indication
        return None

# --- MQTT Callback ---
# This function is called when a message is received on the subscribed topic
def mqtt_callback(topic, msg):
    global message_received_flag, led_solid_task
    print(f"Received message on topic '{topic.decode()}': {msg.decode()}")
    message_received_flag = True
    # If a solid LED task is already running, cancel it before starting a new one
    if led_solid_task:
        try:
            led_solid_task.cancel()
        except Exception as e:
            print(f"Error cancelling previous LED task: {e}") # Should ideally not happen often
    # Schedule the solid LED task to run
    led_solid_task = uasyncio.create_task(handle_solid_led())

# --- Async Task for Solid LED ---
async def handle_solid_led():
    global message_received_flag
    print("LED Solid ON")
    led.on()
    await uasyncio.sleep(1) # Non-blocking sleep for 1 second
    print("LED Solid OFF")
    led.off()
    message_received_flag = False # Reset flag after handling

# --- Main Async Function ---
async def main():
    global led_solid_task # Allow modification

    print("Starting Pico W MQTT Client...")

    # 1. Connect to WiFi
    wlan = connect_wifi(WIFI_SSID, WIFI_PASSWORD)
    if not wlan:
        print("Cannot proceed without WiFi. Halting.")
        # Optional: reset or enter a safe state
        return # Stop execution

    # 2. Initialize MQTT Client
    print(f"Connecting to MQTT Broker: {MQTT_BROKER} as Client ID: {CLIENT_ID.decode()}")
    client = MQTTClient(CLIENT_ID, MQTT_BROKER, port=MQTT_PORT, user=MQTT_USER, password=MQTT_PASSWORD, keepalive=60)
    client.set_callback(mqtt_callback)

    # 3. Connect to MQTT Broker
    try:
        client.connect()
        print(f"Connected to MQTT Broker and subscribing to topic: {MQTT_TOPIC.decode()}")
        client.subscribe(MQTT_TOPIC)
        print("Subscription successful.")
    except OSError as e:
        print(f"Failed to connect or subscribe to MQTT Broker: {e}")
        # Optional: blink LED differently or reset
        return # Stop execution
    except Exception as e:
        print(f"An unexpected error occurred during MQTT setup: {e}")
        return

    # 4. Main Loop - Polling and Blinking
    print("Starting main loop...")
    loop_counter = 0
    while True:
        try:
            # Blink LED on every loop iteration (unless solid LED is active)
            if not message_received_flag:
                led.toggle()

            # Check for incoming MQTT messages (non-blocking)
            client.check_msg()

            # Small delay to prevent busy-waiting and allow other tasks
            await uasyncio.sleep_ms(200) # Adjust sleep time as needed (e.g., 100-500ms)

            loop_counter += 1
            # Optional: Ping MQTT broker periodically to keep connection alive if needed
            # if loop_counter % 100 == 0: # Every ~20 seconds if sleep_ms is 200
            #     try:
            #         client.ping()
            #         print("MQTT Ping Sent")
            #     except OSError as e:
            #         print(f"MQTT Ping failed, attempting reconnect: {e}")
            #         # Add reconnect logic here if desired
            #         await uasyncio.sleep(5) # Wait before retrying loop

        except OSError as e:
            print(f"Network/MQTT Error in main loop: {e}")
            # Basic reconnect attempt (can be made more robust)
            print("Attempting to reconnect...")
            await uasyncio.sleep(5)
            try:
                client.disconnect() # Try a clean disconnect first
            except:
                pass # Ignore errors during disconnect
            try:
                client.connect()
                client.subscribe(MQTT_TOPIC)
                print("Reconnected to MQTT Broker.")
            except Exception as recon_e:
                print(f"Reconnect failed: {recon_e}")
                await uasyncio.sleep(10) # Wait longer before next attempt
        except Exception as e:
            print(f"An unexpected error occurred in main loop: {e}")
            # Consider resetting or logging more details
            await uasyncio.sleep(5)


# --- Run the Application ---
try:
    uasyncio.run(main())
except KeyboardInterrupt:
    print("Program stopped by user.")
except Exception as e:
    print(f"Fatal error: {e}")
finally:
    # Cleanup (optional but good practice)
    try:
        if 'client' in locals() and client:
            client.disconnect()
            print("MQTT client disconnected.")
    except:
        pass # Ignore cleanup errors
    led.off() # Ensure LED is off on exit
    print("Exiting.")