"""Pstryk Energy integration.""" import logging import asyncio from homeassistant.config_entries import ConfigEntry from homeassistant.core import HomeAssistant from homeassistant.components import mqtt from datetime import timedelta from homeassistant.util import dt as dt_util from .mqtt_publisher import PstrykMqttPublisher from .mqtt_common import setup_periodic_mqtt_publish from .services import async_setup_services, async_unload_services from .const import ( DOMAIN, CONF_MQTT_ENABLED, CONF_MQTT_TOPIC_BUY, CONF_MQTT_TOPIC_SELL, CONF_MQTT_48H_MODE, DEFAULT_MQTT_TOPIC_BUY, DEFAULT_MQTT_TOPIC_SELL ) _LOGGER = logging.getLogger(__name__) async def async_setup(hass: HomeAssistant, config: dict) -> bool: """Set up hass.data structure and services.""" hass.data.setdefault(DOMAIN, {}) # Set up services await async_setup_services(hass) return True async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: """Store API key and forward to sensor platform.""" hass.data[DOMAIN].setdefault(entry.entry_id, {})["api_key"] = entry.data.get("api_key") # Register update listener for option changes - only if not already registered if not entry.update_listeners: entry.async_on_unload(entry.add_update_listener(async_reload_entry)) await hass.config_entries.async_forward_entry_setup(entry, "sensor") _LOGGER.debug("Pstryk entry setup: %s", entry.entry_id) # Setup MQTT publisher if enabled mqtt_enabled = entry.options.get(CONF_MQTT_ENABLED, False) mqtt_topic_buy = entry.options.get(CONF_MQTT_TOPIC_BUY, DEFAULT_MQTT_TOPIC_BUY) mqtt_topic_sell = entry.options.get(CONF_MQTT_TOPIC_SELL, DEFAULT_MQTT_TOPIC_SELL) mqtt_48h_mode = entry.options.get(CONF_MQTT_48H_MODE, False) # Store 48h mode in hass.data for coordinators hass.data[DOMAIN][f"{entry.entry_id}_mqtt_48h_mode"] = mqtt_48h_mode if mqtt_enabled: # Check if MQTT is available if not hass.services.has_service("mqtt", "publish"): _LOGGER.error("MQTT integration is not enabled. Cannot setup EVCC bridge.") # Display persistent notification to user hass.components.persistent_notification.create( "MQTT integration is not enabled. EVCC MQTT Bridge for Pstryk Energy " "cannot function. Please configure MQTT integration in Home Assistant.", title="Pstryk Energy MQTT Error", notification_id=f"{DOMAIN}_mqtt_error_{entry.entry_id}" ) # Still return True to allow the rest of the integration to work return True # Create and store the MQTT publisher mqtt_publisher = PstrykMqttPublisher( hass, entry.entry_id, mqtt_topic_buy, mqtt_topic_sell ) hass.data[DOMAIN][f"{entry.entry_id}_mqtt"] = mqtt_publisher # We need to wait until sensors are fully setup async def start_mqtt_publisher(_): """Start the MQTT publisher after a short delay to ensure coordinators are ready.""" await mqtt_publisher.async_initialize() # Wait for coordinators to be created by sensor platform max_wait = 30 # Maximum 30 seconds wait wait_interval = 1 # Check every second waited = 0 while waited < max_wait: buy_coordinator = hass.data[DOMAIN].get(f"{entry.entry_id}_buy") sell_coordinator = hass.data[DOMAIN].get(f"{entry.entry_id}_sell") if buy_coordinator and sell_coordinator: _LOGGER.debug("Coordinators ready, starting MQTT periodic publishing") break await asyncio.sleep(wait_interval) waited += wait_interval else: _LOGGER.warning("Coordinators not ready after %d seconds, MQTT publishing may fail", max_wait) # Use common function for periodic publishing await setup_periodic_mqtt_publish( hass, entry.entry_id, mqtt_topic_buy, mqtt_topic_sell, interval_minutes=60 ) # Schedule the initialization to happen shortly after setup is complete hass.async_create_task(start_mqtt_publisher(None)) _LOGGER.info("EVCC MQTT Bridge enabled for Pstryk Energy (48h mode: %s), publishing to %s and %s", mqtt_48h_mode, mqtt_topic_buy, mqtt_topic_sell) return True async def _cleanup_coordinators(hass: HomeAssistant, entry: ConfigEntry) -> None: """Clean up coordinators and cancel scheduled tasks.""" # Clean up auto retain task retain_key = f"{entry.entry_id}_auto_retain" if retain_key in hass.data[DOMAIN] and callable(hass.data[DOMAIN][retain_key]): hass.data[DOMAIN][retain_key]() hass.data[DOMAIN].pop(retain_key, None) # Clean up MQTT publisher if exists mqtt_publisher = hass.data[DOMAIN].get(f"{entry.entry_id}_mqtt") if mqtt_publisher: _LOGGER.debug("Cleaning up MQTT publisher for entry %s", entry.entry_id) mqtt_publisher.unsubscribe() hass.data[DOMAIN].pop(f"{entry.entry_id}_mqtt", None) # Clean up coordinators for price_type in ("buy", "sell"): key = f"{entry.entry_id}_{price_type}" coordinator = hass.data[DOMAIN].get(key) if coordinator: _LOGGER.debug("Cleaning up %s coordinator for entry %s", price_type, entry.entry_id) # Cancel scheduled updates if hasattr(coordinator, '_unsub_hourly') and coordinator._unsub_hourly: coordinator._unsub_hourly() coordinator._unsub_hourly = None if hasattr(coordinator, '_unsub_midnight') and coordinator._unsub_midnight: coordinator._unsub_midnight() coordinator._unsub_midnight = None if hasattr(coordinator, '_unsub_afternoon') and coordinator._unsub_afternoon: coordinator._unsub_afternoon() coordinator._unsub_afternoon = None # Remove from hass data hass.data[DOMAIN].pop(key, None) # Clean up mqtt 48h mode flag hass.data[DOMAIN].pop(f"{entry.entry_id}_mqtt_48h_mode", None) async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: """Unload sensor platform and clear data.""" # First cancel coordinators' scheduled updates await _cleanup_coordinators(hass, entry) # Then unload the platform unload_ok = await hass.config_entries.async_forward_entry_unload(entry, "sensor") # Finally clean up data if unload_ok: if entry.entry_id in hass.data[DOMAIN]: hass.data[DOMAIN].pop(entry.entry_id) # Clean up any remaining components for key in list(hass.data[DOMAIN].keys()): if key.startswith(f"{entry.entry_id}_"): hass.data[DOMAIN].pop(key, None) # If this is the last entry, unload services entries = hass.config_entries.async_entries(DOMAIN) if len(entries) <= 1: # This is the last or only entry await async_unload_services(hass) return unload_ok async def async_reload_entry(hass: HomeAssistant, entry: ConfigEntry) -> None: """Reload the config entry when options change.""" await async_unload_entry(hass, entry) await async_setup_entry(hass, entry)