"""
Device manager for BLE connections.
"""

import asyncio
from typing import Any, Callable, Optional

from bleak import BleakClient, BleakScanner
from bleak.backends.characteristic import BleakGATTCharacteristic
from bleak.backends.device import BLEDevice
from bleak.backends.service import BleakGATTService
from bleak.exc import BleakError

from airctl.common.logging import get_logger
from airctl.common.utils import get_timestamp
from airctl.config.manager import ConfigManager
from airctl.models.device import ConnectedDevice
from airctl.models.gatt import CharacteristicInfo, CharacteristicProperty, ServiceInfo

logger = get_logger(__name__)


class DeviceManager:
    """
    Manages BLE device connections.
    
    Handles connection lifecycle, GATT operations, and notifications.
    """

    def __init__(self) -> None:
        self._devices: dict[str, BleakClient] = {}
        self._device_info: dict[str, ConnectedDevice] = {}
        self._notification_callbacks: dict[str, dict[str, Callable]] = {}
        self._config_manager = ConfigManager()

    def _resolve_address(self, address_or_alias: str) -> str:
        """
        Resolve an address or alias to a device address.

        Args:
            address_or_alias: Device address or alias

        Returns:
            Resolved device address
        """
        alias_map = self._config_manager.get_aliases()
        return alias_map.get(address_or_alias, address_or_alias)

    async def connect(
        self,
        address: str,
        timeout: float = 30.0,
        alias: Optional[str] = None,
    ) -> ConnectedDevice:
        """
        Connect to a BLE device.

        Args:
            address: Device address
            timeout: Connection timeout
            alias: Optional alias to assign

        Returns:
            Connected device info

        Raises:
            BleakError: On connection failure
        """
        resolved_address = self._resolve_address(address)

        if resolved_address in self._devices:
            client = self._devices[resolved_address]
            if client.is_connected:
                return self._device_info[resolved_address]

        device: Optional[BLEDevice] = None

        try:
            device = await BleakScanner.find_device_by_address(
                resolved_address, timeout=timeout
            )
        except Exception:
            pass

        if device is None:
            raise ValueError(f"Device {resolved_address} not found")

        client = BleakClient(device, timeout=timeout)

        try:
            await client.connect()
        except BleakError as e:
            logger.error(f"Failed to connect to {resolved_address}: {e}")
            raise

        self._devices[resolved_address] = client

        device_info = ConnectedDevice(
            address=resolved_address,
            name=client.name,
            alias=alias,
            mtu_size=client.mtu_size,
            connected_at=get_timestamp(),
            services_discovered=True,
        )
        self._device_info[resolved_address] = device_info
        self._notification_callbacks[resolved_address] = {}

        if alias:
            self._config_manager.set_alias(resolved_address, alias)

        logger.info(f"Connected to {resolved_address} ({client.name})")
        return device_info

    async def disconnect(self, address: str) -> bool:
        """
        Disconnect from a BLE device.

        Args:
            address: Device address or alias

        Returns:
            True if disconnected, False if not connected

        Raises:
            ValueError: If device not found in connected devices
        """
        resolved_address = self._resolve_address(address)

        if resolved_address not in self._devices:
            raise ValueError(f"Device {address} (resolved to {resolved_address}) not connected")

        client = self._devices[resolved_address]

        for uuid in list(self._notification_callbacks.get(resolved_address, {}).keys()):
            try:
                await client.stop_notify(uuid)
            except Exception:
                pass

        try:
            await client.disconnect()
        except Exception as e:
            logger.warning(f"Error disconnecting from {resolved_address}: {e}")

        del self._devices[resolved_address]
        del self._device_info[resolved_address]
        if resolved_address in self._notification_callbacks:
            del self._notification_callbacks[resolved_address]

        logger.info(f"Disconnected from {resolved_address}")
        return True

    async def disconnect_all(self) -> None:
        """Disconnect from all connected devices."""
        for address in list(self._devices.keys()):
            try:
                await self.disconnect(address)
            except Exception as e:
                logger.warning(f"Error disconnecting from {address}: {e}")

    def list_devices(self) -> list[ConnectedDevice]:
        """List all connected devices, removing stale entries."""
        stale_addresses = []
        for address, client in self._devices.items():
            if not client.is_connected:
                stale_addresses.append(address)
        
        for address in stale_addresses:
            logger.warning(f"Removing stale device entry: {address}")
            del self._devices[address]
            if address in self._device_info:
                del self._device_info[address]
            if address in self._notification_callbacks:
                del self._notification_callbacks[address]
        
        return list(self._device_info.values())

    def get_client(self, address: str) -> BleakClient:
        """
        Get the BleakClient for a device.

        Args:
            address: Device address or alias

        Returns:
            BleakClient instance

        Raises:
            ValueError: If device not connected
        """
        resolved_address = self._resolve_address(address)

        if resolved_address not in self._devices:
            raise ValueError(f"Device {resolved_address} not connected")

        client = self._devices[resolved_address]
        
        if not client.is_connected:
            logger.warning(f"Device {resolved_address} disconnected, cleaning up")
            del self._devices[resolved_address]
            if resolved_address in self._device_info:
                del self._device_info[resolved_address]
            if resolved_address in self._notification_callbacks:
                del self._notification_callbacks[resolved_address]
            raise ValueError(f"Device {resolved_address} disconnected")

        return client

    async def get_services(self, address: str) -> list[ServiceInfo]:
        """
        Get GATT services of a device.

        Args:
            address: Device address or alias

        Returns:
            List of service info
        """
        client = self.get_client(address)

        services: list[ServiceInfo] = []
        for service in client.services:
            services.append(ServiceInfo(
                uuid=service.uuid,
                handle=service.handle,
                description=service.description,
            ))

        if address in self._device_info:
            self._device_info[address].services_discovered = True

        return services

    async def get_characteristics(
        self,
        address: str,
        service_uuid: Optional[str] = None,
    ) -> list[CharacteristicInfo]:
        """
        Get GATT characteristics of a device.

        Args:
            address: Device address or alias
            service_uuid: Optional service UUID filter

        Returns:
            List of characteristic info
        """
        client = self.get_client(address)

        characteristics: list[CharacteristicInfo] = []
        for service in client.services:
            if service_uuid and service.uuid.lower() != service_uuid.lower():
                continue

            for char in service.characteristics:
                properties = [
                    CharacteristicProperty(p)
                    for p in char.properties
                    if p in [e.value for e in CharacteristicProperty]
                ]

                characteristics.append(CharacteristicInfo(
                    uuid=char.uuid,
                    handle=char.handle,
                    service_uuid=service.uuid,
                    service_description=service.description,
                    properties=properties,
                    description=char.description,
                    max_write_without_response_size=char.max_write_without_response_size,
                ))

        return characteristics

    async def read_characteristic(
        self,
        address: str,
        uuid: Optional[str] = None,
        handle: Optional[int] = None,
    ) -> bytes:
        """
        Read a characteristic value.

        Args:
            address: Device address or alias
            uuid: Characteristic UUID
            handle: Characteristic handle (alternative to UUID)

        Returns:
            Characteristic value as bytes
        """
        client = self.get_client(address)
        
        if handle is not None:
            return await client.read_gatt_char(handle)
        elif uuid:
            services = client.services
            readable_char = None
            for service in services:
                for char in service.characteristics:
                    if char.uuid.lower() == uuid.lower():
                        if "read" in char.properties:
                            readable_char = char
                            break
            if readable_char:
                return await client.read_gatt_char(readable_char.handle)
            else:
                raise ValueError(f"Characteristic {uuid} not found or not readable")
        else:
            raise ValueError("Either uuid or handle must be specified")

    async def write_characteristic(
        self,
        address: str,
        data: bytes,
        uuid: Optional[str] = None,
        handle: Optional[int] = None,
        response: bool = False,
    ) -> None:
        """
        Write to a characteristic.

        Args:
            address: Device address or alias
            data: Data to write
            uuid: Characteristic UUID
            handle: Characteristic handle (alternative to UUID)
            response: Whether to request write response
        """
        client = self.get_client(address)
        
        if handle is not None:
            await client.write_gatt_char(handle, data, response=response)
        elif uuid:
            services = client.services
            writable_char = None
            for service in services:
                for char in service.characteristics:
                    if char.uuid.lower() == uuid.lower():
                        if "write" in char.properties:
                            writable_char = char
                            break
            if writable_char:
                await client.write_gatt_char(writable_char.handle, data, response=response)
            else:
                raise ValueError(f"Characteristic {uuid} not found or not writable")
        else:
            raise ValueError("Either uuid or handle must be specified")

    async def subscribe_notification(
        self,
        address: str,
        uuid: str,
        callback: Callable[[str, bytes], Any],
    ) -> None:
        """
        Subscribe to characteristic notifications.

        Args:
            address: Device address or alias
            uuid: Characteristic UUID
            callback: Async callback function(char_uuid, data)
        """
        client = self.get_client(address)
        resolved_address = self._resolve_address(address)

        async def _handler(char: BleakGATTCharacteristic, data: bytearray) -> None:
            await callback(char.uuid, bytes(data))

        await client.start_notify(uuid, _handler)

        if resolved_address not in self._notification_callbacks:
            self._notification_callbacks[resolved_address] = {}
        self._notification_callbacks[resolved_address][uuid] = callback

        logger.debug(f"Subscribed to notifications on {uuid}")

    async def unsubscribe_notification(
        self,
        address: str,
        uuid: str,
    ) -> None:
        """
        Unsubscribe from characteristic notifications.

        Args:
            address: Device address or alias
            uuid: Characteristic UUID
        """
        client = self.get_client(address)
        resolved_address = self._resolve_address(address)

        await client.stop_notify(uuid)

        if resolved_address in self._notification_callbacks:
            self._notification_callbacks[resolved_address].pop(uuid, None)

        logger.debug(f"Unsubscribed from notifications on {uuid}")
