fix: multiple fixes

Signed-off-by: Rafal Zielinski <sq4ind@gmail.com>

custom_components/adguard_hub/__init__.py

@@ -1,20 +1,24 @@
 """
-🛡️ AdGuard Control Hub for Home Assistant.
+AdGuard Control Hub for Home Assistant.
 
-Transform your AdGuard Home into a smart network management powerhouse with 
+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 typing import Dict, Any
+
 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
 from .const import DOMAIN, PLATFORMS, SCAN_INTERVAL, CONF_SSL, CONF_VERIFY_SSL
-from .api import AdGuardHomeAPI
+from .services import AdGuardControlHubServices
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -23,6 +27,7 @@ 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],
@@ -34,16 +39,26 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
 
     # 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])
+        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 Exception as err:
         _LOGGER.error("Failed to connect to AdGuard Home: %s", err)
-        raise ConfigEntryNotReady(f"Unable to connect: {err}")
+        raise ConfigEntryNotReady(f"Unable to connect: {err}") from err
 
     # Create update coordinator
     coordinator = AdGuardControlHubCoordinator(hass, api)
-    await coordinator.async_config_entry_first_refresh()
+
+    try:
+        await coordinator.async_config_entry_first_refresh()
+    except Exception 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, {})
@@ -53,9 +68,24 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     }
 
     # Set up platforms
-    await hass.config_entries.async_forward_entry_setups(entry, 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)
+        raise ConfigEntryNotReady(f"Failed to set up platforms: {err}") from err
 
-    _LOGGER.info("AdGuard Control Hub setup complete")
+    # Register services (only once, not per config entry)
+    if not hass.services.has_service(DOMAIN, "block_services"):
+        services = AdGuardControlHubServices(hass)
+        services.register_services()
+
+        # Store services instance for cleanup
+        hass.data.setdefault(f"{DOMAIN}_services", services)
+
+    _LOGGER.info("AdGuard Control Hub setup complete for %s:%s", 
+                 entry.data[CONF_HOST], entry.data[CONF_PORT])
     return True
 
 
@@ -64,8 +94,19 @@ async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     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)
 
+        # Unregister services if this was the last entry
+        if not hass.data[DOMAIN]:  # No more entries
+            services = hass.data.get(f"{DOMAIN}_services")
+            if services:
+                services.unregister_services()
+                hass.data.pop(f"{DOMAIN}_services", None)
+
+            # Also clean up the empty domain entry
+            hass.data.pop(DOMAIN, None)
+
     return unload_ok
 
 
@@ -81,36 +122,54 @@ class AdGuardControlHubCoordinator(DataUpdateCoordinator):
             update_interval=timedelta(seconds=SCAN_INTERVAL),
         )
         self.api = api
-        self._clients = {}
-        self._statistics = {}
-        self._protection_status = {}
+        self._clients: Dict[str, Any] = {}
+        self._statistics: Dict[str, Any] = {}
+        self._protection_status: Dict[str, Any] = {}
 
-    async def _async_update_data(self):
+    async def _async_update_data(self) -> Dict[str, Any]:
         """Fetch data from AdGuard Home."""
         try:
             # Fetch all data concurrently for better performance
-            results = await asyncio.gather(
+            tasks = [
                 self.api.get_clients(),
                 self.api.get_statistics(),
                 self.api.get_status(),
-                return_exceptions=True,
-            )
+            ]
 
+            results = await asyncio.gather(*tasks, return_exceptions=True)
             clients, statistics, status = results
 
-            # Handle any exceptions
+            # Handle any exceptions in individual requests
             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)
+                    _LOGGER.warning(
+                        "Error fetching %s from %s:%s: %s",
+                        endpoint_names[i],
+                        self.api.host,
+                        self.api.port,
+                        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 {}
+            if not isinstance(clients, Exception):
+                self._clients = {
+                    client["name"]: client
+                    for client in clients.get("clients", [])
+                    if client.get("name")  # Ensure client has a name
+                }
+            else:
+                _LOGGER.warning("Failed to update clients data, keeping previous data")
+
+            if not isinstance(statistics, Exception):
+                self._statistics = statistics
+            else:
+                _LOGGER.warning("Failed to update statistics data, keeping previous data")
+
+            if not isinstance(status, Exception):
+                self._protection_status = status
+            else:
+                _LOGGER.warning("Failed to update status data, keeping previous data")
 
             return {
                 "clients": self._clients,
@@ -118,20 +177,40 @@ class AdGuardControlHubCoordinator(DataUpdateCoordinator):
                 "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}")
+            raise UpdateFailed(f"Error communicating with AdGuard Control Hub: {err}") from err
 
     @property
-    def clients(self):
+    def clients(self) -> Dict[str, Any]:
         """Return clients data."""
         return self._clients
 
     @property
-    def statistics(self):
+    def statistics(self) -> Dict[str, Any]:
         """Return statistics data."""
         return self._statistics
 
     @property
-    def protection_status(self):
+    def protection_status(self) -> Dict[str, Any]:
         """Return protection status data."""
         return self._protection_status
+
+    def get_client(self, client_name: str) -> Dict[str, Any] | None:
+        """Get a specific client by name."""
+        return self._clients.get(client_name)
+
+    def has_client(self, client_name: str) -> bool:
+        """Check if a client exists."""
+        return client_name in self._clients
+
+    @property
+    def client_count(self) -> int:
+        """Return the number of clients."""
+        return len(self._clients)
+
+    @property
+    def is_protection_enabled(self) -> bool:
+        """Return True if protection is enabled."""
+        return self._protection_status.get("protection_enabled", False)