#!/usr/bin/env python3
"""
xTool SafetyPro AP2 — BLE control (F0F7 / M-code protocol).

Frame:  0xF0 | id0 id1 id2 | 0x01 | seq | <ASCII M-code> | 0x0A | (sum(body)&0x7F) | 0xF7
        body = id[3] + 0x01 + seq + cmd + 0x0A ; checksum over body (0xF0 excluded)
        seq = 1-byte counter from 0, +1 per frame.
AP2 V2 id = 45 73 60  (V3/Max = 4C 73 6B).

CONTROL COMMANDS (AP2 V2 / PURIFIER_V2), confirmed by firmware RE + live power meter:
  set speed/gear : M9039 A<n>   n=0 off, 1..4 = speeds (A1~46W, A2~85W, A3~110W, A4~150W max)
  identify       : M99          status : M9033        buzzer : M9046 F0 / F1
  (auto mode M9039 D1 C<n> exists per the V3 community repo; untested on V2)
Transport: AP2 advertises connectably whenever powered (no pairing/bonding needed);
NUS RX 6e400002 write, TX 6e400003 notify.

Modes:
    frames                 print example frames (offline)
    set <n> / fan <n>      set gear A<n> (0=off..4=max), watch power
    sweep                  step A1..A4 then off, watch power curve
    off                    M9039 A0
    status                 M99 + M9033
"""
import asyncio
import json
import os
import sys
import time
import urllib.request

from bleak import BleakClient, BleakScanner

# Set these for your setup via env vars (or a .env you source); the committed
# defaults are dummies. AP2_ADDR is your AP2's BLE MAC (run `find` to discover it).
SHELLY = os.environ.get("AP2_SHELLY_IP", "192.168.1.50")
AP2_ADDR = os.environ.get("AP2_ADDR", "AA:BB:CC:DD:EE:FF")
NUS = "6e400001-b5a3-f393-e0a9-e50e24dcca9e"
RX = "6e400002-b5a3-f393-e0a9-e50e24dcca9e"
TX = "6e400003-b5a3-f393-e0a9-e50e24dcca9e"
PREFIX = bytes([0x45, 0x73, 0x60])   # PURIFIER_V2 (AP2)


def build_frame(cmd: str, seq: int = 0, prefix: bytes = PREFIX) -> bytes:
    body = bytearray(prefix)
    body += bytes([0x01, seq & 0xFF])
    body += cmd.encode("ascii")
    body.append(0x0A)
    chk = sum(body) & 0x7F
    return bytes([0xF0]) + bytes(body) + bytes([chk, 0xF7])


def asc(b):
    return "".join(chr(c) if 32 <= c < 127 else "." for c in b)


def shelly_power():
    try:
        d = json.loads(urllib.request.urlopen(f"http://{SHELLY}/sensor/power", timeout=4).read())
        return float(d.get("value", 0.0) or 0.0)
    except Exception:
        return -1.0


def shelly_switch(on: bool):
    a = "turn_on" if on else "turn_off"
    try:
        req = urllib.request.Request(f"http://{SHELLY}/switch/Relay/{a}", data=b"", method="POST")
        urllib.request.urlopen(req, timeout=5).read()
        return True
    except Exception as e:  # noqa: BLE001
        print(f"  shelly {a} failed: {e}"); return False


async def watch_power(label, secs=24.0):
    print(f"    [power] watching {secs:.0f}s ({label}); fan ON ~= >100W")
    t0 = time.time(); hi = 0.0
    while time.time() - t0 < secs:
        p = shelly_power(); hi = max(hi, p)
        print(f"      +{time.time()-t0:4.0f}s  {p:6.1f} W")
        await asyncio.sleep(2.5)
    print(f"    [power] peak {hi:.1f} W")
    return hi


async def find_ap2():
    dev = await BleakScanner.find_device_by_address(AP2_ADDR, timeout=10.0)
    if dev:
        return dev
    found = await BleakScanner.discover(timeout=8.0, return_adv=True)
    for d, adv in sorted(found.values(), key=lambda kv: kv[1].rssi or -999, reverse=True):
        if not (adv.local_name or d.name):
            return d
    return None


def _mk_send(client, notifs):
    st = {"seq": 0}

    async def send(cmd, wait=1.5):
        f = build_frame(cmd, st["seq"])
        nb = len(notifs)
        try:
            await client.write_gatt_char(RX, f, response=True)
            tag = "ACK"
        except Exception as e:  # noqa: BLE001
            tag = f"ERR {type(e).__name__}"
        print(f"  -> seq{st['seq']:<3} {cmd:<12} {f.hex(' ')}  [{tag}]")
        st["seq"] = st["seq"] + 1 if st["seq"] < 0x7E else 1
        await asyncio.sleep(wait)
        return len(notifs) > nb

    return send


def _mk_cb(notifs):
    def cb(_, data):
        b = bytes(data)
        notifs.append((time.time(), b))
        print(f"    <- NOTIFY {b.hex(' ')} |{asc(b)}|")
    return cb


async def persist_test():
    print("=== PERSISTENCE TEST (bond -> power-cycle -> reconnect w/o button) ===")
    print("Phase 1: connect in pairing mode and BOND")
    dev = await find_ap2()
    if not dev:
        print("!! AP2 not found. Hold power 5s (link LED blinking) first."); return 1
    notifs = []
    async with BleakClient(dev, timeout=15.0) as c:
        if NUS not in [s.uuid.lower() for s in c.services]:
            print("!! not the AP2."); return 1
        await c.start_notify(TX, _mk_cb(notifs))
        try:
            await c.pair()
            print("  pair() returned (WinRT reports paired/bonded)")
        except Exception as e:  # noqa: BLE001
            print(f"  pair() failed: {type(e).__name__}: {e}")
        send = _mk_send(c, notifs)
        await send("M99")
        await send("M9039 A1"); await send("M9039 W2")  # brief on, proves bonded control
        await asyncio.sleep(2)
        await send("M9039 W0"); await send("M9039 A0")
        await c.stop_notify(TX)
    print("  Phase 1 done (disconnected). Bond should be saved on the AP2.\n")

    print("Phase 2: power-cycle the AP2 via Shelly (simulates normal power-up, NO button)")
    shelly_switch(False); print("  power OFF"); await asyncio.sleep(8)
    shelly_switch(True);  print("  power ON; waiting ~16s for boot..."); await asyncio.sleep(16)

    print("\nPhase 3: try to reconnect WITHOUT pressing the pairing button")
    dev2 = await BleakScanner.find_device_by_address(AP2_ADDR, timeout=25.0)
    if not dev2:
        print("  >> AP2 is NOT advertising after power-cycle (no button).")
        print("     => WinRT bond did not enable auto-advertise. We'll need a proper bond")
        print("        from a real BLE stack (nRF5340 DK / ESP32) or directed-adv reconnect.")
        return 0
    print(f"  >> AP2 IS advertising without the button! ({dev2.address}) connecting...")
    notifs2 = []
    async with BleakClient(dev2, timeout=15.0) as c2:
        await c2.start_notify(TX, _mk_cb(notifs2))
        s2 = _mk_send(c2, notifs2)
        got = await s2("M99")
        await s2("M9039 A1"); await s2("M9039 W3")
        await watch_power("reconnect W3", 18)
        await s2("M9039 W0"); await s2("M9039 A0")
        await c2.stop_notify(TX)
    print("\n  *** PERSISTENCE CONFIRMED: controlled the AP2 after power-cycle, no button! ***"
          if notifs2 else "\n  reconnected but no telemetry — check power trace above.")
    return 0


async def find_mode():
    """Discover the AP2's MAC: it's a blank-name device exposing Nordic UART."""
    print("Scanning for the AP2 (blank device exposing Nordic UART)...")
    found = await BleakScanner.discover(timeout=8.0, return_adv=True)
    cands = [(d, a) for d, a in found.values() if not (a.local_name or d.name)]
    cands.sort(key=lambda x: x[1].rssi if x[1].rssi is not None else -999, reverse=True)
    for d, a in cands[:10]:
        try:
            async with BleakClient(d, timeout=8.0) as c:
                if NUS in [s.uuid.lower() for s in c.services]:
                    print(f"\n  FOUND AP2: {d.address}  (RSSI {a.rssi})")
                    print("  -> use this as ble_client mac_address in your ESPHome YAML")
                    return 0
        except Exception:
            continue
    print("  AP2 not found — make sure it is powered and within range.")
    return 1


async def main():
    mode = sys.argv[1] if len(sys.argv) > 1 else "set"
    if mode == "find":
        return await find_mode()
    if mode == "persist":
        return await persist_test()
    arg = sys.argv[2] if len(sys.argv) > 2 else None

    if mode == "frames":
        for c in ["M99", "M9033", "M9039 A0", "M9039 A1", "M9039 A2",
                  "M9039 A3", "M9039 A4"]:
            print(f"{c:<12}: {build_frame(c).hex(' ')}")
        return

    print("Finding AP2 (must be in pairing mode)...")
    dev = await find_ap2()
    if not dev:
        print("!! AP2 not found. Hold power 5s (link LED blinking) and retry.")
        return 1

    notifs = []

    def cb(_, data):
        b = bytes(data)
        notifs.append((time.time(), b))
        print(f"    <- NOTIFY {b.hex(' ')} |{asc(b)}|")

    async with BleakClient(dev, timeout=15.0) as c:
        print(f"connected {c.is_connected}")
        if NUS not in [s.uuid.lower() for s in c.services]:
            print("!! no Nordic UART — not the AP2."); return 1
        await c.start_notify(TX, cb)
        print("  notify enabled")
        seq = 0

        async def send(cmd, wait=1.5):
            nonlocal seq
            f = build_frame(cmd, seq)
            nb = len(notifs)
            try:
                await c.write_gatt_char(RX, f, response=True)
                print(f"  -> seq{seq:<3} {cmd:<12} {f.hex(' ')}  [ACK]")
            except Exception as e:  # noqa: BLE001
                print(f"  -> seq{seq:<3} {cmd:<12} {f.hex(' ')}  WRITE-ERR {type(e).__name__}: {e}")
            seq = (seq + 1) if seq < 0x7E else 1
            await asyncio.sleep(wait)
            return len(notifs) > nb

        # handshake (resend M99 until reply)
        for _ in range(3):
            if await send("M99", 1.5):
                print("  >> M99 reply — online"); break
        else:
            print("  !! no M99 reply")

        n = (arg or "4")
        if mode in ("fan", "set"):       # set gear A<n> (0=off, 1..4)
            await send(f"M9039 A{n}")
            await watch_power(f"gear A{n}", 26)
        elif mode == "sweep":            # step through all speeds
            for g in [1, 2, 3, 4]:
                await send(f"M9039 A{g}")
                await watch_power(f"gear A{g}", 16)
            await send("M9039 A0")
        elif mode == "off":
            await send("M9039 A0")
            await watch_power("off", 14)
        elif mode == "status":
            await send("M9033", 2.0)
        await asyncio.sleep(1.0)
        await c.stop_notify(TX)

    print(f"\nDone. {len(notifs)} notification(s).")
    return 0


class _Tee:
    def __init__(self, *s): self.s = s
    def write(self, x):
        for st in self.s: st.write(x); st.flush()
    def flush(self):
        for st in self.s: st.flush()


if __name__ == "__main__":
    os.makedirs("logs", exist_ok=True)
    _logf = open(os.path.join("logs", "ap2_ble.log"), "a", encoding="utf-8")
    _logf.write(f"\n===== run {time.strftime('%Y-%m-%d %H:%M:%S')} args={sys.argv[1:]} =====\n")
    sys.stdout = _Tee(sys.__stdout__, _logf)
    sys.exit(asyncio.run(main()) or 0)