from __future__ import annotations

import asyncio
from contextlib import suppress

from bleak import BleakClient, BleakScanner
from bleak.backends.characteristic import BleakGATTCharacteristic
from bleak.backends.device import BLEDevice

from .transport import (
    AidlabScanner,
    AidlabTransport,
    AidlabTransportFactory,
    DeviceInfo,
)


class BleakTransport(AidlabTransport):
    def __init__(self, device: str | BLEDevice):
        self._device = device
        self._client: BleakClient | None = None
        self._disconnect_queue: asyncio.Queue[None] = asyncio.Queue()
        self._notify_slots: dict[str, _NotifySlot] = {}

    @property
    def mtu_size(self) -> int:
        client = self._client
        if client is None:
            return 23
        return getattr(client, "mtu_size", 23)

    async def connect(self) -> None:
        def _on_disconnect(_client: BleakClient) -> None:
            with suppress(asyncio.QueueFull):
                self._disconnect_queue.put_nowait(None)

        client = BleakClient(self._device, disconnected_callback=_on_disconnect)
        self._client = client
        await client.connect(timeout=10)

    async def disconnect(self) -> None:
        await self._require_client().disconnect()

    async def read_char(self, uuid: str) -> bytes:
        return bytes(await self._require_client().read_gatt_char(uuid))

    async def write_char(self, uuid: str, data: bytes, response: bool = True) -> None:
        await self._require_client().write_gatt_char(uuid, data, response)

    async def notifications(self, uuid: str):
        slot = self._notify_slots.get(uuid)
        if slot is None:
            slot = _NotifySlot(uuid=uuid, transport=self)
            self._notify_slots[uuid] = slot
        return await slot.subscribe()

    async def disconnections(self):
        return _queue_iter(self._disconnect_queue)

    def _require_client(self) -> BleakClient:
        client = self._client
        if client is None:
            raise RuntimeError("BLE client is not connected")
        return client


async def _queue_iter(queue: asyncio.Queue[None] | asyncio.Queue[bytes]):
    while True:
        item = await queue.get()
        yield item


class _NotifySlot:
    def __init__(self, uuid: str, transport: BleakTransport):
        self._uuid = uuid
        self._transport = transport
        self._queues: list[asyncio.Queue[bytes]] = []
        self._active = False
        self._lock = asyncio.Lock()

    async def subscribe(self):
        queue: asyncio.Queue[bytes] = asyncio.Queue()
        self._queues.append(queue)
        await self._ensure_started()

        async def _gen():
            try:
                while True:
                    data = await queue.get()
                    yield data
            finally:
                await self._remove_queue(queue)

        return _gen()

    async def _ensure_started(self) -> None:
        async with self._lock:
            if self._active:
                return
            self._active = True
            client = self._transport._require_client()

            def _on_data(_sender: BleakGATTCharacteristic, data: bytearray) -> None:
                payload = bytes(data)
                for q in list(self._queues):
                    with suppress(asyncio.QueueFull):
                        q.put_nowait(payload)

            await client.start_notify(self._uuid, _on_data)

    async def _remove_queue(self, queue: asyncio.Queue[bytes]) -> None:
        async with self._lock:
            with suppress(ValueError):
                self._queues.remove(queue)
            if self._queues:
                return
            if not self._active:
                return
            self._active = False
            client = self._transport._client
            if client is None:
                return
            with suppress(Exception):
                await client.stop_notify(self._uuid)


class BleakAidlabScanner(AidlabScanner):
    async def scan(self, timeout: float = 10.0) -> list[DeviceInfo]:
        devices = await BleakScanner.discover(timeout=timeout)
        compatible_devices: list[DeviceInfo] = []

        for device in devices:
            if (device.name == "Aidlab 2" or device.name == "Aidlab") and device.address:
                compatible_devices.append(DeviceInfo(address=device.address, name=device.name, metadata=device))

        return compatible_devices


class BleakTransportFactory(AidlabTransportFactory):
    def create(self, device: DeviceInfo) -> AidlabTransport:
        native_device = device.metadata if isinstance(device.metadata, BLEDevice) else device.address
        return BleakTransport(native_device)
