WRITE_UUID = "0000ae3b-0000-1000-8000-00805f9b34fb"
await client.write_gatt_char(WRITE_UUID, bytes.fromhex("ACAD06A000000003BCBD"), response=False)  # level 3

devices = await BleakScanner.discover(timeout=7.0)
for d in devices:
    name = (d.name or "").strip()
    if any(name.startswith(p) for p in ["HRK-016B", "FitRx Wireles", "FitRx 7772"]):
        self.device = d
        break

self.client = BleakClient(self.device.address)
await self.client.connect()
await self.client.start_notify(self.notify_uuid, self._notification_handler)

async def set_intensity(self, level: int):
    level = max(0, min(level, 9))
    await self._send_payload(f"ACAD06A0000000{level:02d}BCBD")

async def set_timer(self, minutes: int):
    if minutes == 0:
        await self._send_payload("ACAD06A200000000BCBD")
    else:
        await self._send_payload(f"ACAD06A2010100{minutes:02X}BCBD")

async def _ramp(self, start_lvl, end_lvl, duration_sec):
    steps = abs(end_lvl - start_lvl)
    if steps == 0 or duration_sec <= 0:
        await self.massager.set_intensity(end_lvl)
        return
    interval = duration_sec / steps
    direction = 1 if end_lvl > start_lvl else -1
    current = start_lvl
    for _ in range(steps):
        current += direction
        await self.massager.set_intensity(current)
        await asyncio.sleep(interval)

async def main(cfg):
    task = asyncio.current_task()
    asyncio.get_running_loop().add_signal_handler(signal.SIGINT, task.cancel)
    await SessionOrchestrator(cfg).run()

# ...inside run():
finally:
    try:
        await self.massager.stop()
    except Exception as e:
        print(f"!! SOFTWARE STOP FAILED ({e}).")
        print("!! Turn the unit off with its power button or remove the pads.")

#!/usr/bin/env python3
"""
FitRx / HRK-016B session controller (BLE, macOS).

Runs ramp / hold / ramp-down / rest cycles using the unit's own mode presets
(1-8) and intensity levels (0-9), with the hardware timer set as a backstop.
Logs raw notification bytes from the device for study; it does not interpret
them, because the device's status format hasn't been confirmed yet.

The real emergency stop is the unit's power button or removing a pad.
The software stop below depends on Bluetooth working.
"""

import argparse
import asyncio
import signal
import sys
import time
from dataclasses import dataclass

from bleak import BleakClient, BleakScanner

# GATT UUIDs from the decompiled app
UUID_MAP = {
    "standard": {
        "service": "0000ae3a-0000-1000-8000-00805f9b34fb",
        "write": "0000ae3b-0000-1000-8000-00805f9b34fb",
        "notify": "0000ae3c-0000-1000-8000-00805f9b34fb",
    },
    # Untested: no hardware with this profile has been tried.
    "7772": {
        "service": "0000fff0-0000-1000-8000-00805f9b34fb",
        "write": "0000fff5-0000-1000-8000-00805f9b34fb",
        "notify": "0000fff4-0000-1000-8000-00805f9b34fb",
    },
}

TIMER_MIN, TIMER_MAX = 5, 40  # the unit's advertised timer range, in minutes


@dataclass
class SessionConfig:
    mode_preset: int = 1
    target_intensity: int = 4
    ramp_up_seconds: float = 5.0
    ramp_down_seconds: float = 2.0
    work_seconds: float = 10.0
    rest_seconds: float = 20.0
    total_minutes: int = 20


class FitRxMassager:
    def __init__(self):
        self.client = None
        self.device = None
        self.variant = "standard"
        self.write_uuid = None
        self.notify_uuid = None
        self._t0 = time.monotonic()

    def _notification_handler(self, sender, data: bytearray):
        """Log raw notification bytes with a timestamp. No interpretation."""
        t = time.monotonic() - self._t0
        print(f"[RX {t:8.2f}s] {data.hex(' ').upper()}")

    async def scan_and_connect(self):
        print("Searching for FitRx / HRK-016B...")
        devices = await BleakScanner.discover(timeout=7.0)

        for d in devices:
            name = (d.name or "").strip()
            if any(name.startswith(p) for p in ["HRK-016B", "FitRx Wireles", "FitRx 7772"]):
                self.device = d
                break

        if not self.device:
            raise RuntimeError("FitRx device not found. Make sure the unit is on and blinking.")

        print(f"Found {self.device.name} [{self.device.address}]")
        self.variant = "7772" if "7772" in self.device.name else "standard"
        if self.variant == "7772":
            print("Warning: the 7772 profile is untested.")
        self.write_uuid = UUID_MAP[self.variant]["write"]
        self.notify_uuid = UUID_MAP[self.variant]["notify"]

        self.client = BleakClient(self.device.address)
        await self.client.connect()
        print(f"Connected ({self.variant} profile).")

        await self.client.start_notify(self.notify_uuid, self._notification_handler)
        print("Logging notifications.")

    async def _send_payload(self, hex_string: str):
        if not self.client or not self.client.is_connected:
            raise ConnectionError("BLE client is not connected.")
        await self.client.write_gatt_char(self.write_uuid, bytes.fromhex(hex_string), response=False)
        await asyncio.sleep(0.08)

    async def set_mode(self, mode: int):
        mode = max(1, min(mode, 8))
        print(f"[TX] mode {mode}")
        if self.variant == "standard":
            await self._send_payload(f"ACAD06A1000000{mode:02d}BCBD")
        else:
            await self._send_payload(f"AT+CMODE=0,{mode:02d}\r\n".encode().hex())

    async def set_intensity(self, level: int):
        level = max(0, min(level, 9))
        print(f"[TX] level {level}")
        if self.variant == "standard":
            await self._send_payload(f"ACAD06A0000000{level:02d}BCBD")
        else:
            await self._send_payload(f"41542b435354523d312C{f'{level:X}'.encode().hex()}0D0A")

    async def set_timer(self, minutes: int):
        print(f"[TX] timer {minutes} min" if minutes else "[TX] timer off")
        if self.variant == "standard":
            if minutes == 0:
                await self._send_payload("ACAD06A200000000BCBD")
            else:
                await self._send_payload(f"ACAD06A2010100{minutes:02X}BCBD")
        else:
            if minutes == 0:
                await self._send_payload("41542B4344434C4B3D3030303030300A0D0A")
            else:
                await self._send_payload(f"41542B4344434C4B3D3030{f'{minutes:02d}'.encode().hex()}30300D0A")

    async def stop(self):
        print("[Software stop] Setting level 0 and clearing timer...")
        await self.set_intensity(0)
        await self.set_timer(0)

    async def disconnect(self):
        if self.client and self.client.is_connected:
            try:
                await self.client.stop_notify(self.notify_uuid)
            except Exception:
                pass
            await self.client.disconnect()
            print("Disconnected.")


class SessionOrchestrator:
    def __init__(self, config: SessionConfig):
        self.cfg = config
        self.massager = FitRxMassager()

    async def _ramp(self, start_lvl: int, end_lvl: int, duration_sec: float):
        steps = abs(end_lvl - start_lvl)
        if steps == 0 or duration_sec <= 0:
            await self.massager.set_intensity(end_lvl)
            return
        interval = duration_sec / steps
        direction = 1 if end_lvl > start_lvl else -1
        current = start_lvl
        for _ in range(steps):
            current += direction
            await self.massager.set_intensity(current)
            await asyncio.sleep(interval)

    async def run(self):
        c = self.cfg
        try:
            await self.massager.scan_and_connect()

            # Hardware timer first, so the unit shuts itself off even if this script dies.
            await self.massager.set_timer(c.total_minutes)
            await self.massager.set_mode(c.mode_preset)
            await self.massager.set_intensity(0)

            print("\nSession:")
            print(f"  Duration     : {c.total_minutes} min (hardware timer set)")
            print(f"  Mode preset  : {c.mode_preset}")
            print(f"  Peak level   : {c.target_intensity} of 9")
            print(f"  Work / rest  : {c.work_seconds}s / {c.rest_seconds}s")
            print(f"  Ramp up/down : {c.ramp_up_seconds}s / {c.ramp_down_seconds}s")
            print("  Ctrl+C to stop. Power button or removing a pad is the real emergency stop.\n")

            total_seconds = c.total_minutes * 60
            elapsed = 0.0
            cycle = 1
            while elapsed < total_seconds:
                print(f"[Cycle {cycle}] Ramp up 0 -> {c.target_intensity} over {c.ramp_up_seconds}s")
                await self._ramp(0, c.target_intensity, c.ramp_up_seconds)
                elapsed += c.ramp_up_seconds

                print(f"[Cycle {cycle}] Hold level {c.target_intensity} for {c.work_seconds}s")
                await asyncio.sleep(c.work_seconds)
                elapsed += c.work_seconds

                print(f"[Cycle {cycle}] Ramp down {c.target_intensity} -> 0 over {c.ramp_down_seconds}s")
                await self._ramp(c.target_intensity, 0, c.ramp_down_seconds)
                elapsed += c.ramp_down_seconds

                if elapsed >= total_seconds:
                    break

                print(f"[Cycle {cycle}] Rest at level 0 for {c.rest_seconds}s\n")
                await asyncio.sleep(c.rest_seconds)
                elapsed += c.rest_seconds
                cycle += 1

            print("\nSession complete.")
        except asyncio.CancelledError:
            print("\nInterrupted.")
        except Exception as e:
            print(f"\nError: {e}")
        finally:
            try:
                await self.massager.stop()
            except Exception as e:
                print(f"\n!! SOFTWARE STOP FAILED ({e}).")
                print("!! Turn the unit off with its power button or remove the pads.")
            try:
                await self.massager.disconnect()
            except Exception:
                pass


def parse_args() -> SessionConfig:
    p = argparse.ArgumentParser(description="FitRx BLE session runner with ramps and duty cycles")
    p.add_argument("--mode", type=int, default=1, choices=range(1, 9),
                   help="Hardware mode preset 1-8 (default: 1)")
    p.add_argument("--target-intensity", type=int, default=4, choices=range(1, 10),
                   help="Peak level 1-9 (default: 4)")
    p.add_argument("--duty-cycle", nargs=2, type=float, metavar=("WORK", "REST"), default=[10.0, 20.0],
                   help="Hold and rest times in seconds (default: 10 20)")
    p.add_argument("--ramp", type=float, default=5.0, help="Ramp-up seconds (default: 5)")
    p.add_argument("--ramp-down", type=float, default=2.0, help="Ramp-down seconds (default: 2)")
    p.add_argument("--timer", type=int, default=20, choices=range(TIMER_MIN, TIMER_MAX + 1),
                   metavar=f"{TIMER_MIN}-{TIMER_MAX}",
                   help=f"Session length in minutes, {TIMER_MIN}-{TIMER_MAX} (default: 20)")
    a = p.parse_args()

    work, rest = a.duty_cycle
    if work <= 0 or rest < 0 or a.ramp < 0 or a.ramp_down < 0:
        p.error("work time must be positive; rest and ramp times can't be negative")

    return SessionConfig(
        mode_preset=a.mode,
        target_intensity=a.target_intensity,
        ramp_up_seconds=a.ramp,
        ramp_down_seconds=a.ramp_down,
        work_seconds=work,
        rest_seconds=rest,
        total_minutes=a.timer,
    )


async def main(cfg: SessionConfig):
    # Turn Ctrl+C into a task cancellation so the stop in `finally` always gets
    # sent, on every Python version, instead of the loop being torn down first.
    task = asyncio.current_task()
    asyncio.get_running_loop().add_signal_handler(signal.SIGINT, task.cancel)
    await SessionOrchestrator(cfg).run()


if __name__ == "__main__":
    cfg = parse_args()
    try:
        asyncio.run(main(cfg))
    except KeyboardInterrupt:
        print("\nAborted. If the unit is still running, use its power button.")
        sys.exit(1)