# onecontrol_client.py
# Lippert OneControl BLE Client
# Based on reverse engineered protocol from decompiled Xamarin app

import asyncio
import logging
import struct
from collections.abc import Callable
from bleak import BleakClient
from bleak_retry_connector import establish_connection, BleakClientWithServiceCache
from .cobs_protocol import CobsEncoder, CobsDecoder
from enum import IntEnum, StrEnum


class EventType(StrEnum):
    SWITCH    = "switch"
    TANK      = "tank"
    BATTERY   = "battery"
    COVER_RAW = "cover_raw"

_LOGGER = logging.getLogger(__name__)

AUTH_SERVICE   = "00000010-0200-a58e-e411-afe28044e62c"
SEED_CHAR      = "00000012-0200-a58e-e411-afe28044e62c"
KEY_CHAR       = "00000013-0200-a58e-e411-afe28044e62c"
RV_LINK_CYPHER = 612643285


def _tea_encrypt(cypher: int, seed: int) -> int:
    M = 0xFFFFFFFF
    DELTA = 2654435769
    num = DELTA
    for _ in range(32):
        t = (((cypher << 4) & M) + 1131376761) & M
        t ^= (cypher + num) & M
        t ^= ((cypher >> 5) + 1919510376) & M
        seed = (seed + t) & M
        t = (((seed << 4) & M) + 1948272964) & M
        t ^= (seed + num) & M
        t ^= ((seed >> 5) + 1400073827) & M
        cypher = (cypher + t) & M
        num = (num + DELTA) & M
    return seed


async def _perform_auth(client: BleakClient) -> bool:
    for attempt in range(3):
        if not client.is_connected:
            _LOGGER.warning("Auth: connection lost before attempt %d", attempt + 1)
            return False

        try:
            seed_bytes = await client.read_gatt_char(SEED_CHAR)
        except Exception as e:
            _LOGGER.warning("Auth: failed to read seed (attempt %d): %s", attempt + 1, e)
            if not client.is_connected:
                return False
            await asyncio.sleep(2)
            continue

        if seed_bytes.lower() == b"unlocked":
            _LOGGER.debug("Auth: already unlocked")
            return True

        if len(seed_bytes) < 4:
            _LOGGER.warning("Auth: seed too short (%d bytes), retrying", len(seed_bytes))
            await asyncio.sleep(2)
            continue

        seed = struct.unpack_from(">I", seed_bytes, 0)[0]
        if seed == 0:
            _LOGGER.warning("Auth: seed is zero, retrying")
            await asyncio.sleep(2)
            continue

        key = _tea_encrypt(RV_LINK_CYPHER, seed)
        _LOGGER.debug("Auth: seed=0x%08x → key=0x%08x", seed, key)

        try:
            await client.write_gatt_char(KEY_CHAR, struct.pack(">I", key), response=False)
            await asyncio.sleep(0.5)
            verify = await client.read_gatt_char(SEED_CHAR)
            if verify.lower() == b"unlocked":
                _LOGGER.debug("Auth: confirmed unlocked")
                return True
            _LOGGER.warning("Auth: unexpected verify response: %r", verify)
            return True  # key was accepted even if verify string differs
        except Exception as e:
            _LOGGER.error("Auth: failed to write key (attempt %d): %s", attempt + 1, e)
            return False

    _LOGGER.error("Auth: all attempts exhausted")
    return False


class CommandType(IntEnum):
    GET_DEVICES    = 1
    ACTION_SWITCH  = 64
    ACTION_MOVEMENT = 65
    ACTION_DIMMABLE = 67
    ACTION_RGB     = 68
    ACTION_HVAC    = 69


class SwitchState(IntEnum):
    OFF    = 0
    ON     = 1
    TOGGLE = 2


class MovementState(IntEnum):
    RETRACT = 0
    EXTEND  = 1
    STOP    = 2


# State callback type: (event_type: str, data: dict) -> None
# event_type values: "switch", "tank", "battery", "cover"
StateCallback = Callable[[str, dict], None]


class OneControlClient:
    SERVICE_UUID = "00000030-0200-A58E-E411-AFE28044E62C"
    WRITE_CHAR   = "00000033-0200-A58E-E411-AFE28044E62C"
    READ_CHAR    = "00000034-0200-A58E-E411-AFE28044E62C"

    _QUIET_EVENTS = {1, 3, 4, 26, 32}

    def __init__(self, address: str, state_callback: StateCallback | None = None,
                 disconnect_callback: Callable[[], None] | None = None,
                 ble_device=None):
        self.address = address
        self._ble_device = ble_device
        self._state_cb = state_callback
        self._disconnect_cb = disconnect_callback
        self.client: BleakClient | None = None
        self.encoder = CobsEncoder()
        self.decoder = CobsDecoder()
        self._seq = 0
        self._pending_commands: dict[int, CommandType] = {}

    def set_ble_device(self, ble_device) -> None:
        """Update the BLE device reference (from HA scanner)."""
        self._ble_device = ble_device

    async def connect(self) -> None:
        # Panel often drops the first BLE connection immediately.
        # Retry the full connect→auth sequence with a fresh client each time.
        # Do NOT call pair() — panel uses proprietary TEA auth, not BLE pairing,
        # and the pair request may trigger the disconnect.
        last_err: Exception | None = None
        for attempt in range(5):
            try:
                if self._ble_device:
                    self.client = await establish_connection(
                        BleakClientWithServiceCache,
                        self._ble_device,
                        self.address,
                        disconnected_callback=self._on_disconnect,
                        max_attempts=2,
                    )
                else:
                    self.client = BleakClient(self.address, disconnected_callback=self._on_disconnect)
                    await self.client.connect(timeout=30.0)

                _LOGGER.info("BLE connected (attempt %d), starting auth...", attempt + 1)

                # Brief settle time — just enough for BlueZ service discovery
                await asyncio.sleep(0.5)

                if not self.client.is_connected:
                    _LOGGER.info("Panel dropped connection on attempt %d — retrying", attempt + 1)
                    await asyncio.sleep(2.0)
                    continue

                if not await _perform_auth(self.client):
                    if self.client.is_connected:
                        # Auth truly failed (not just disconnect) — still retry,
                        # the panel may need a warmup connection first
                        await self.client.disconnect()
                    _LOGGER.warning("Auth failed on attempt %d — retrying full sequence", attempt + 1)
                    await asyncio.sleep(2.0)
                    continue

                await asyncio.sleep(0.5)
                await self.client.start_notify(self.READ_CHAR, self._notification_handler)
                _LOGGER.info("OneControl connected and authenticated (attempt %d)", attempt + 1)
                return

            except Exception as e:
                last_err = e
                _LOGGER.warning("Connect attempt %d failed: %s", attempt + 1, e)
                if self.client and self.client.is_connected:
                    try:
                        await self.client.disconnect()
                    except Exception:
                        pass
                await asyncio.sleep(2.0)

        raise RuntimeError(f"OneControl: failed to connect after 4 attempts: {last_err}")

    def _on_disconnect(self, _client: BleakClient) -> None:
        _LOGGER.warning("OneControl: disconnected unexpectedly")
        if self._disconnect_cb:
            self._disconnect_cb()

    async def disconnect(self) -> None:
        if self.client and self.client.is_connected:
            await self.client.disconnect()

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

    def _notification_handler(self, sender, data: bytearray) -> None:
        try:
            decoded = self.decoder.decode(data)
            if len(decoded) < 4:
                return

            event_type = decoded[0]
            seq        = struct.unpack(">H", decoded[1:3])[0]
            resp_type  = decoded[3]
            payload    = decoded[4:]
            completed  = resp_type >= 128

            if event_type == 2:
                cmd_type = self._pending_commands.get(seq)
                if cmd_type == CommandType.GET_DEVICES and completed:
                    self._pending_commands.pop(seq, None)
                elif cmd_type is not None and completed:
                    self._pending_commands.pop(seq, None)
            else:
                self._handle_status_event(event_type, decoded)

        except Exception as e:
            _LOGGER.debug("OneControl: error decoding notification: %s", e)

    def _handle_status_event(self, event_type: int, decoded: bytes) -> None:
        if event_type == 7 and len(decoded) >= 6:
            v_raw    = struct.unpack_from(">H", decoded, 1)[0]
            features = decoded[5]
            if features & 1 and v_raw != 0xFFFF:
                voltage = v_raw / 256.0
                self._emit(EventType.BATTERY, {"voltage": voltage})

        elif event_type == 12:
            i = 2
            while i + 1 < len(decoded):
                dev_id = decoded[i]
                pct    = decoded[i + 1]
                self._emit(EventType.TANK, {"device_id": dev_id, "pct": pct})
                i += 2

        elif event_type == 6:
            i = 2
            while i + 1 < len(decoded):
                dev_id     = decoded[i]
                state_byte = decoded[i + 1]
                on         = bool(state_byte & 0x01)
                self._emit(EventType.SWITCH, {"device_id": dev_id, "on": on})
                i += 7

        elif event_type == 14:
            # RelayHBridgeMomentaryType2 — covers/slides/awnings
            # payload structure not fully decoded; pass raw for now
            self._emit(EventType.COVER_RAW, {"data": decoded})

        elif event_type not in self._QUIET_EVENTS:
            _LOGGER.debug("OneControl: unhandled event %d: %s", event_type, decoded.hex())

    def _emit(self, event_type: str, data: dict) -> None:
        if self._state_cb:
            self._state_cb(event_type, data)

    def _next_seq(self) -> int:
        self._seq = (self._seq + 1) & 0xFFFF
        return self._seq

    async def send_command(self, cmd_type: CommandType, table_id: int, payload: bytes) -> None:
        if not self.is_connected:
            raise RuntimeError("Not connected")
        seq = self._next_seq()
        self._pending_commands[seq] = cmd_type
        packet  = struct.pack(">HBB", seq, cmd_type, table_id) + payload
        encoded = self.encoder.encode(packet)
        # Panel's command char is a streaming (Write Without Response) endpoint,
        # like the OEM app. Bleak 3.x defaults to write-with-response when the
        # char advertises the "write" property, which the panel rejects with
        # ATT 0x0E (Unlikely Error). Force write-without-response.
        await self.client.write_gatt_char(self.WRITE_CHAR, encoded, response=False)

    async def get_devices(self) -> None:
        await self.send_command(CommandType.GET_DEVICES, 1, bytes([0, 255]))

    async def set_switch(self, device_id: int, on: bool) -> None:
        state = SwitchState.ON if on else SwitchState.OFF
        await self.send_command(CommandType.ACTION_SWITCH, 1, bytes([state, device_id]))

    async def control_movement(self, device_id: int, state: MovementState) -> None:
        await self.send_command(CommandType.ACTION_MOVEMENT, 1, bytes([state, device_id]))

    async def set_dimmer(self, device_id: int, level: int) -> None:
        await self.send_command(CommandType.ACTION_DIMMABLE, 1, bytes([level, device_id]))