"""
ble_scanner.py - Discover and dump BLE services/characteristics from nearby devices.

Run this first to find your KAIYU LED controller and identify its BLE services.

Usage:
    python scanner/ble_scanner.py                          # scan only
    python scanner/ble_scanner.py --connect 7C:77:6C:6B:8E:C8   # scan then connect immediately
    python scanner/ble_scanner.py --probe 7C:77:6C:6B:8E:C8     # scan then probe with test commands
    python scanner/ble_scanner.py --filter KAIYU          # scan with name filter

IMPORTANT: --connect and --probe now re-scan first so the device is still
           advertising when we try to connect. Do not scan separately then connect.
"""

import asyncio
import argparse
import json
from bleak import BleakScanner, BleakClient


async def scan_and_get_device(address: str, timeout: float = 10.0):
    """
    Scan for a specific address and return the BLEDevice object while it's
    still in the scanner's cache (i.e. still advertising).
    Returns None if not found.
    """
    print(f"Scanning for {address} ({timeout}s)...")
    target = None

    def detection_callback(device, adv):
        nonlocal target
        if device.address.upper() == address.upper():
            target = device
            print(f"Found: {device.address}  name={device.name or '(unknown)'}  rssi={adv.rssi}")

    scanner = BleakScanner(detection_callback=detection_callback)
    await scanner.start()
    await asyncio.sleep(timeout)
    await scanner.stop()
    return target


async def scan_devices(filter_name: str = None, timeout: float = 10.0):
    """Scan for nearby BLE devices and print them."""
    print(f"Scanning for BLE devices ({timeout}s)...\n")

    devices = await BleakScanner.discover(timeout=timeout, return_adv=True)

    found = []
    for device, adv in devices.values():
        name = device.name or "(unknown)"
        if filter_name and filter_name.lower() not in name.lower():
            continue
        found.append((device, adv))

    if not found:
        print("No devices found." if not filter_name else f"No devices matching '{filter_name}' found.")
        return []

    print(f"{'#':<4} {'Name':<30} {'Address':<20} {'RSSI':<8} {'Manufacturer'}")
    print("-" * 80)
    for i, (device, adv) in enumerate(found):
        name = device.name or "(unknown)"
        mfr = ""
        if adv.manufacturer_data:
            mfr = " ".join(f"{k:#06x}:{v.hex()}" for k, v in adv.manufacturer_data.items())
        print(f"{i:<4} {name:<30} {device.address:<20} {adv.rssi:<8} {mfr}")

    print()
    return found


async def inspect_device(address: str, timeout: float = 10.0):
    """Scan for device, then immediately connect and dump all services/characteristics."""

    device = await scan_and_get_device(address, timeout=timeout)
    if device is None:
        print(f"\nDevice {address} not found during scan.")
        print("Make sure it is powered on and not connected to another device (e.g. your phone).")
        return

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

    async with BleakClient(device) as client:
        print(f"Connected!\n")
        print("=" * 60)
        print("SERVICES & CHARACTERISTICS")
        print("=" * 60)

        dump = {}

        for service in client.services:
            print(f"\nSERVICE: {service.uuid}")
            print(f"  Description: {service.description}")
            dump[service.uuid] = {
                "description": service.description,
                "characteristics": {}
            }

            for char in service.characteristics:
                props = ", ".join(char.properties)
                print(f"  CHAR: {char.uuid}")
                print(f"    Properties : {props}")
                print(f"    Handle     : {char.handle}")

                value_hex = None
                value_str = None

                if "read" in char.properties:
                    try:
                        value = await client.read_gatt_char(char.uuid)
                        value_hex = value.hex()
                        try:
                            value_str = value.decode("utf-8")
                        except Exception:
                            value_str = None
                        print(f"    Value (hex): {value_hex}")
                        if value_str:
                            print(f"    Value (str): {value_str}")
                    except Exception as e:
                        print(f"    Value      : (read error: {e})")

                dump[service.uuid]["characteristics"][char.uuid] = {
                    "properties": list(char.properties),
                    "handle": char.handle,
                    "value_hex": value_hex,
                    "value_str": value_str,
                }

                for descriptor in char.descriptors:
                    print(f"    DESC: {descriptor.uuid} handle={descriptor.handle}")

        # Save dump to JSON
        dump_file = f"scanner/device_{address.replace(':', '-')}.json"
        with open(dump_file, "w") as f:
            json.dump(dump, f, indent=2)
        print(f"\nDump saved to {dump_file}")

        print("\n" + "=" * 60)
        print("LIKELY COMMAND CHARACTERISTICS (writable):")
        print("=" * 60)
        for svc_uuid, svc in dump.items():
            for char_uuid, char in svc["characteristics"].items():
                props = char["properties"]
                if "write" in props or "write-without-response" in props:
                    print(f"  -> {char_uuid}  [{', '.join(props)}]")


async def try_known_protocols(address: str, timeout: float = 10.0):
    """Scan for device then probe it with known protocol commands."""

    device = await scan_and_get_device(address, timeout=timeout)
    if device is None:
        print(f"\nDevice {address} not found during scan.")
        print("Make sure it is powered on and not connected to your phone.")
        return

    KNOWN_WRITE_UUIDS = [
        "0000ffe1-0000-1000-8000-00805f9b34fb",  # LEDBLE / MagicLight
        "0000fff3-0000-1000-8000-00805f9b34fb",  # ELK-BLEDOM
        "0000ffd9-0000-1000-8000-00805f9b34fb",  # Some controllers
    ]

    test_commands = {
        "Turn ON  (7e)":      bytes.fromhex("7e000400000000ffef"),
        "Turn OFF (7e)":      bytes.fromhex("7e000400000000 00ef".replace(" ", "")),
        "Red      (7e)":      bytes.fromhex("7e000503ff000000ef"),
        "Green    (7e)":      bytes.fromhex("7e00050300ff0000ef"),
        "Blue     (7e)":      bytes.fromhex("7e0005030000ff00ef"),
        "White    (7e)":      bytes.fromhex("7e000503ffffffff ef".replace(" ", "")),
        "Turn ON  (56/cc)":   bytes.fromhex("cc2333"),
        "Turn OFF (56/cc)":   bytes.fromhex("cc2433"),
        "Red      (56)":      bytes.fromhex("56ff000000f0aa"),
        "Green    (56)":      bytes.fromhex("5600ff000000f0aa".replace(" ", "")),
        "Blue     (56)":      bytes.fromhex("560000ff00f0aa"),
    }

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

        writable = None
        for svc in client.services:
            for char in svc.characteristics:
                if char.uuid in KNOWN_WRITE_UUIDS:
                    writable = char.uuid
                    print(f"Found known write characteristic: {char.uuid}")
                    break

        if not writable:
            print("No known UUID found. Available writable characteristics:")
            for svc in client.services:
                for char in svc.characteristics:
                    if "write" in char.properties or "write-without-response" in char.properties:
                        print(f"  {char.uuid}  [{', '.join(char.properties)}]")
                        if not writable:
                            writable = char.uuid
                            print(f"  ^ Using this one")

        if not writable:
            print("No writable characteristic found at all!")
            return

        print(f"\nSending test commands to: {writable}")
        print("Watch the lights and press Enter for each one.\n")

        for name, cmd in test_commands.items():
            input(f"  Press Enter to send '{name}' ({cmd.hex()})...")
            try:
                await client.write_gatt_char(writable, cmd, response=False)
                print(f"    Sent OK")
            except Exception as e:
                # Try with response=True as fallback
                try:
                    await client.write_gatt_char(writable, cmd, response=True)
                    print(f"    Sent OK (with response)")
                except Exception as e2:
                    print(f"    Error: {e2}")


def main():
    parser = argparse.ArgumentParser(description="KAIYU BLE Scanner & Inspector")
    parser.add_argument("--filter", "-f", help="Filter devices by name")
    parser.add_argument("--connect", "-c", help="Scan then connect to address and dump services")
    parser.add_argument("--probe", "-p", help="Scan then probe address with known commands")
    parser.add_argument("--timeout", "-t", type=float, default=10.0, help="Scan timeout in seconds (default 10)")
    args = parser.parse_args()

    if args.connect:
        asyncio.run(inspect_device(args.connect, timeout=args.timeout))
    elif args.probe:
        asyncio.run(try_known_protocols(args.probe, timeout=args.timeout))
    else:
        async def scan_and_pick():
            found = await scan_devices(filter_name=args.filter, timeout=args.timeout)
            if found:
                print("Tip: run with --connect <ADDRESS> to inspect a device")
                print("     run with --probe <ADDRESS> to test known protocols")
                print("\nIMPORTANT: run --connect immediately after the device appears")
                print("           in a scan, or it may stop advertising.")
        asyncio.run(scan_and_pick())


if __name__ == "__main__":
    main()