#!/usr/bin/env python3
"""
BLE Sensor Setup Script for WellNuo WP sensors
Connects to sensor, unlocks it, and configures WiFi
"""

import asyncio
import sys
from bleak import BleakClient, BleakScanner

# Sensor BLE UUIDs (from app code)
SERVICE_UUID = "4fafc201-1fb5-459e-8fcc-c5c9c331914b"
CHAR_TX_UUID = "beb5483e-36e1-4688-b7f5-ea07361b26a8"  # Write to sensor
CHAR_RX_UUID = "beb5483e-36e1-4688-b7f5-ea07361b26a9"  # Read from sensor (notifications)

# Sensor details
DEVICE_ADDRESS = "14:2B:2F:81:A1:4E"  # WP_497_81a14c
DEVICE_PIN = "7856"

# Global for notification response
response_data = None
response_event = asyncio.Event()

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

async def send_command(client, command, timeout=10):
    """Send command and wait for response"""
    global response_data
    response_data = None
    response_event.clear()

    print(f"[TX] {command}")
    await client.write_gatt_char(CHAR_TX_UUID, command.encode('utf-8'))

    try:
        await asyncio.wait_for(response_event.wait(), timeout=timeout)
        return response_data
    except asyncio.TimeoutError:
        print(f"[TIMEOUT] No response after {timeout}s")
        return None

async def scan_for_sensor():
    """Scan for WP sensors"""
    print("Scanning for BLE devices...")
    devices = await BleakScanner.discover(timeout=5.0)

    wp_devices = []
    for d in devices:
        if d.name and d.name.startswith("WP_"):
            wp_devices.append(d)
            print(f"  Found: {d.name} ({d.address})")

    return wp_devices

async def get_wifi_list(client):
    """Get list of available WiFi networks"""
    print("\n=== Getting WiFi networks ===")
    response = await send_command(client, "W|list", timeout=15)
    if response:
        # Parse networks from response
        # Format: mac,xxxxx|W|list|SSID1,RSSI1|SSID2,RSSI2|...
        if "|W|list|" in response:
            parts = response.split("|W|list|")
            if len(parts) > 1:
                networks = parts[1].split("|")
                print("\nAvailable WiFi networks:")
                for i, net in enumerate(networks, 1):
                    if "," in net:
                        ssid, rssi = net.rsplit(",", 1)
                        print(f"  {i}. {ssid} (signal: {rssi})")
                return networks
    return []

async def setup_wifi(client, ssid, password):
    """Configure WiFi on sensor"""
    print(f"\n=== Setting WiFi: {ssid} ===")

    # Step 1: Unlock with PIN
    print("\n1. Unlocking sensor...")
    response = await send_command(client, f"pin|{DEVICE_PIN}")
    if not response or "ok" not in response.lower():
        print(f"ERROR: Unlock failed! Response: {response}")
        return False
    print("   Unlocked!")

    # Step 2: Send WiFi credentials
    print(f"\n2. Sending WiFi credentials...")
    print(f"   SSID: {ssid}")
    print(f"   Password: {'*' * len(password)}")

    response = await send_command(client, f"W|{ssid},{password}", timeout=20)

    if response:
        if "ok" in response.lower():
            print("   SUCCESS! WiFi configured.")
            return True
        elif "fail" in response.lower():
            print(f"   FAILED! Sensor rejected credentials.")
            print(f"   Response: {response}")
            return False

    print("   No response from sensor")
    return False

async def main():
    print("=" * 50)
    print("WellNuo Sensor WiFi Setup")
    print("=" * 50)

    # Check if WiFi credentials provided
    if len(sys.argv) < 3:
        print("\nUsage: python ble-sensor-setup.py <SSID> <PASSWORD>")
        print("\nExample:")
        print("  python ble-sensor-setup.py MyWiFi mypassword123")
        print("\nWill scan for sensors first...")

        # Just scan and show available networks
        devices = await scan_for_sensor()
        if not devices:
            print("\nNo WP sensors found. Make sure sensor is powered on and in range.")
            return

        # Connect to first sensor and get WiFi list
        device = devices[0]
        print(f"\nConnecting to {device.name}...")

        async with BleakClient(device.address) as client:
            print("Connected!")

            # Start notifications
            await client.start_notify(CHAR_RX_UUID, notification_handler)

            # Get WiFi list
            await get_wifi_list(client)

            await client.stop_notify(CHAR_RX_UUID)
        return

    ssid = sys.argv[1]
    password = sys.argv[2]

    # Scan for sensors
    devices = await scan_for_sensor()
    if not devices:
        print("\nNo WP sensors found. Make sure sensor is powered on and in range.")
        return

    # Use first found sensor or specified address
    device = devices[0]
    address = DEVICE_ADDRESS if any(d.address == DEVICE_ADDRESS for d in devices) else device.address

    print(f"\nConnecting to {address}...")

    async with BleakClient(address) as client:
        print("Connected!")
        print(f"MTU: {client.mtu_size}")

        # Start notifications
        await client.start_notify(CHAR_RX_UUID, notification_handler)

        # Setup WiFi
        success = await setup_wifi(client, ssid, password)

        if success:
            print("\n" + "=" * 50)
            print("WiFi setup SUCCESSFUL!")
            print("Sensor should now connect to the network.")
            print("=" * 50)
        else:
            print("\n" + "=" * 50)
            print("WiFi setup FAILED!")
            print("Check that:")
            print("  1. SSID is correct (case-sensitive)")
            print("  2. Password is correct")
            print("  3. WiFi network is 2.4GHz (not 5GHz)")
            print("  4. Sensor is in range of WiFi")
            print("=" * 50)

        await client.stop_notify(CHAR_RX_UUID)

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