"""Discovery module for Lunatone interfaces."""

import asyncio
import dataclasses
import json
import socket
from enum import IntFlag, StrEnum, auto
from types import TracebackType
from typing import Any, AsyncIterator, Final

DISCOVERY_ADDRESS: Final[str] = "255.255.255.255"
DISCOVERY_PORT: Final[int] = 5555
DISCOVERY_MESSAGE: Final[str] = "discovery"
DISCOVERY_LISTEN_TIMEOUT_SECONDS: Final[float] = 3.0


class LunatoneDiscoveryFeature(IntFlag):
    """Features."""

    REST_API = auto()


class LunatoneDiscoveryType(StrEnum):
    """Discovery types."""

    DALI2_IOT = "dali-2-iot"
    DALI2_IOT4 = "dali-2-iot4"
    DALI2_DISPLAY = "dali-2-display"
    DALI_WIFI = "dali-wifi"
    AIR_WLAN_SWITCH = "air-wlan-switch"

    @property
    def supported_features(self) -> set[LunatoneDiscoveryFeature]:
        return _SUPPORTED_FEATURES[self]


_SUPPORTED_FEATURES: dict[
    LunatoneDiscoveryType, set[LunatoneDiscoveryFeature]
] = {
    LunatoneDiscoveryType.DALI2_IOT: (LunatoneDiscoveryFeature.REST_API),
    LunatoneDiscoveryType.DALI2_IOT4: (LunatoneDiscoveryFeature.REST_API),
    LunatoneDiscoveryType.DALI2_DISPLAY: (LunatoneDiscoveryFeature.REST_API),
    LunatoneDiscoveryType.DALI_WIFI: (),
    LunatoneDiscoveryType.AIR_WLAN_SWITCH: (),
}


@dataclasses.dataclass(frozen=True)
class LunatoneDiscoveryFilters:
    features: list[LunatoneDiscoveryFeature] = dataclasses.field(default_factory=list)
    types: list[LunatoneDiscoveryType] = dataclasses.field(default_factory=list)

    def match_features(self, type_: LunatoneDiscoveryType) -> bool:
        return all(f in type_.supported_features for f in self.features)

    def match_types(self, type_: LunatoneDiscoveryType) -> bool:
        return all(t == type_ for t in self.types)

    def match_all(self, type_: LunatoneDiscoveryType) -> bool:
        if self.match_features(type_) and self.match_types(type_):
            return True
        return False


@dataclasses.dataclass(frozen=True)
class LunatoneDiscoveryInfo:
    """LunatoneDiscoveryInfo."""

    host: str
    name: str
    type_: LunatoneDiscoveryType


class LunatoneDiscoveryUDPProtocol(asyncio.DatagramProtocol):
    """UDP protocol implementation for discovering Lunatone interfaces."""

    def __init__(self, queue: asyncio.Queue) -> None:
        """Initialize the UDP protocol instance."""
        super().__init__()
        self._queue: asyncio.Queue[LunatoneDiscoveryInfo] = queue

    def datagram_received(self, data: bytes, addr: tuple[str | Any, int]) -> None:
        """Handle an incoming UDP datagram."""
        try:
            message = json.loads(data.decode())
        except json.JSONDecodeError:
            return
        if not ("type" in message and "name" in message):  # Validate message
            return
        if message["type"] in LunatoneDiscoveryType:
            info = LunatoneDiscoveryInfo(
                addr[0], message["name"], LunatoneDiscoveryType(message["type"])
            )
            self._queue.put_nowait(info)


class LunatoneDiscovery:
    """Discover Lunatone devices."""

    def __init__(self, loop: asyncio.AbstractEventLoop | None = None) -> None:
        self.loop: asyncio.AbstractEventLoop = loop or asyncio.get_running_loop()
        self._queue: asyncio.Queue[LunatoneDiscoveryInfo] = asyncio.Queue()
        self._transport: asyncio.DatagramTransport | None = None

    async def __aenter__(self) -> "LunatoneDiscovery":
        return self

    async def __aexit__(
        self,
        exc_type: type[BaseException] | None,
        exc_val: BaseException | None,
        exc_tb: TracebackType | None,
    ) -> None:
        self.close()

    def close(self) -> None:
        if self._transport:
            self._transport.close()

    async def async_discover_devices_stream(
        self,
        timeout: float | None = None,
        local_ip: str | None = None,
        filters: LunatoneDiscoveryFilters | None = None
    ) -> AsyncIterator[LunatoneDiscoveryInfo]:
        """Yield devices one by one as they arrive (streaming style)."""
        if timeout is None:
            timeout = DISCOVERY_LISTEN_TIMEOUT_SECONDS
        if self._transport is None:
            self._transport, _ = await self.loop.create_datagram_endpoint(
                lambda: LunatoneDiscoveryUDPProtocol(self._queue),
                local_addr=(local_ip or "0.0.0.0", DISCOVERY_PORT),
                family=socket.AF_INET,
                allow_broadcast=True,
            )
        self._transport.sendto(
            DISCOVERY_MESSAGE.encode(), (DISCOVERY_ADDRESS, DISCOVERY_PORT)
        )
        end = self.loop.time() + timeout
        while self.loop.time() < end:
            try:
                device = await asyncio.wait_for(
                    self._queue.get(), timeout=end - self.loop.time()
                )
                if filters and filters.match_all(device.type):
                    yield device
            except KeyError:
                continue
            except asyncio.TimeoutError:
                break

    async def async_discover_devices(
        self,
        timeout: float | None = None,
        local_ip: str | None = None,
        filters: LunatoneDiscoveryFilters | None = None
    ) -> list[LunatoneDiscoveryInfo]:
        """Collect all devices and return them at once (batch style)."""
        return [d async for d in self.async_discover_devices_stream(timeout, local_ip, filters)]
