#!/usr/bin/env python3
"""
BLE Reboot for WellNuo WP sensors
Reboot sensor to clear stuck WiFi state
"""

import asyncio
from bleak import BleakClient, BleakScanner

# Sensor BLE UUIDs
SERVICE_UUID = "4fafc201-1fb5-459e-8fcc-c5c9c331914b"
CHAR_UUID = "beb5483e-36e1-4688-b7f5-ea07361b26a8"

DEVICE_PIN = "7856"

response_data = None
response_event = asyncio.Event()

def notification_handler(sender, data):
    global response_data
    decoded = data.decode('utf-8', errors='replace')
    print(f"  [NOTIFY] {decoded}")
    response_data = decoded
    response_event.set()

async def send_and_wait(client, command, timeout=10):
    global response_data
    response_data = None
    response_event.clear()

    print(f"\n>>> Sending: {command}")
    await client.write_gatt_char(CHAR_UUID, command.encode('utf-8'))

    try:
        await asyncio.wait_for(response_event.wait(), timeout=timeout)
        return response_data
    except asyncio.TimeoutError:
        try:
            data = await client.read_gatt_char(CHAR_UUID)
            decoded = data.decode('utf-8', errors='replace')
            print(f"  [READ] {decoded}")
            return decoded
        except:
            print(f"  [TIMEOUT]")
            return None

async def main():
    print("=" * 60)
    print("WellNuo Sensor Reboot Tool")
    print("=" * 60)

    print("\nScanning for sensors...")
    devices = await BleakScanner.discover(timeout=5.0)

    wp_device = None
    for d in devices:
        if d.name and d.name.startswith("WP_"):
            wp_device = d
            print(f"  Found: {d.name} ({d.address})")
            break

    if not wp_device:
        print("No WP sensor found!")
        return

    print(f"\nConnecting to {wp_device.name}...")
    async with BleakClient(wp_device.address) as client:
        print("Connected!")

        await client.start_notify(CHAR_UUID, notification_handler)

        # Read initial status
        print("\n--- Current status ---")
        data = await client.read_gatt_char(CHAR_UUID)
        print(f"Status: {data.decode('utf-8', errors='replace')}")

        # Unlock
        print("\n--- Step 1: Unlock ---")
        response = await send_and_wait(client, f"pin|{DEVICE_PIN}")
        if not response or "ok" not in response.lower():
            print("Unlock failed!")
            return
        print("✓ Unlocked!")

        # Get current WiFi status
        print("\n--- Step 2: Current WiFi status ---")
        await send_and_wait(client, "a")  # GET_WIFI_STATUS command

        # Reboot
        print("\n--- Step 3: Sending REBOOT command ---")
        await send_and_wait(client, "s", timeout=3)  # REBOOT command is 's'
        print("Reboot command sent!")

        await client.stop_notify(CHAR_UUID)

    print("\n--- Waiting 10 seconds for sensor to reboot ---")
    await asyncio.sleep(10)

    # Try to reconnect
    print("\n--- Attempting to reconnect ---")
    devices = await BleakScanner.discover(timeout=10.0)

    wp_device = None
    for d in devices:
        if d.name and d.name.startswith("WP_"):
            wp_device = d
            print(f"  Found: {d.name} ({d.address})")
            break

    if not wp_device:
        print("Sensor not found after reboot - may still be booting")
        return

    print(f"\nConnecting to {wp_device.name}...")
    async with BleakClient(wp_device.address) as client:
        print("Connected!")

        await client.start_notify(CHAR_UUID, notification_handler)

        # Read status after reboot
        print("\n--- Status after reboot ---")
        data = await client.read_gatt_char(CHAR_UUID)
        status = data.decode('utf-8', errors='replace')
        print(f"Status: {status}")

        # Unlock
        print("\n--- Unlock ---")
        await send_and_wait(client, f"pin|{DEVICE_PIN}")

        # Try WiFi list
        print("\n--- WiFi List ---")
        response = await send_and_wait(client, "w", timeout=20)  # GET_WIFI_LIST is 'w'
        if response:
            print(f"Response: {response}")
            if "|w|" in response:
                parts = response.split("|w|")
                if len(parts) > 1:
                    count_and_networks = parts[1]
                    items = count_and_networks.split("|")
                    print(f"\nWiFi networks ({items[0]} found):")
                    for item in items[1:]:
                        if "," in item:
                            ssid, rssi = item.rsplit(",", 1)
                            print(f"  📶 {ssid} (signal: {rssi})")

        await client.stop_notify(CHAR_UUID)

    print("\n" + "=" * 60)
    print("Reboot complete!")
    print("=" * 60)

if __name__ == "__main__":
    asyncio.run(main())