adguard-control-hub/custom_components/adguard_hub/__init__.py

"""
🛡️ AdGuard Control Hub for Home Assistant.

Transform your AdGuard Home into a smart network management powerhouse with 
complete client control, service blocking, and automation capabilities.
"""
import asyncio
import logging
from datetime import timedelta
from homeassistant.config_entries import ConfigEntry
from homeassistant.const import CONF_HOST, CONF_PASSWORD, CONF_PORT, CONF_USERNAME
from homeassistant.core import HomeAssistant
from homeassistant.exceptions import ConfigEntryNotReady
from homeassistant.helpers.aiohttp_client import async_get_clientsession
from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, UpdateFailed
from .const import DOMAIN, PLATFORMS, SCAN_INTERVAL, CONF_SSL, CONF_VERIFY_SSL
from .api import AdGuardHomeAPI

_LOGGER = logging.getLogger(__name__)

async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
    """Set up AdGuard Control Hub from a config entry."""
    session = async_get_clientsession(hass, entry.data.get(CONF_VERIFY_SSL, True))

    api = AdGuardHomeAPI(
        host=entry.data[CONF_HOST],
        port=entry.data[CONF_PORT],
        username=entry.data.get(CONF_USERNAME),
        password=entry.data.get(CONF_PASSWORD),
        ssl=entry.data.get(CONF_SSL, False),
        session=session,
    )

    # Test the connection
    try:
        await api.test_connection()
        _LOGGER.info("Successfully connected to AdGuard Home at %s:%s", 
                    entry.data[CONF_HOST], entry.data[CONF_PORT])
    except Exception as err:
        _LOGGER.error("Failed to connect to AdGuard Home: %s", err)
        raise ConfigEntryNotReady(f"Unable to connect: {err}")

    # Create update coordinator
    coordinator = AdGuardControlHubCoordinator(hass, api)
    await coordinator.async_config_entry_first_refresh()

    # Store data
    hass.data.setdefault(DOMAIN, {})
    hass.data[DOMAIN][entry.entry_id] = {
        "coordinator": coordinator,
        "api": api,
    }

    # Set up platforms
    await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS)

    _LOGGER.info("AdGuard Control Hub setup complete")
    return True

async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
    """Unload AdGuard Control Hub config entry."""
    unload_ok = await hass.config_entries.async_unload_platforms(entry, PLATFORMS)

    if unload_ok:
        hass.data[DOMAIN].pop(entry.entry_id)

    return unload_ok

class AdGuardControlHubCoordinator(DataUpdateCoordinator):
    """AdGuard Control Hub data update coordinator."""

    def __init__(self, hass: HomeAssistant, api: AdGuardHomeAPI):
        """Initialize the coordinator."""
        super().__init__(
            hass,
            _LOGGER,
            name=f"{DOMAIN}_coordinator",
            update_interval=timedelta(seconds=SCAN_INTERVAL),
        )
        self.api = api
        self._clients = {}
        self._statistics = {}
        self._protection_status = {}

    async def _async_update_data(self):
        """Fetch data from AdGuard Home."""
        try:
            # Fetch all data concurrently for better performance
            results = await asyncio.gather(
                self.api.get_clients(),
                self.api.get_statistics(), 
                self.api.get_status(),
                return_exceptions=True,
            )

            clients, statistics, status = results

            # Handle any exceptions
            for i, result in enumerate(results):
                if isinstance(result, Exception):
                    endpoint_names = ["clients", "statistics", "status"]
                    _LOGGER.warning("Error fetching %s: %s", endpoint_names[i], result)

            # Update stored data (use empty dict if fetch failed)
            self._clients = {
                client["name"]: client 
                for client in (clients.get("clients", []) if not isinstance(clients, Exception) else [])
            }
            self._statistics = statistics if not isinstance(statistics, Exception) else {}
            self._protection_status = status if not isinstance(status, Exception) else {}

            return {
                "clients": self._clients,
                "statistics": self._statistics,
                "status": self._protection_status,
            }

        except Exception as err:
            raise UpdateFailed(f"Error communicating with AdGuard Control Hub: {err}")

    @property
    def clients(self):
        """Return clients data."""
        return self._clients

    @property 
    def statistics(self):
        """Return statistics data."""
        return self._statistics

    @property
    def protection_status(self):
        """Return protection status data."""
        return self._protection_status