"""
acaia-lunar-ble — Python library for the Acaia Lunar coffee scale over BLE.

Communicates with the Acaia Lunar scale using the proprietary BLE protocol.
Built on top of bleak for cross-platform Bluetooth Low Energy support.
Tested with the Lunar 2021 (Model AL008). Models AL009 and AL010 should
work as well but have not been tested. Does not support other Acaia scales.

Protocol notes:
    - All commands are sent to the HEARTBEAT_CHAR_UUID characteristic (write-without-response).
    - Data arrives via notifications on NOTIFY_CHAR_UUID.
    - Initialization order matters: subscribe → identify → notification request → heartbeat.
    - A heartbeat packet must be sent every ~3 seconds to keep the connection alive.

Message framing (both directions):
    [0xEF] [0xDD] [cmd] [payload...] [cksum1] [cksum2]
    - cksum1 = sum of even-indexed payload bytes & 0xFF
    - cksum2 = sum of odd-indexed payload bytes & 0xFF

Inbound message types (cmd byte):
    cmd 0x0C — event data, sub-type at byte 4:
        0x05 — weight (unstable)
        0x07 — timer
        0x08 — button event (tare, start, stop, reset)
        0x0B — weight (stable) or heartbeat response
    cmd 0x08 — settings (battery, units, auto-off, beep)
"""

from __future__ import annotations

import asyncio
import logging
from collections.abc import Callable
from dataclasses import dataclass

from bleak import BleakClient, BleakScanner

__version__ = "1.0.1"

logger = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# BLE protocol constants
# ---------------------------------------------------------------------------

HEARTBEAT_CHAR_UUID = "49535343-8841-43f4-a8d4-ecbe34729bb3"
"""Write characteristic used for all commands (identify, notification request, heartbeat)."""

NOTIFY_CHAR_UUID = "49535343-1e4d-4bd9-ba61-23c647249616"
"""Notify characteristic that streams weight and other data from the scale."""

HEADER1 = 0xEF
HEADER2 = 0xDD

# Command bytes (outbound)
CMD_HEARTBEAT = 0x00
CMD_TARE = 0x04
CMD_GET_SETTINGS = 0x06
CMD_SETTINGS_RESPONSE = 0x08
CMD_IDENTIFY = 0x0B
CMD_EVENT = 0x0C
CMD_TIMER = 0x0D

# Event sub-types (inbound, at byte 4 when cmd == CMD_EVENT)
EVENT_WEIGHT = 0x05
EVENT_TIMER = 0x07
EVENT_BUTTON = 0x08
EVENT_WEIGHT_STABLE = 0x0B

IDENTIFY = bytes([
    0xEF, 0xDD, 0x0B, 0x30, 0x31, 0x32, 0x33, 0x34, 0x35, 0x36, 0x37, 0x38,
    0x39, 0x30, 0x31, 0x32, 0x33, 0x34, 0x9A, 0x6D,
])
"""Handshake packet that identifies the client to the scale."""

NOTIFICATION_REQUEST = bytes([
    0xEF, 0xDD, 0x0C, 0x09, 0x00, 0x01, 0x01,
    0x02, 0x02, 0x05, 0x03, 0x04, 0x15, 0x06,
])
"""Packet that enables data streaming (weight, battery, timer, button events, settings)."""

HEARTBEAT = bytes([0xEF, 0xDD, 0x00, 0x02, 0x00, 0x02, 0x00])
"""Keep-alive packet. Must be sent at regular intervals (~3 s) to prevent disconnect."""

DEFAULT_HEARTBEAT_INTERVAL = 2.75
DEFAULT_SCAN_RETRY_INTERVAL = 3.0


# ---------------------------------------------------------------------------
# Data classes
# ---------------------------------------------------------------------------

@dataclass
class ScaleSettings:
    """Settings reported by the scale."""
    battery: int          # 0-100 percent
    units: str            # "grams" or "ounces"
    auto_off_minutes: int # minutes until auto power-off
    beep: bool            # whether beep is enabled


# ---------------------------------------------------------------------------
# Packet encoding
# ---------------------------------------------------------------------------

def _encode_packet(cmd: int, payload: list[int]) -> bytes:
    """Build a framed packet: [0xEF] [0xDD] [cmd] [payload...] [cksum1] [cksum2]."""
    buf = bytearray(5 + len(payload))
    buf[0] = HEADER1
    buf[1] = HEADER2
    buf[2] = cmd
    cksum1 = 0
    cksum2 = 0
    for i, val in enumerate(payload):
        val &= 0xFF
        buf[3 + i] = val
        if i % 2 == 0:
            cksum1 += val
        else:
            cksum2 += val
    buf[3 + len(payload)] = cksum1 & 0xFF
    buf[4 + len(payload)] = cksum2 & 0xFF
    return bytes(buf)


# Pre-built command packets
TARE_PACKET = _encode_packet(CMD_TARE, [0])
START_TIMER_PACKET = _encode_packet(CMD_TIMER, [0, 0])
STOP_TIMER_PACKET = _encode_packet(CMD_TIMER, [0, 2])
RESET_TIMER_PACKET = _encode_packet(CMD_TIMER, [0, 1])
GET_SETTINGS_PACKET = _encode_packet(CMD_GET_SETTINGS, [0] * 16)


# ---------------------------------------------------------------------------
# Decoding
# ---------------------------------------------------------------------------

def _decode_weight_payload(payload: bytes | bytearray) -> float | None:
    """Decode a weight payload (6+ bytes) into grams.

    Byte layout (relative to payload start):
        [0]  — weight low byte
        [1]  — weight high byte
        [4]  — decimal scale exponent (weight / 10^value)
        [5]  — flags; bit 1 (0x02) = negative
    """
    if len(payload) < 6:
        return None
    raw = ((payload[1] & 0xFF) << 8) | (payload[0] & 0xFF)
    scale = 10 ** payload[4]
    sign = -1 if (payload[5] & 0x02) else 1
    return sign * (raw / scale)


def decode_weight(data: bytearray) -> float | None:
    """Decode a full 17-byte BLE notification into a weight value in grams.

    This is a convenience wrapper around the internal weight decoder.
    Returns ``None`` if *data* is not a weight message.
    """
    if len(data) < 7 or data[0] != HEADER1 or data[1] != HEADER2:
        return None
    if data[2] != CMD_EVENT:
        return None
    event_type = data[4]
    if event_type == EVENT_WEIGHT:
        return _decode_weight_payload(data[5:])
    if event_type == EVENT_WEIGHT_STABLE:
        return _decode_weight_payload(data[8:])  # 3 extra bytes before weight payload
    return None


def _decode_timer_payload(payload: bytes | bytearray) -> float | None:
    """Decode a timer payload into seconds.

    Byte layout: [0]=minutes, [1]=seconds, [2]=tenths of a second.
    """
    if len(payload) < 3:
        return None
    return payload[0] * 60 + payload[1] + payload[2] / 10.0


def _decode_settings(data: bytearray) -> ScaleSettings | None:
    """Decode a settings response (cmd 0x08) into a ScaleSettings object.

    Byte layout (payload starts at byte 3):
        [0] — unknown
        [1] — battery (& 0x7F = percentage)
        [2] — units (2=grams, 5=ounces)
        [3] — unknown
        [4] — auto-off (* 5 = minutes)
        [5] — unknown
        [6] — beep (1=on)
    """
    if len(data) < 10:
        return None
    payload = data[3:]
    return ScaleSettings(
        battery=payload[1] & 0x7F,
        units="grams" if payload[2] == 2 else "ounces",
        auto_off_minutes=payload[4] * 5,
        beep=payload[6] == 1,
    )


_BUTTON_MAP = {
    (0, 5): "tare",
    (8, 5): "start",
    (10, 7): "stop",
    (9, 7): "reset",
}


def _decode_button_event(payload: bytes | bytearray) -> str | None:
    """Decode a button event sub-type into a string.

    Returns one of: "tare", "start", "stop", "reset", or None.
    """
    if len(payload) < 2:
        return None
    return _BUTTON_MAP.get((payload[0], payload[1]))


# ---------------------------------------------------------------------------
# ScaleConnection
# ---------------------------------------------------------------------------

class ScaleConnection:
    """Async BLE connection to an Acaia Lunar coffee scale.

    Parameters
    ----------
    address:
        The BLE address (MAC or UUID) of the scale.
    weight_callback:
        Called with each weight reading in grams (``float``).
    timer_callback:
        Called with timer value in seconds (``float``).
    battery_callback:
        Called with battery percentage (``int``, 0-100).
    settings_callback:
        Called with a ``ScaleSettings`` object when settings are received.
    button_callback:
        Called with button name (``str``): "tare", "start", "stop", or "reset".
    notification_callback:
        Called with every raw BLE notification payload (``bytearray``).
    """

    def __init__(
        self,
        address: str,
        *,
        weight_callback: Callable[[float], None] | None = None,
        timer_callback: Callable[[float], None] | None = None,
        battery_callback: Callable[[int], None] | None = None,
        settings_callback: Callable[[ScaleSettings], None] | None = None,
        button_callback: Callable[[str], None] | None = None,
        notification_callback: Callable[[bytearray], None] | None = None,
    ):
        self.address = address
        self.client: BleakClient | None = None
        self._services = None
        self._heartbeat_task: asyncio.Task | None = None
        self._weight_callback = weight_callback
        self._timer_callback = timer_callback
        self._battery_callback = battery_callback
        self._settings_callback = settings_callback
        self._button_callback = button_callback
        self._notification_callback = notification_callback

        # Latest values from the stream
        self.weight: float | None = None
        self.timer: float | None = None
        self.settings: ScaleSettings | None = None

    # -- context manager -----------------------------------------------------

    async def __aenter__(self) -> ScaleConnection:
        await self.connect()
        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb) -> None:
        await self.disconnect()

    # -- properties ----------------------------------------------------------

    @property
    def is_connected(self) -> bool:
        return self.client is not None and self.client.is_connected

    # -- discovery -----------------------------------------------------------

    @staticmethod
    async def discover(timeout: float = 5.0) -> tuple[list, list]:
        """Scan for nearby Acaia scales.

        Returns a tuple ``(acaia_devices, all_devices)`` where *acaia_devices*
        contains only devices whose name contains ``ACAIA`` or ``LUNAR``
        (case-insensitive), and *all_devices* is the full list of discovered
        BLE devices.
        """
        all_devices = await BleakScanner.discover(timeout=timeout)
        acaia_devices = [
            d for d in all_devices
            if d.name and any(k in d.name.upper() for k in ("ACAIA", "LUNAR"))
        ]
        for d in acaia_devices:
            logger.info("Discovered Acaia device: %s (%s)", d.name, d.address)
        return acaia_devices, all_devices

    # -- scan / connect / disconnect -----------------------------------------

    async def scan_until_found(self) -> None:
        """Block until the scale with the configured address is found via BLE scan."""
        while True:
            logger.info("Scanning for BLE devices...")
            dev = await BleakScanner.find_device_by_address(self.address, timeout=1.0)
            if dev:
                logger.info("Found scale: %s at %s", dev.name, dev.address)
                return
            logger.info("Not found. Retrying in %.1fs...", DEFAULT_SCAN_RETRY_INTERVAL)
            await asyncio.sleep(DEFAULT_SCAN_RETRY_INTERVAL)

    async def connect(self) -> bool:
        """Connect to the scale and start data streaming.

        The initialization order is critical — the scale will not send data if
        the steps are performed out of order:

        1. Subscribe to notifications
        2. Send IDENTIFY packet
        3. Send NOTIFICATION_REQUEST
        4. Start heartbeat loop
        """
        logger.info("Connecting to %s...", self.address)
        self.client = BleakClient(self.address)
        try:
            await self.client.connect()
        except Exception as e:
            logger.error("Connection failed: %s", e)
            self.client = None
            return False

        if not self.client.is_connected:
            logger.error("Connection failed.")
            self.client = None
            return False

        logger.info("Connected! %s", self.client.address)

        # The order and timing of these steps is important!
        # The scale needs time to process each command before the next.
        await self.start_notifications()
        await asyncio.sleep(0.2)
        await self.send_identify()
        await asyncio.sleep(0.3)
        await self.send_notification_request()
        await asyncio.sleep(0.3)
        await self.start_heartbeat(interval=DEFAULT_HEARTBEAT_INTERVAL)

        return True

    async def disconnect(self) -> None:
        """Stop streaming and disconnect from the scale."""
        if self.client:
            try:
                await self.client.stop_notify(NOTIFY_CHAR_UUID)
            except Exception as e:
                logger.debug("Error stopping notifications: %s", e)
            await self.stop_heartbeat()
            try:
                await self.client.disconnect()
            except Exception as e:
                logger.debug("Error disconnecting: %s", e)
            logger.info("Disconnected")
            self.client = None
            self._services = None

    async def ensure_connected(self) -> bool:
        """Reconnect if the connection was lost. Returns ``True`` if connected."""
        if self.is_connected:
            return True
        logger.info("Client not connected, reconnecting...")
        return await self.connect()

    # -- service inspection (debug) ------------------------------------------

    async def get_services(self):
        """Fetch and cache GATT services from the device."""
        if not await self.ensure_connected():
            logger.warning("Cannot get services: not connected")
            return None
        if self._services is None:
            self._services = self.client.services
        return self._services

    async def print_services(self) -> None:
        """Print all GATT services, characteristics, and descriptors."""
        services = await self.get_services()
        if services is None:
            return
        for service in services:
            print(f"[Service] {service.uuid}: {service.description}")
            for char in service.characteristics:
                props = ",".join(char.properties)
                print(f"  [Characteristic] {char.uuid}: ({props})")
                for desc in char.descriptors:
                    print(f"    [Descriptor] {desc.uuid}")

    # -- heartbeat -----------------------------------------------------------

    async def _heartbeat_loop(self, interval: float) -> None:
        """Internal loop that sends heartbeat packets while connected."""
        logger.debug("Heartbeat loop started (interval=%.2fs)", interval)
        try:
            while self.is_connected:
                try:
                    await self.client.write_gatt_char(
                        HEARTBEAT_CHAR_UUID, HEARTBEAT, response=False,
                    )
                except Exception as e:
                    logger.warning("Heartbeat error: %s", e)
                    break
                await asyncio.sleep(interval)
        finally:
            logger.debug("Heartbeat loop stopped")
            self._heartbeat_task = None

    async def start_heartbeat(self, interval: float = DEFAULT_HEARTBEAT_INTERVAL) -> None:
        """Start the heartbeat background task if not already running."""
        if self._heartbeat_task is not None:
            logger.debug("Heartbeat already running")
            return
        if not await self.ensure_connected():
            logger.warning("Cannot start heartbeat: not connected")
            return
        self._heartbeat_task = asyncio.get_running_loop().create_task(
            self._heartbeat_loop(interval)
        )

    async def stop_heartbeat(self) -> None:
        """Cancel the heartbeat background task."""
        if self._heartbeat_task is not None:
            logger.debug("Cancelling heartbeat loop...")
            self._heartbeat_task.cancel()
            try:
                await self._heartbeat_task
            except asyncio.CancelledError:
                pass
            self._heartbeat_task = None

    # -- scale commands ------------------------------------------------------

    async def _send_command(self, packet: bytes, name: str) -> bool:
        """Send a pre-built command packet to the scale."""
        if not await self.ensure_connected():
            logger.warning("Cannot send %s: not connected", name)
            return False
        try:
            await self.client.write_gatt_char(
                HEARTBEAT_CHAR_UUID, packet, response=False,
            )
            logger.info("%s sent", name)
            return True
        except Exception as e:
            logger.error("Error sending %s: %s", name, e)
            return False

    async def send_identify(self) -> bool:
        """Send the IDENTIFY handshake packet to the scale."""
        return await self._send_command(IDENTIFY, "Identify")

    async def send_notification_request(self) -> bool:
        """Send the NOTIFICATION_REQUEST to enable data streaming."""
        return await self._send_command(NOTIFICATION_REQUEST, "Notification request")

    async def tare(self) -> bool:
        """Zero the scale."""
        return await self._send_command(TARE_PACKET, "Tare")

    async def start_timer(self) -> bool:
        """Start the scale's built-in timer."""
        return await self._send_command(START_TIMER_PACKET, "Start timer")

    async def stop_timer(self) -> bool:
        """Stop/pause the scale's built-in timer."""
        return await self._send_command(STOP_TIMER_PACKET, "Stop timer")

    async def reset_timer(self) -> bool:
        """Reset the scale's built-in timer to zero."""
        return await self._send_command(RESET_TIMER_PACKET, "Reset timer")

    async def get_settings(self) -> bool:
        """Request current settings from the scale (battery, units, auto-off, beep).

        The response arrives asynchronously via the ``settings_callback``.
        """
        return await self._send_command(GET_SETTINGS_PACKET, "Get settings")

    # -- notifications -------------------------------------------------------

    def _handle_weight(self, payload: bytes | bytearray) -> None:
        """Decode weight payload, store it, and invoke callback."""
        weight = _decode_weight_payload(payload)
        if weight is not None:
            self.weight = weight
            if self._weight_callback:
                self._weight_callback(weight)

    def _handle_timer(self, payload: bytes | bytearray) -> None:
        """Decode timer payload, store it, and invoke callback."""
        timer = _decode_timer_payload(payload)
        if timer is not None:
            self.timer = timer
            if self._timer_callback:
                self._timer_callback(timer)

    def _notify_handler(self, sender, data: bytearray) -> None:
        """BLE notification callback — parses messages and invokes user callbacks."""
        if self._notification_callback:
            self._notification_callback(data)

        if len(data) < 5 or data[0] != HEADER1 or data[1] != HEADER2:
            return

        cmd = data[2]

        # Settings response (battery, units, auto-off, beep)
        if cmd == CMD_SETTINGS_RESPONSE:
            settings = _decode_settings(data)
            if settings is not None:
                self.settings = settings
                logger.debug("Settings: %s", settings)
                if self._battery_callback:
                    self._battery_callback(settings.battery)
                if self._settings_callback:
                    self._settings_callback(settings)
            return

        # Event data (weight, timer, button, heartbeat response)
        if cmd != CMD_EVENT or len(data) < 6:
            return

        event_type = data[4]

        if event_type == EVENT_WEIGHT:
            self._handle_weight(data[5:])
            # Timer data is embedded after the weight payload at bytes [11:15]
            if len(data) >= 15 and data[11] == EVENT_TIMER:
                self._handle_timer(data[12:])

        elif event_type == EVENT_WEIGHT_STABLE:
            # Heartbeat response — may contain weight (sub-type 5) or timer (sub-type 7)
            if len(data) >= 8:
                sub_type = data[7]
                if sub_type == EVENT_WEIGHT:
                    self._handle_weight(data[8:])
                elif sub_type == EVENT_TIMER:
                    self._handle_timer(data[8:])

        elif event_type == EVENT_TIMER:
            self._handle_timer(data[5:])

        elif event_type == EVENT_BUTTON:
            button = _decode_button_event(data[5:])
            if button is not None:
                logger.debug("Button event: %s", button)
                if self._button_callback:
                    self._button_callback(button)

    async def start_notifications(self) -> bool:
        """Subscribe to the scale's notify characteristic."""
        if not await self.ensure_connected():
            logger.warning("Cannot subscribe: not connected")
            return False
        try:
            await self.client.start_notify(NOTIFY_CHAR_UUID, self._notify_handler)
            logger.info("Subscribed to %s", NOTIFY_CHAR_UUID)
            return True
        except Exception as e:
            logger.error("Error subscribing to notifications: %s", e)
            return False


__all__ = [
    "ScaleConnection",
    "ScaleSettings",
    "decode_weight",
    "HEARTBEAT_CHAR_UUID",
    "NOTIFY_CHAR_UUID",
    "IDENTIFY",
    "NOTIFICATION_REQUEST",
    "HEARTBEAT",
    "__version__",
]
