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.
"""
import asyncio
import logging
from datetime import timedelta
from typing import Any, Dict

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 .api import AdGuardHomeAPI, AdGuardConnectionError, AdGuardHomeError
from .const import CONF_SSL, CONF_VERIFY_SSL, DOMAIN, PLATFORMS, SCAN_INTERVAL
from .services import AdGuardControlHubServices

_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))

    # Create API instance
    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),
        verify_ssl=entry.data.get(CONF_VERIFY_SSL, True),
        session=session,
    )

    # Test the connection
    try:
        if not await api.test_connection():
            raise ConfigEntryNotReady("Unable to connect to AdGuard Home")

        _LOGGER.info(
            "Successfully connected to AdGuard Home at %s:%s",
            entry.data[CONF_HOST],
            entry.data[CONF_PORT]
        )
    except AdGuardHomeError as err:
        _LOGGER.error("Failed to connect to AdGuard Home: %s", err)
        raise ConfigEntryNotReady(f"Unable to connect: {err}") from err
    except Exception as err:
        _LOGGER.exception("Unexpected error connecting to AdGuard Home")
        raise ConfigEntryNotReady(f"Unexpected error: {err}") from err

    # Create update coordinator
    coordinator = AdGuardControlHubCoordinator(hass, api)

    try:
        await coordinator.async_config_entry_first_refresh()
    except UpdateFailed as err:
        _LOGGER.error("Failed to perform initial data refresh: %s", err)
        raise ConfigEntryNotReady(f"Failed to fetch initial data: {err}") from err

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

    # Set up platforms
    try:
        await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS)
    except Exception as err:
        _LOGGER.error("Failed to set up platforms: %s", err)
        # Clean up on failure
        hass.data[DOMAIN].pop(entry.entry_id, None)
        raise ConfigEntryNotReady(f"Failed to set up platforms: {err}") from err

    # Register services (only once)
    services_key = f"{DOMAIN}_services"
    if services_key not in hass.data:
        services = AdGuardControlHubServices(hass)
        services.register_services()
        hass.data[services_key] = services

    _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:
        # Remove this entry's data
        hass.data[DOMAIN].pop(entry.entry_id, None)

        # Unregister services if this was the last entry
        if not hass.data[DOMAIN]:
            services_key = f"{DOMAIN}_services"
            services = hass.data.get(services_key)
            if services:
                services.unregister_services()
                hass.data.pop(services_key, None)
            hass.data.pop(DOMAIN, None)

    return unload_ok


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

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

    async def _async_update_data(self) -> Dict[str, Any]:
        """Fetch data from AdGuard Home."""
        try:
            # Fetch all data concurrently
            tasks = [
                self.api.get_clients(),
                self.api.get_statistics(),
                self.api.get_status(),
            ]

            results = await asyncio.gather(*tasks, return_exceptions=True)
            clients, statistics, status = results

            # FIXED: Check if ALL calls failed with connection errors
            connection_errors = 0
            total_calls = len(results)

            # Update stored data (use previous data if fetch failed)
            if not isinstance(clients, Exception):
                self._clients = {
                    client["name"]: client
                    for client in clients.get("clients", [])
                    if client.get("name")
                }
            else:
                _LOGGER.warning("Failed to update clients data: %s", clients)
                if isinstance(clients, AdGuardConnectionError):
                    connection_errors += 1

            if not isinstance(statistics, Exception):
                self._statistics = statistics
            else:
                _LOGGER.warning("Failed to update statistics data: %s", statistics)
                if isinstance(statistics, AdGuardConnectionError):
                    connection_errors += 1

            if not isinstance(status, Exception):
                self._protection_status = status
            else:
                _LOGGER.warning("Failed to update status data: %s", status)
                if isinstance(status, AdGuardConnectionError):
                    connection_errors += 1

            # FIXED: Only raise UpdateFailed if ALL calls failed with connection errors
            if connection_errors == total_calls:
                raise UpdateFailed("Connection error to AdGuard Home: All API calls failed")

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

        except AdGuardConnectionError as err:
            raise UpdateFailed(f"Connection error to AdGuard Home: {err}") from err
        except Exception as err:
            raise UpdateFailed(f"Error communicating with AdGuard Control Hub: {err}") from err

    @property
    def clients(self) -> Dict[str, Any]:
        """Return clients data."""
        return self._clients

    @property
    def statistics(self) -> Dict[str, Any]:
        """Return statistics data."""
        return self._statistics

    @property
    def protection_status(self) -> Dict[str, Any]:
        """Return protection status data."""
        return self._protection_status