import asyncio from collections.abc import Callable, Coroutine from functools import partial import json from json import JSONDecodeError import logging import os import shutil from typing import Any, NotRequired, TypedDict, cast import aiohttp from aiohttp import ClientError from awesomeversion import AwesomeVersion from homeassistant.core import HomeAssistant from homeassistant.helpers.aiohttp_client import async_get_clientsession from homeassistant.helpers.storage import STORAGE_DIR from homeassistant.loader import async_get_integration from custom_components.powercalc.const import API_URL, BUILT_IN_LIBRARY_DIR, DOMAIN from custom_components.powercalc.helpers import async_cache, clear_async_cache from custom_components.powercalc.power_profile.error import LibraryLoadingError, ProfileDownloadError from custom_components.powercalc.power_profile.loader.protocol import Loader from custom_components.powercalc.power_profile.power_profile import DeviceType, DiscoveryBy _LOGGER = logging.getLogger(__name__) ENDPOINT_LIBRARY = f"{API_URL}/library" ENDPOINT_DOWNLOAD = f"{API_URL}/download" TIMEOUT_SECONDS = 30 class LibraryModel(TypedDict): id: str name: NotRequired[str] aliases: NotRequired[list[str]] legacy_ids: NotRequired[list[str]] hash: str device_type: NotRequired[DeviceType] discovery_by: NotRequired[DiscoveryBy] min_version: NotRequired[str] class LibraryManufacturer(TypedDict): name: str dir_name: str aliases: NotRequired[list[str]] models: list[LibraryModel] class RemoteLoader(Loader): retry_timeout = 3 def __init__(self, hass: HomeAssistant) -> None: self.hass = hass self.library_contents: dict[str, Any] = {} self.model_infos: dict[str, LibraryModel] = {} self.manufacturer_models: dict[str, list[LibraryModel]] = {} self.model_lookup: dict[str, dict[str, list[LibraryModel]]] = {} self.manufacturer_lookup: dict[str, set[str]] = {} self.profile_hashes: dict[str, str] = {} async def initialize(self) -> None: """Initialize the loader.""" integration = await async_get_integration(self.hass, DOMAIN) powercalc_version = AwesomeVersion(str(integration.version)) self._clear_caches() self.library_contents = await self.load_library_json() self.profile_hashes = await self.hass.async_add_executor_job(self._load_profile_hashes) self.model_infos.clear() self.model_lookup.clear() self.manufacturer_models.clear() self.manufacturer_lookup.clear() manufacturers: list[LibraryManufacturer] = self.library_contents.get("manufacturers", []) for manufacturer in manufacturers: self._index_manufacturer(manufacturer, powercalc_version) def _index_manufacturer(self, manufacturer: LibraryManufacturer, powercalc_version: AwesomeVersion) -> None: """Register a manufacturer, its aliases and all of its supported models in the lookup tables.""" manufacturer_name = str(manufacturer.get("dir_name")) models: list[LibraryModel] = manufacturer.get("models", []) or [] # manufacturer alias map (alias -> {canonical manufacturer_name}) self.manufacturer_lookup.setdefault(manufacturer_name.lower(), set()).add(manufacturer_name) for alias in manufacturer.get("aliases", []) or []: self.manufacturer_lookup.setdefault(str(alias).lower(), set()).add(manufacturer_name) # per-manufacturer model lookup kept_models: list[LibraryModel] = [] lookup: dict[str, list[LibraryModel]] = {} for model in models: model_id = str(model.get("id")) self.model_infos[f"{manufacturer_name}/{model_id}"] = model if self._is_unsupported_version(manufacturer_name, model_id, model, powercalc_version): continue kept_models.append(model) self._add_model_to_lookup(lookup, model, model_id.lower()) self.manufacturer_models[manufacturer_name] = kept_models self.model_lookup[manufacturer_name] = lookup @staticmethod def _is_unsupported_version( manufacturer_name: str, model_id: str, model: LibraryModel, powercalc_version: AwesomeVersion, ) -> bool: """Check whether the model requires a newer powercalc version than the one installed.""" min_version = model.get("min_version") if not min_version or powercalc_version >= AwesomeVersion(min_version): return False _LOGGER.debug( "Skipping model %s/%s as it requires powercalc version %s (current: %s)", manufacturer_name, model_id, min_version, powercalc_version, ) return True @staticmethod def _add_model_to_lookup(lookup: dict[str, list[LibraryModel]], model: LibraryModel, model_id_lower: str) -> None: """Bucket a model by its id and aliases. Exact ids take priority over aliases.""" # Exact id bucket first (highest priority) lookup.setdefault(model_id_lower, []).insert(0, model) # Alias buckets afterwards (lower priority) for alias in model.get("aliases", []) or []: alias_lower = str(alias).lower() if alias_lower == model_id_lower: continue # Append to the end to ensure aliased models are always last lookup.setdefault(alias_lower, []).append(model) def _clear_caches(self) -> None: """Clear cached lookups backed by mutable library state.""" clear_async_cache(self.get_manufacturer_listing) clear_async_cache(self.find_manufacturers) clear_async_cache(self.get_model_listing) clear_async_cache(self.find_model) clear_async_cache(self.find_model_migration) clear_async_cache(self.load_model) async def load_library_json(self) -> dict[str, Any]: """Load library.json file""" local_path = self.hass.config.path(STORAGE_DIR, BUILT_IN_LIBRARY_DIR, "library.json") def _load_local_library_json() -> dict[str, Any]: """Load library.json file from local storage""" if not os.path.exists(local_path): raise ProfileDownloadError("Local library.json file not found") with open(local_path) as f: return cast(dict[str, Any], json.load(f)) async def _download_remote_library_json() -> dict[str, Any] | None: """ Download library.json from Github. On success, save it to local storage as a fallback for internet connection issues. """ _LOGGER.debug("Loading library.json from github") session = async_get_clientsession(self.hass) try: async with asyncio.timeout(TIMEOUT_SECONDS), session.get(ENDPOINT_LIBRARY) as resp: if resp.status != 200: raise ProfileDownloadError( f"Failed to download library.json, unexpected status code: {resp.status}", ) data = await resp.read() except (TimeoutError, ClientError) as err: raise ProfileDownloadError(f"Failed to download library.json: {err}") from err def _save_to_local_storage(data: bytes) -> None: os.makedirs(os.path.dirname(local_path), exist_ok=True) with open(local_path, "wb") as f: f.write(data) await self.hass.async_add_executor_job(_save_to_local_storage, data) return cast(dict[str, Any], json.loads(data)) try: return cast(dict[str, Any], await self.download_with_retry(_download_remote_library_json)) except ProfileDownloadError: _LOGGER.debug("Failed to download library.json, falling back to local copy") return await self.hass.async_add_executor_job(_load_local_library_json) @async_cache async def get_manufacturer_listing( self, device_types: set[DeviceType] | None, discovery_by: DiscoveryBy | None = None, ) -> set[tuple[str, str]]: """Get listing of available manufacturers.""" return { (manufacturer["dir_name"], manufacturer["full_name"]) for manufacturer in self.library_contents.get("manufacturers", []) if any( self._model_matches_filters(model, device_types, discovery_by) for model in manufacturer.get("models", []) ) } @async_cache async def find_manufacturers(self, search: str) -> set[str]: """Find the manufacturer in the library.""" return self.manufacturer_lookup.get(search.lower(), set()) @async_cache async def get_model_listing( self, manufacturer: str, device_types: set[DeviceType] | None, discovery_by: DiscoveryBy | None = None, ) -> set[tuple[str, str]]: """Get listing of available models and display names for a given manufacturer.""" models = self.manufacturer_models.get(manufacturer) if not models: return set() return { (model["id"], str(model.get("name") or model["id"])) for model in self.manufacturer_models.get(manufacturer, []) if self._model_matches_filters(model, device_types, discovery_by) } @staticmethod def _model_matches_filters( model: LibraryModel, device_types: set[DeviceType] | None, discovery_by: DiscoveryBy | None, ) -> bool: model_device_type = DeviceType(model.get("device_type", DeviceType.LIGHT)) if device_types and model_device_type not in device_types: return False model_discovery_by = DiscoveryBy(model.get("discovery_by", DiscoveryBy.ENTITY)) return not discovery_by or model_discovery_by == discovery_by @async_cache async def find_model(self, manufacturer: str, search: set[str]) -> list[str]: """Find matching model IDs in the library.""" models = self.model_lookup.get(manufacturer, {}) return [ model["id"] for phrase in search if (phrase_lower := phrase.lower()) in models for model in models[phrase_lower] ] @async_cache async def find_model_migration(self, manufacturer: str, model: str) -> str | None: """Find the canonical model id for a legacy profile id.""" model_lower = model.lower() matches = { str(model_data.get("id")) for manufacturer_data in self.library_contents.get("manufacturers", []) if str(manufacturer_data.get("dir_name", "")).lower() == manufacturer for model_data in manufacturer_data.get("models", []) or [] if model_lower in {str(legacy_id).lower() for legacy_id in model_data.get("legacy_ids", []) or []} } if len(matches) != 1: return None return next(iter(matches)) @async_cache async def load_model( self, manufacturer: str, model: str, force_update: bool = False, retry_count: int = 0, ) -> tuple[dict[str, Any], str] | None: """Load a model, downloading it if necessary, with retry logic.""" model_info = self._get_library_model(manufacturer, model) storage_path = self.get_storage_path(manufacturer, model) model_path = os.path.join(storage_path, "model.json") if await self._needs_update(model_info, manufacturer, model, model_path, force_update): await self._download_profile_with_retry(manufacturer, model, storage_path, model_path) try: json_data = await self._load_model_json(model_path) except JSONDecodeError as e: return await self._handle_json_decode_error(e, manufacturer, model, retry_count) return json_data, storage_path def _get_library_model(self, manufacturer: str, model: str) -> LibraryModel: """Retrieve model info, or raise an error if not found.""" model_info = self.model_infos.get(f"{manufacturer}/{model}") if not model_info: raise LibraryLoadingError(f"Model not found in library: {manufacturer}/{model}") return model_info async def _needs_update( self, model_info: LibraryModel, manufacturer: str, model: str, model_path: str, force_update: bool, ) -> bool: """Check if the model needs to be updated.""" if force_update: return True path_exists = await self.hass.async_add_executor_job(os.path.exists, model_path) if not path_exists: return True existing_hash = self.profile_hashes.get(f"{manufacturer}/{model}") new_hash = model_info.get("hash") return existing_hash != new_hash async def _download_profile_with_retry( self, manufacturer: str, model: str, storage_path: str, model_path: str, ) -> None: """Attempt to download the profile, with retry logic and error handling.""" try: model_info = self._get_library_model(manufacturer, model) model_hash = str(model_info.get("hash")) callback = partial(self.download_profile, manufacturer, model, storage_path, model_hash) await self.download_with_retry(callback) self.profile_hashes[f"{manufacturer}/{model}"] = model_hash await self.hass.async_add_executor_job(self._write_profile_hashes, self.profile_hashes) except ProfileDownloadError as e: path_exists, storage_path_exists = await self.hass.async_add_executor_job( self._profile_paths_exist, model_path, storage_path, ) if not path_exists: if storage_path_exists: await self.hass.async_add_executor_job(shutil.rmtree, storage_path) # pragma: no cover raise e _LOGGER.debug("Failed to download profile, falling back to local profile") @staticmethod def _profile_paths_exist(model_path: str, storage_path: str) -> tuple[bool, bool]: """Check profile paths from the executor.""" return os.path.exists(model_path), os.path.exists(storage_path) async def _load_model_json(self, model_path: str) -> dict[str, Any]: """Load the JSON data from the model file.""" def _load_json() -> dict[str, Any]: with open(model_path) as f: return cast(dict[str, Any], json.load(f)) return await self.hass.async_add_executor_job(_load_json) async def _handle_json_decode_error( self, error: JSONDecodeError, manufacturer: str, model: str, retry_count: int, ) -> tuple[dict[str, Any], str] | None: """Handle JSON decode errors with retry logic.""" _LOGGER.error("model.json file is not valid JSON for manufacturer: %s, model: %s", manufacturer, model) if retry_count < 2: _LOGGER.debug("Retrying to load model.json file") return await self.load_model(manufacturer, model, True, retry_count + 1) raise LibraryLoadingError("Failed to load model.json file") from error def get_storage_path(self, manufacturer: str, model: str) -> str: """Retrieve the storage path for a given manufacturer and model.""" return str(self.hass.config.path(STORAGE_DIR, BUILT_IN_LIBRARY_DIR, manufacturer, model)) async def download_with_retry( self, callback: Callable[[], Coroutine[Any, Any, None | dict[str, Any]]], ) -> None | dict[str, Any]: """Download a file from a remote endpoint with retries""" max_retries = 3 retry_count = 0 while retry_count < max_retries: try: return await callback() except (ClientError, TimeoutError, ProfileDownloadError) as e: _LOGGER.debug(e) retry_count += 1 if retry_count == max_retries: raise ProfileDownloadError( f"Failed to download even after {max_retries} retries, falling back to local copy", ) from e await asyncio.sleep(self.retry_timeout) _LOGGER.warning("Failed to download, retrying... (Attempt %d of %d)", retry_count + 1, max_retries) return None # pragma: no cover async def download_profile(self, manufacturer: str, model: str, storage_path: str, model_hash: str) -> None: """ Download the profile from Github using the Powercalc download API Saves the profile to manufacturer/model directory in .storage/powercalc_profiles folder """ _LOGGER.debug("Downloading profile: %s/%s from github", manufacturer, model) endpoint = f"{ENDPOINT_DOWNLOAD}/{manufacturer}/{model}" def _save_file(data: bytes, directory: str) -> None: """Save file from Github to local storage directory""" path = os.path.join(storage_path, directory) os.makedirs(os.path.dirname(path), exist_ok=True) with open(path, "wb") as f: f.write(data) session = async_get_clientsession(self.hass) try: async with asyncio.timeout(TIMEOUT_SECONDS): async with session.get(endpoint, params={"hash": model_hash}) as resp: if resp.status != 200: raise ProfileDownloadError(f"Failed to download profile: {manufacturer}/{model}") resources = await resp.json() await self.hass.async_add_executor_job(lambda: os.makedirs(storage_path, exist_ok=True)) # Download the files for resource in resources: url = resource.get("url") async with session.get(url) as resp: if resp.status != 200: raise ProfileDownloadError(f"Failed to download github URL: {url}") contents = await resp.read() await self.hass.async_add_executor_job(_save_file, contents, resource.get("path")) except (TimeoutError, aiohttp.ClientError) as e: raise ProfileDownloadError(f"Failed to download profile: {manufacturer}/{model}") from e def _get_profile_hashes_path(self) -> str: """Retrieve the local storage path for the profile hashes file.""" return str(self.hass.config.path(STORAGE_DIR, BUILT_IN_LIBRARY_DIR, ".profile_hashes")) def _load_profile_hashes(self) -> dict[str, str]: """Load profile hashes from local storage""" path = self._get_profile_hashes_path() if not os.path.exists(path): return {} with open(path) as f: return json.load(f) # type: ignore[no-any-return] def _write_profile_hashes(self, hashes: dict[str, str]) -> None: """Write profile hashes to local storage""" path = self._get_profile_hashes_path() with open(path, "w") as json_file: json.dump(hashes, json_file, indent=4)